add redis support (not tested)

This commit is contained in:
粒粒橙
2020-05-03 19:54:31 +08:00
parent e47227c141
commit c9e77992c9
6 changed files with 143 additions and 0 deletions
+7
View File
@@ -0,0 +1,7 @@
// +build redis
package build
import (
_ "github.com/p4gefau1t/trojan-go/stat/redis"
)
+9
View File
@@ -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"`
+5
View File
@@ -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
)
+10
View File
@@ -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=
+2
View File
@@ -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"
+110
View File
@@ -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)
}