db5bd942d3
All checks were successful
Build Multi-Platform Binaries / build (push) Successful in 13m17s
147 lines
3.0 KiB
Go
147 lines
3.0 KiB
Go
package db
|
|
|
|
import (
|
|
"database/sql"
|
|
"encoding/json"
|
|
"sync"
|
|
|
|
_ "modernc.org/sqlite"
|
|
|
|
"github.com/gotunnel/pkg/protocol"
|
|
)
|
|
|
|
// SQLiteStore SQLite 存储实现
|
|
type SQLiteStore struct {
|
|
db *sql.DB
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
// NewSQLiteStore 创建 SQLite 存储
|
|
func NewSQLiteStore(dbPath string) (*SQLiteStore, error) {
|
|
db, err := sql.Open("sqlite", dbPath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
s := &SQLiteStore{db: db}
|
|
if err := s.init(); err != nil {
|
|
db.Close()
|
|
return nil, err
|
|
}
|
|
|
|
return s, nil
|
|
}
|
|
|
|
// init 初始化数据库表
|
|
func (s *SQLiteStore) init() error {
|
|
_, err := s.db.Exec(`
|
|
CREATE TABLE IF NOT EXISTS clients (
|
|
id TEXT PRIMARY KEY,
|
|
rules TEXT NOT NULL DEFAULT '[]'
|
|
);
|
|
`)
|
|
return err
|
|
}
|
|
|
|
// Close 关闭数据库连接
|
|
func (s *SQLiteStore) Close() error {
|
|
return s.db.Close()
|
|
}
|
|
|
|
// GetAllClients 获取所有客户端
|
|
func (s *SQLiteStore) GetAllClients() ([]Client, error) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
rows, err := s.db.Query(`SELECT id, rules FROM clients`)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var clients []Client
|
|
for rows.Next() {
|
|
var c Client
|
|
var rulesJSON string
|
|
if err := rows.Scan(&c.ID, &rulesJSON); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := json.Unmarshal([]byte(rulesJSON), &c.Rules); err != nil {
|
|
c.Rules = []protocol.ProxyRule{}
|
|
}
|
|
clients = append(clients, c)
|
|
}
|
|
return clients, nil
|
|
}
|
|
|
|
// GetClient 获取单个客户端
|
|
func (s *SQLiteStore) GetClient(id string) (*Client, error) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
var c Client
|
|
var rulesJSON string
|
|
err := s.db.QueryRow(`SELECT id, rules FROM clients WHERE id = ?`, id).Scan(&c.ID, &rulesJSON)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if err := json.Unmarshal([]byte(rulesJSON), &c.Rules); err != nil {
|
|
c.Rules = []protocol.ProxyRule{}
|
|
}
|
|
return &c, nil
|
|
}
|
|
|
|
// CreateClient 创建客户端
|
|
func (s *SQLiteStore) CreateClient(c *Client) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
rulesJSON, err := json.Marshal(c.Rules)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = s.db.Exec(`INSERT INTO clients (id, rules) VALUES (?, ?)`, c.ID, string(rulesJSON))
|
|
return err
|
|
}
|
|
|
|
// UpdateClient 更新客户端
|
|
func (s *SQLiteStore) UpdateClient(c *Client) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
rulesJSON, err := json.Marshal(c.Rules)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = s.db.Exec(`UPDATE clients SET rules = ? WHERE id = ?`, string(rulesJSON), c.ID)
|
|
return err
|
|
}
|
|
|
|
// DeleteClient 删除客户端
|
|
func (s *SQLiteStore) DeleteClient(id string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
_, err := s.db.Exec(`DELETE FROM clients WHERE id = ?`, id)
|
|
return err
|
|
}
|
|
|
|
// ClientExists 检查客户端是否存在
|
|
func (s *SQLiteStore) ClientExists(id string) (bool, error) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
var count int
|
|
err := s.db.QueryRow(`SELECT COUNT(*) FROM clients WHERE id = ?`, id).Scan(&count)
|
|
return count > 0, err
|
|
}
|
|
|
|
// GetClientRules 获取客户端规则
|
|
func (s *SQLiteStore) GetClientRules(id string) ([]protocol.ProxyRule, error) {
|
|
c, err := s.GetClient(id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return c.Rules, nil
|
|
}
|