GoProxy/storage/storage.go

package storage

import (
	"database/sql"
	"fmt"
	"log"
	"math/rand"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

type Proxy struct {
	ID        int64
	Address   string // host:port
	Protocol  string // http, socks5
	FailCount int
	LastCheck time.Time
	CreatedAt time.Time
}

type Storage struct {
	db *sql.DB
}

func New(dbPath string) (*Storage, error) {
	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, fmt.Errorf("open db: %w", err)
	}

	db.SetMaxOpenConns(1) // SQLite 单写

	s := &Storage{db: db}
	if err := s.initSchema(); err != nil {
		return nil, err
	}
	return s, nil
}

func (s *Storage) initSchema() error {
	_, err := s.db.Exec(`
		CREATE TABLE IF NOT EXISTS proxies (
			id         INTEGER PRIMARY KEY AUTOINCREMENT,
			address    TEXT NOT NULL UNIQUE,
			protocol   TEXT NOT NULL,
			fail_count INTEGER NOT NULL DEFAULT 0,
			last_check DATETIME,
			created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
		)
	`)
	return err
}

// AddProxy 新增代理，已存在则忽略
func (s *Storage) AddProxy(address, protocol string) error {
	_, err := s.db.Exec(
		`INSERT OR IGNORE INTO proxies (address, protocol) VALUES (?, ?)`,
		address, protocol,
	)
	return err
}

// AddProxies 批量新增
func (s *Storage) AddProxies(proxies []Proxy) error {
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	stmt, err := tx.Prepare(`INSERT OR IGNORE INTO proxies (address, protocol) VALUES (?, ?)`)
	if err != nil {
		tx.Rollback()
		return err
	}
	defer stmt.Close()

	for _, p := range proxies {
		if _, err := stmt.Exec(p.Address, p.Protocol); err != nil {
			log.Printf("insert proxy %s error: %v", p.Address, err)
		}
	}
	return tx.Commit()
}

// GetRandom 随机取一个可用代理
func (s *Storage) GetRandom() (*Proxy, error) {
	rows, err := s.db.Query(
		`SELECT id, address, protocol, fail_count, last_check, created_at
		 FROM proxies WHERE fail_count < 3
		 ORDER BY RANDOM() LIMIT 1`,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	if rows.Next() {
		p := &Proxy{}
		var lastCheck sql.NullTime
		if err := rows.Scan(&p.ID, &p.Address, &p.Protocol, &p.FailCount, &lastCheck, &p.CreatedAt); err != nil {
			return nil, err
		}
		if lastCheck.Valid {
			p.LastCheck = lastCheck.Time
		}
		return p, nil
	}
	return nil, fmt.Errorf("no available proxy")
}

// GetAll 获取所有可用代理
func (s *Storage) GetAll() ([]Proxy, error) {
	rows, err := s.db.Query(
		`SELECT id, address, protocol, fail_count, last_check, created_at
		 FROM proxies WHERE fail_count < 3`,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var proxies []Proxy
	for rows.Next() {
		p := Proxy{}
		var lastCheck sql.NullTime
		if err := rows.Scan(&p.ID, &p.Address, &p.Protocol, &p.FailCount, &lastCheck, &p.CreatedAt); err != nil {
			return nil, err
		}
		if lastCheck.Valid {
			p.LastCheck = lastCheck.Time
		}
		proxies = append(proxies, p)
	}
	return proxies, nil
}

// GetRandomExclude 排除指定地址随机取一个
func (s *Storage) GetRandomExclude(excludes []string) (*Proxy, error) {
	proxies, err := s.GetAll()
	if err != nil {
		return nil, err
	}

	excludeMap := make(map[string]bool)
	for _, e := range excludes {
		excludeMap[e] = true
	}

	var available []Proxy
	for _, p := range proxies {
		if !excludeMap[p.Address] {
			available = append(available, p)
		}
	}

	if len(available) == 0 {
		// 没有可排除的了，随机取任意一个
		return s.GetRandom()
	}

	p := available[rand.Intn(len(available))]
	return &p, nil
}

// Delete 立即删除指定代理
func (s *Storage) Delete(address string) error {
	_, err := s.db.Exec(`DELETE FROM proxies WHERE address = ?`, address)
	return err
}

// IncrFail 增加失败次数
func (s *Storage) IncrFail(address string) error {
	_, err := s.db.Exec(
		`UPDATE proxies SET fail_count = fail_count + 1, last_check = CURRENT_TIMESTAMP WHERE address = ?`,
		address,
	)
	return err
}

// ResetFail 重置失败次数（验证通过）
func (s *Storage) ResetFail(address string) error {
	_, err := s.db.Exec(
		`UPDATE proxies SET fail_count = 0, last_check = CURRENT_TIMESTAMP WHERE address = ?`,
		address,
	)
	return err
}

// DeleteInvalid 删除失败次数超过阈值的代理
func (s *Storage) DeleteInvalid(maxFailCount int) (int64, error) {
	res, err := s.db.Exec(`DELETE FROM proxies WHERE fail_count >= ?`, maxFailCount)
	if err != nil {
		return 0, err
	}
	return res.RowsAffected()
}

// Count 返回可用代理数量
func (s *Storage) Count() (int, error) {
	var count int
	err := s.db.QueryRow(`SELECT COUNT(*) FROM proxies WHERE fail_count < 3`).Scan(&count)
	return count, err
}

// CountByProtocol 按协议统计数量
func (s *Storage) CountByProtocol(protocol string) (int, error) {
	var count int
	err := s.db.QueryRow(`SELECT COUNT(*) FROM proxies WHERE fail_count < 3 AND protocol = ?`, protocol).Scan(&count)
	return count, err
}

// GetByProtocol 按协议获取代理列表
func (s *Storage) GetByProtocol(protocol string) ([]Proxy, error) {
	rows, err := s.db.Query(
		`SELECT id, address, protocol, fail_count, last_check, created_at
		 FROM proxies WHERE fail_count < 3 AND protocol = ?
		 ORDER BY created_at DESC`, protocol,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var proxies []Proxy
	for rows.Next() {
		p := Proxy{}
		var lastCheck sql.NullTime
		if err := rows.Scan(&p.ID, &p.Address, &p.Protocol, &p.FailCount, &lastCheck, &p.CreatedAt); err != nil {
			return nil, err
		}
		if lastCheck.Valid {
			p.LastCheck = lastCheck.Time
		}
		proxies = append(proxies, p)
	}
	return proxies, nil
}

// Close 关闭数据库
func (s *Storage) Close() error {
	return s.db.Close()
}