From aab331f6bc3ef181496e75f66fcd7c0760ca5f9a Mon Sep 17 00:00:00 2001 From: Sousa Chinensis <84500424+sousa-chinensis@users.noreply.github.com> Date: Tue, 1 Jun 2021 15:46:06 +0000 Subject: [PATCH] features: SQLite Authenticator --- component/sqlite.go | 7 ++ docs/content/basic/full-config.md | 36 ++++++++++ go.mod | 1 + go.sum | 2 + statistic/sqlite/config.go | 25 +++++++ statistic/sqlite/sqlite.go | 110 ++++++++++++++++++++++++++++++ tunnel/trojan/config.go | 19 ++++-- tunnel/trojan/server.go | 4 ++ 8 files changed, 197 insertions(+), 7 deletions(-) create mode 100644 component/sqlite.go create mode 100644 statistic/sqlite/config.go create mode 100644 statistic/sqlite/sqlite.go diff --git a/component/sqlite.go b/component/sqlite.go new file mode 100644 index 000000000..d678ef6b3 --- /dev/null +++ b/component/sqlite.go @@ -0,0 +1,7 @@ +// +build sqlite full mini + +package build + +import ( + _ "github.com/p4gefau1t/trojan-go/statistic/sqlite" +) diff --git a/docs/content/basic/full-config.md b/docs/content/basic/full-config.md index 2548b1b56..d10017079 100644 --- a/docs/content/basic/full-config.md +++ b/docs/content/basic/full-config.md @@ -110,6 +110,11 @@ weight: 30 "password": "", "check_rate": 60 }, + "sqlite": { + "enabled": false, + "path": "", + "check_rate": 60 + }, "api": { "enabled": false, "api_addr": "", @@ -329,6 +334,37 @@ CREATE TABLE users ( ); ``` +### ```sqlite```数据库选项 + +trojan-go另设有使用sqlite的方式，但更推荐的方式是使用API。 + +```enabled```表示是否启用sqlite数据库进行用户验证。 + +```path```表示数据库文件的路径（相对当前工作目录）。 + +```check_rate```是trojan-go从MySQL获取用户数据并更新缓存的间隔时间，单位为秒。 + +其他选项可以顾名思义，不再赘述。 + +users表结构和trojan版本定义一致，下面是一个创建users表的例子。注意这里的password指的是密码经过SHA224散列之后的值（字符串），流量download, upload, quota的单位是字节。你可以通过修改数据库users表中的用户记录的方式，添加和删除用户，或者指定用户的流量配额。trojan-go会根据所有的用户流量配额，自动更新当前有效的用户列表。如果download+upload>quota，trojan-go服务器将拒绝该用户的连接。 + +```mysql +CREATE TABLE users ( + id [INT UNSIGNED] NOT NULL, + username VARCHAR (64) NOT NULL, + password CHAR (56) NOT NULL, + quota BIGINT NOT NULL + DEFAULT 0, + download [BIGINT UNSIGNED] NOT NULL + DEFAULT 0, + upload [BIGINT UNSIGNED] NOT NULL + DEFAULT 0, + PRIMARY KEY ( + id + ) +); +``` + ### ```forward_proxy```前置代理选项 前置代理选项允许使用其他代理承载trojan-go的流量 diff --git a/go.mod b/go.mod index 5f0cb5ddd..b03feff13 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/p4gefau1t/trojan-go go 1.16 require ( + github.com/mattn/go-sqlite3 v1.14.7 github.com/go-sql-driver/mysql v1.6.0 github.com/patrickmn/go-cache v2.1.0+incompatible // indirect github.com/refraction-networking/utls v0.0.0-20201210053706-2179f286686b diff --git a/go.sum b/go.sum index 9778cfacc..7eb0b550c 100644 --- a/go.sum +++ b/go.sum @@ -192,6 +192,8 @@ github.com/marten-seemann/qtls-go1-16 v0.1.3 h1:XEZ1xGorVy9u+lJq+WXNE+hiqRYLNvJG github.com/marten-seemann/qtls-go1-16 v0.1.3/go.mod h1:gNpI2Ol+lRS3WwSOtIUUtRwZEQMXjYK+dQSBFbethAk= github.com/mattn/go-colorable v0.0.9/go.mod h1:9vuHe8Xs5qXnSaW/c/ABM9alt+Vo+STaOChaDxuIBZU= github.com/mattn/go-isatty v0.0.3/go.mod h1:M+lRXTBqGeGNdLjl/ufCoiOlB5xdOkqRJdNxMWT7Zi4= +github.com/mattn/go-sqlite3 v1.14.7 h1:fxWBnXkxfM6sRiuH3bqJ4CfzZojMOLVc0UTsTglEghA= +github.com/mattn/go-sqlite3 v1.14.7/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU= github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= github.com/microcosm-cc/bluemonday v1.0.1/go.mod h1:hsXNsILzKxV+sX77C5b8FSuKF00vh2OMYv+xgHpAMF4= github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg= diff --git a/statistic/sqlite/config.go b/statistic/sqlite/config.go new file mode 100644 index 000000000..7d4c5ee1a --- /dev/null +++ b/statistic/sqlite/config.go @@ -0,0 +1,25 @@ +package sqlite + +import ( + "github.com/p4gefau1t/trojan-go/config" +) + +type SQLiteConfig struct { + Enabled bool `json:"enabled" yaml:"enabled"` + Database string `json:"path" yaml:"path"` + CheckRate int `json:"check_rate" yaml:"check-rate"` +} + +type Config struct { + SQLite SQLiteConfig `json:"sqlite" yaml:"sqlite"` +} + +func init() { + config.RegisterConfigCreator(Name, func() interface{} { + return &Config{ + SQLite: SQLiteConfig{ + CheckRate: 30, + }, + } + }) +} diff --git a/statistic/sqlite/sqlite.go b/statistic/sqlite/sqlite.go new file mode 100644 index 000000000..321d0e98a --- /dev/null +++ b/statistic/sqlite/sqlite.go @@ -0,0 +1,110 @@ +package sqlite + +import ( + "context" + "database/sql" + "strings" + "time" + + // SQLite Driver + _ "github.com/mattn/go-sqlite3" + + "github.com/p4gefau1t/trojan-go/common" + "github.com/p4gefau1t/trojan-go/config" + "github.com/p4gefau1t/trojan-go/log" + "github.com/p4gefau1t/trojan-go/statistic" + "github.com/p4gefau1t/trojan-go/statistic/memory" +) + +const Name = "SQLITE" + +type Authenticator struct { + *memory.Authenticator + db *sql.DB + updateDuration time.Duration + ctx context.Context +} + +func (a *Authenticator) updater() { + for { + for _, user := range a.ListUsers() { + // swap upload and download for users + hash := user.Hash() + sent, recv := user.ResetTraffic() + + s, err := a.db.Exec("UPDATE `users` SET `upload`=`upload`+?, `download`=`download`+? WHERE `password`=?;", recv, sent, hash) + if err != nil { + log.Error(common.NewError("failed to update data to user table").Base(err)) + continue + } + if r, err := s.RowsAffected(); err != nil { + if r == 0 { + a.DelUser(hash) + } + } + } + log.Info("buffered data has been written into the database") + + // update memory + rows, err := a.db.Query("SELECT password,quota,download,upload FROM users") + if err != nil || rows.Err() != nil { + log.Error(common.NewError("failed to pull data from the database").Base(err)) + time.Sleep(a.updateDuration) + continue + } + for rows.Next() { + var hash string + var quota, download, upload int64 + err := rows.Scan(&hash, "a, &download, &upload) + if err != nil { + log.Error(common.NewError("failed to obtain data from the query result").Base(err)) + break + } + if download+upload < quota || quota < 0 { + a.AddUser(hash) + } else { + a.DelUser(hash) + } + } + + select { + case <-time.After(a.updateDuration): + case <-a.ctx.Done(): + log.Debug("SQLite daemon exiting...") + return + } + } +} + +func connectDatabase(driverName, dbName string) (*sql.DB, error) { + path := strings.Join([]string{dbName, "?charset=utf8"}, "") + return sql.Open(driverName, path) +} + +func NewAuthenticator(ctx context.Context) (statistic.Authenticator, error) { + cfg := config.FromContext(ctx, Name).(*Config) + db, err := connectDatabase( + "sqlite3", + cfg.SQLite.Database, + ) + if err != nil { + return nil, common.NewError("Failed to connect to database server").Base(err) + } + memoryAuth, err := memory.NewAuthenticator(ctx) + if err != nil { + return nil, err + } + a := &Authenticator{ + db: db, + ctx: ctx, + updateDuration: time.Duration(cfg.SQLite.CheckRate) * time.Second, + Authenticator: memoryAuth.(*memory.Authenticator), + } + go a.updater() + log.Debug("sqlite authenticator created") + return a, nil +} + +func init() { + statistic.RegisterAuthenticatorCreator(Name, NewAuthenticator) +} diff --git a/tunnel/trojan/config.go b/tunnel/trojan/config.go index 32ac3a755..7de208a30 100644 --- a/tunnel/trojan/config.go +++ b/tunnel/trojan/config.go @@ -3,19 +3,24 @@ package trojan import "github.com/p4gefau1t/trojan-go/config" type Config struct { - LocalHost string `json:"local_addr" yaml:"local-addr"` - LocalPort int `json:"local_port" yaml:"local-port"` - RemoteHost string `json:"remote_addr" yaml:"remote-addr"` - RemotePort int `json:"remote_port" yaml:"remote-port"` - DisableHTTPCheck bool `json:"disable_http_check" yaml:"disable-http-check"` - MySQL MySQLConfig `json:"mysql" yaml:"mysql"` - API APIConfig `json:"api" yaml:"api"` + LocalHost string `json:"local_addr" yaml:"local-addr"` + LocalPort int `json:"local_port" yaml:"local-port"` + RemoteHost string `json:"remote_addr" yaml:"remote-addr"` + RemotePort int `json:"remote_port" yaml:"remote-port"` + DisableHTTPCheck bool `json:"disable_http_check" yaml:"disable-http-check"` + MySQL MySQLConfig `json:"mysql" yaml:"mysql"` + SQLite SQLiteConfig `json:"sqlite" yaml:"sqlite"` + API APIConfig `json:"api" yaml:"api"` } type MySQLConfig struct { Enabled bool `json:"enabled" yaml:"enabled"` } +type SQLiteConfig struct { + Enabled bool `json:"enabled" yaml:"enabled"` +} + type APIConfig struct { Enabled bool `json:"enabled" yaml:"enabled"` } diff --git a/tunnel/trojan/server.go b/tunnel/trojan/server.go index 52e46cc38..113cc1b62 100644 --- a/tunnel/trojan/server.go +++ b/tunnel/trojan/server.go @@ -15,6 +15,7 @@ import ( "github.com/p4gefau1t/trojan-go/statistic" "github.com/p4gefau1t/trojan-go/statistic/memory" "github.com/p4gefau1t/trojan-go/statistic/mysql" + "github.com/p4gefau1t/trojan-go/statistic/sqlite" "github.com/p4gefau1t/trojan-go/tunnel" "github.com/p4gefau1t/trojan-go/tunnel/mux" ) @@ -214,6 +215,9 @@ func NewServer(ctx context.Context, underlay tunnel.Server) (*Server, error) { if cfg.MySQL.Enabled { log.Debug("mysql enabled") auth, err = statistic.NewAuthenticator(ctx, mysql.Name) + } else if cfg.SQLite.Enabled { + log.Debug("sqlite enabled") + auth, err = statistic.NewAuthenticator(ctx, sqlite.Name) } else { log.Debug("auth by config file") auth, err = statistic.NewAuthenticator(ctx, memory.Name)