From c9e77992c99765fb5fa5e3434caca47edc57c559 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=B2=92=E7=B2=92=E6=A9=99?= Date: Sun, 3 May 2020 19:54:31 +0800 Subject: [PATCH] add redis support (not tested) --- build/redis.go | 7 +++ conf/conf.go | 9 ++++ go.mod | 5 ++ go.sum | 10 ++++ stat/mysql/mysql.go | 2 + stat/redis/redis.go | 110 ++++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 143 insertions(+) create mode 100644 build/redis.go create mode 100644 stat/redis/redis.go diff --git a/build/redis.go b/build/redis.go new file mode 100644 index 0000000..0b4c014 --- /dev/null +++ b/build/redis.go @@ -0,0 +1,7 @@ +// +build redis + +package build + +import ( + _ "github.com/p4gefau1t/trojan-go/stat/redis" +) diff --git a/conf/conf.go b/conf/conf.go index 4338aad..15921fb 100644 --- a/conf/conf.go +++ b/conf/conf.go @@ -72,6 +72,14 @@ type MySQLConfig struct { CheckRate int `json:"check_rate"` } +type RedisConfig struct { + Enabled bool `json:"enabled"` + ServerHost string `json:"server_addr"` + ServerPort int `json:"server_port"` + Password string `json:"password"` + CheckRate int `json:"check_rate"` +} + type ForwardProxyConfig struct { Enabled bool `json:"enabled"` ProxyHost string `json:"proxy_addr"` @@ -141,6 +149,7 @@ type GlobalConfig struct { TLS TLSConfig `json:"ssl"` TCP TCPConfig `json:"tcp"` MySQL MySQLConfig `json:"mysql"` + Redis RedisConfig `json:"redis"` Mux MuxConfig `json:"mux"` Router RouterConfig `json:"router"` Websocket WebsocketConfig `json:"websocket"` diff --git a/go.mod b/go.mod index 0783f76..6c4d50d 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,9 @@ require ( github.com/go-sql-driver/mysql v1.5.0 github.com/golang/protobuf v1.4.0 github.com/mattn/go-sqlite3 v2.0.3+incompatible // indirect + github.com/mediocregopher/radix/v3 v3.5.0 + github.com/onsi/ginkgo v1.10.1 // indirect + github.com/onsi/gomega v1.7.0 // indirect github.com/proullon/ramsql v0.0.0-20181213202341-817cee58a244 github.com/refraction-networking/utls v0.0.0-20190909200633-43c36d3c1f57 github.com/smartystreets/goconvey v1.6.4 @@ -18,5 +21,7 @@ require ( golang.org/x/sys v0.0.0-20200420163511-1957bb5e6d1f golang.org/x/time v0.0.0-20191024005414-555d28b269f0 google.golang.org/grpc v1.29.1 + gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect + gopkg.in/yaml.v2 v2.2.4 // indirect v2ray.com/core v4.19.1+incompatible ) diff --git a/go.sum b/go.sum index 60575e4..c8410be 100644 --- a/go.sum +++ b/go.sum @@ -191,6 +191,8 @@ github.com/mattn/go-sqlite3 v2.0.3+incompatible h1:gXHsfypPkaMZrKbD5209QV9jbUTJK github.com/mattn/go-sqlite3 v2.0.3+incompatible/go.mod h1:FPy6KqzDD04eiIsT53CuJW3U88zkxoIYsOqkbpncsNc= github.com/mattn/go-tty v0.0.0-20180219170247-931426f7535a/go.mod h1:XPvLUNfbS4fJH25nqRHfWLMa1ONC8Amw+mIA639KxkE= github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= +github.com/mediocregopher/radix/v3 v3.5.0 h1:8QHQmNh2ne9aFxTD3z63u/bkPPiOtknHoz80oP8EA/E= +github.com/mediocregopher/radix/v3 v3.5.0/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8= github.com/miekg/dns v1.1.4/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg= github.com/miekg/dns v1.1.27 h1:aEH/kqUzUxGJ/UHcEKdJY+ugH6WEzsEBBSPa8zuy1aM= github.com/miekg/dns v1.1.27/go.mod h1:KNUDUusw/aVsxyTYZM1oqvCicbwhgbNgztCETuNZ7xM= @@ -212,8 +214,12 @@ github.com/olekukonko/tablewriter v0.0.1/go.mod h1:vsDQFd/mU46D+Z4whnwzcISnGGzXW github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/ginkgo v1.7.0 h1:WSHQ+IS43OoUrWtD1/bbclrwK8TTH5hzp+umCiuxHgs= github.com/onsi/ginkgo v1.7.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.10.1 h1:q/mM8GF/n0shIN8SaAZ0V+jnLPzen6WIVZdiwrRlMlo= +github.com/onsi/ginkgo v1.10.1/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/gomega v1.4.3 h1:RE1xgDvH7imwFD45h+u2SgIfERHlS2yNG4DObb5BSKU= github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= +github.com/onsi/gomega v1.7.0 h1:XPnZz8VVBHjVsy1vzJmRwIcSwiUO+JFfrv/xGiigmME= +github.com/onsi/gomega v1.7.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= github.com/openzipkin/zipkin-go v0.1.6/go.mod h1:QgAqvLzwWbR/WpD4A3cGpPtJrZXNIiJc5AZX7/PBEpw= github.com/oracle/oci-go-sdk v7.0.0+incompatible/go.mod h1:VQb79nF8Z2cwLkLS35ukwStZIg5F66tcBccjip/j888= github.com/ovh/go-ovh v0.0.0-20181109152953-ba5adb4cf014/go.mod h1:joRatxRJaZBsY3JAOEMcoOp05CnZzsx4scTxi95DHyQ= @@ -505,6 +511,8 @@ gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLks gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo= +gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= gopkg.in/fsnotify.v1 v1.4.7 h1:xOHLXZwVvI9hhs+cLKq5+I5onOuwQLhQwiu63xxlHs4= gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= @@ -522,6 +530,8 @@ gopkg.in/yaml.v2 v2.0.0-20170812160011-eb3733d160e7/go.mod h1:JAlM8MvJe8wmxCU4Bl gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.4 h1:/eiJrUcujPVeJ3xlSWaiNi3uSVmDGBK1pDHUHAnao1I= +gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= h12.io/socks v1.0.0/go.mod h1:MdYbo5/eB9ka7u5dzW2Qh0iSyJENwB3KI5H5ngenFGA= honnef.co/go/tools v0.0.0-20180728063816-88497007e858/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= diff --git a/stat/mysql/mysql.go b/stat/mysql/mysql.go index e9885ac..b5954a7 100644 --- a/stat/mysql/mysql.go +++ b/stat/mysql/mysql.go @@ -7,7 +7,9 @@ import ( "strings" "time" + // MySQL Driver _ "github.com/go-sql-driver/mysql" + "github.com/p4gefau1t/trojan-go/common" "github.com/p4gefau1t/trojan-go/conf" "github.com/p4gefau1t/trojan-go/log" diff --git a/stat/redis/redis.go b/stat/redis/redis.go new file mode 100644 index 0000000..bdf89c0 --- /dev/null +++ b/stat/redis/redis.go @@ -0,0 +1,110 @@ +package redis + +import ( + "context" + "regexp" + "strconv" + "time" + + "github.com/mediocregopher/radix/v3" + "github.com/p4gefau1t/trojan-go/common" + "github.com/p4gefau1t/trojan-go/conf" + "github.com/p4gefau1t/trojan-go/log" + "github.com/p4gefau1t/trojan-go/stat" + "github.com/p4gefau1t/trojan-go/stat/memory" +) + +type DBAuth struct { + *memory.MemoryAuthenticator + db *radix.Pool + updateDuration time.Duration + ctx context.Context +} + +func validateHash(i string) bool { + re := regexp.MustCompile("^[0-9a-fA-F]{56}$") + return re.MatchString(i) +} + +func (a *DBAuth) updater() { + for { + users := a.ListUsers() + for _, user := range users { + // fetch user flow + hash := user.Hash() + sent, recv := user.GetAndReset() + + // check if user exists in DB + var exist bool + if err := a.db.Do(radix.Cmd(&exist, "EXISTS", hash)); err != nil { + log.Error(common.NewError("failed to check user in DB").Base(err)) + } + + // remove the user if not + if !exist { + a.DelUser(hash) + continue + } + + // update flow to DB + pipe := radix.Pipeline( + radix.Cmd(nil, "HINCRBY", hash, "upload", strconv.FormatUint(recv, 10)), + radix.Cmd(nil, "HINCRBY", hash, "download", strconv.FormatUint(sent, 10)), + ) + if err := a.db.Do(pipe); err != nil { + log.Error(common.NewError("failed to execute pipeline").Base(err)) + } + } + log.Info("buffered data has been written into the database") + + //update memory + var keys []string + if err := a.db.Do(radix.Cmd(&keys, "KEYS")); err != nil { + log.Error(common.NewError("failed to pull data from the database").Base(err)) + time.Sleep(a.updateDuration) + continue + } + for _, key := range keys { + if validateHash(key) { + a.AddUser(key) + } + } + + select { + case <-time.After(a.updateDuration): + case <-a.ctx.Done(): + log.Debug("db daemon exiting...") + return + } + } +} + +func NewDBAuth(ctx context.Context, config *conf.GlobalConfig) (stat.Authenticator, error) { + addr := config.Redis.ServerHost + ":" + strconv.Itoa(config.Redis.ServerPort) + conn := func(network, addr string) (radix.Conn, error) { + return radix.Dial(network, addr, + radix.DialAuthPass("mySuperSecretPassword"), + ) + } + db, err := radix.NewPool("tcp", addr, 10, radix.PoolConnFunc(conn)) + if err != nil { + return nil, common.NewError("failed to connect to database server").Base(err) + } + + memoryAuth, err := memory.NewMemoryAuth(ctx, config) + if err != nil { + return nil, err + } + a := &DBAuth{ + db: db, + ctx: ctx, + updateDuration: time.Duration(config.Redis.CheckRate) * time.Second, + MemoryAuthenticator: memoryAuth.(*memory.MemoryAuthenticator), + } + go a.updater() + return a, nil +} + +func init() { + stat.RegisterAuthCreator("redis", NewDBAuth) +}