diff --git a/go.mod b/go.mod index 895740e..52fb495 100644 --- a/go.mod +++ b/go.mod @@ -17,6 +17,7 @@ require ( golang.org/x/crypto v0.0.0-20200423211502-4bdfaf469ed5 golang.org/x/net v0.0.0-20200421231249-e086a090c8fd 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 v2ray.com/core v4.19.1+incompatible ) diff --git a/go.sum b/go.sum index 88d3c30..fcfa70c 100644 --- a/go.sum +++ b/go.sum @@ -415,6 +415,7 @@ golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20190921001708-c4c64cad1fd0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= +golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs= golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/tools v0.0.0-20180828015842-6cd1fcedba52/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= diff --git a/protocol/trojan/inbound.go b/protocol/trojan/inbound.go index 5716f33..bf79561 100644 --- a/protocol/trojan/inbound.go +++ b/protocol/trojan/inbound.go @@ -31,14 +31,14 @@ type TrojanInboundConnSession struct { func (i *TrojanInboundConnSession) Write(p []byte) (int, error) { n, err := i.rwc.Write(p) i.sent += uint64(n) - i.meter.Count(uint64(n), 0) + i.meter.Count(n, 0) return n, err } func (i *TrojanInboundConnSession) Read(p []byte) (int, error) { n, err := i.rwc.Read(p) i.recv += uint64(n) - i.meter.Count(0, uint64(n)) + i.meter.Count(0, n) return n, err } diff --git a/protocol/trojan/outbound.go b/protocol/trojan/outbound.go index d5da1eb..d7985d3 100644 --- a/protocol/trojan/outbound.go +++ b/protocol/trojan/outbound.go @@ -29,14 +29,14 @@ func (o *TrojanOutboundConnSession) SetMeter(meter stat.TrafficMeter) { func (o *TrojanOutboundConnSession) Write(p []byte) (int, error) { n, err := o.rwc.Write(p) - o.meter.Count(uint64(n), 0) + o.meter.Count(n, 0) o.sent += uint64(n) return n, err } func (o *TrojanOutboundConnSession) Read(p []byte) (int, error) { n, err := o.rwc.Read(p) - o.meter.Count(0, uint64(n)) + o.meter.Count(0, n) o.recv += uint64(n) return n, err } diff --git a/stat/db/db.go b/stat/db/db.go index 463846d..da99681 100644 --- a/stat/db/db.go +++ b/stat/db/db.go @@ -32,10 +32,11 @@ func (a *DBAuth) updater() { } for _, user := range users { //swap upload and download for users - s, err := tx.Prepare("UPDATE users SET upload=upload+? WHERE password=?;") - common.Must(err) hash := user.Hash() sent, recv := user.GetAndReset() + + s, err := tx.Prepare("UPDATE users SET upload=upload+? WHERE password=?;") + common.Must(err) _, err = s.Exec(recv, hash) s, err = tx.Prepare("UPDATE users SET download=download+? WHERE password=?;") @@ -100,10 +101,15 @@ func NewDBAuth(ctx context.Context, config *conf.GlobalConfig) (stat.Authenticat 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.MySQL.CheckRate) * time.Second, + db: db, + ctx: ctx, + updateDuration: time.Duration(config.MySQL.CheckRate) * time.Second, + MemoryAuthenticator: memoryAuth.(*memory.MemoryAuthenticator), } go a.updater() return a, nil diff --git a/stat/memory/memory.go b/stat/memory/memory.go index a2b297c..ec9070d 100644 --- a/stat/memory/memory.go +++ b/stat/memory/memory.go @@ -9,18 +9,24 @@ import ( "github.com/p4gefau1t/trojan-go/common" "github.com/p4gefau1t/trojan-go/conf" "github.com/p4gefau1t/trojan-go/stat" + "golang.org/x/time/rate" ) type MemoryTrafficMeter struct { stat.TrafficMeter - sent uint64 - recv uint64 - lastSent uint64 - lastRecv uint64 - hash string - ctx context.Context - cancel context.CancelFunc + sent uint64 + recv uint64 + lastSent uint64 + lastRecv uint64 + speedLock sync.Mutex + sendSpeed uint64 + recvSpeed uint64 + hash string + sendLimiter *rate.Limiter + recvLimiter *rate.Limiter + ctx context.Context + cancel context.CancelFunc } func (m *MemoryTrafficMeter) Close() error { @@ -29,11 +35,29 @@ func (m *MemoryTrafficMeter) Close() error { return nil } -func (m *MemoryTrafficMeter) Count(sent, recv uint64) { +func (m *MemoryTrafficMeter) Count(sent, recv int) { + if m.sendLimiter != nil && sent != 0 { + m.sendLimiter.WaitN(m.ctx, sent) + } else if m.recvLimiter != nil && recv != 0 { + m.recvLimiter.WaitN(m.ctx, recv) + } atomic.AddUint64(&m.sent, uint64(sent)) atomic.AddUint64(&m.recv, uint64(recv)) } +func (m *MemoryTrafficMeter) LimitSpeed(sent, recv int) { + if sent == 0 { + m.sendLimiter = nil + } else { + m.sendLimiter = rate.NewLimiter(rate.Limit(sent), sent*2) + } + if recv == 0 { + m.recvLimiter = nil + } else { + m.recvLimiter = rate.NewLimiter(rate.Limit(recv), recv*2) + } +} + func (m *MemoryTrafficMeter) Hash() string { return m.hash } @@ -63,19 +87,22 @@ func (m *MemoryTrafficMeter) speedUpdater() { case <-m.ctx.Done(): return case <-time.After(time.Second): - lastSent, lastRecv := m.Get() - atomic.StoreUint64(&m.lastSent, lastSent) - atomic.StoreUint64(&m.lastRecv, lastRecv) + m.speedLock.Lock() + sent, recv := m.Get() + m.sendSpeed = sent - m.lastSent + m.recvSpeed = recv - m.lastRecv + m.lastSent = sent + m.lastRecv = recv + m.speedLock.Unlock() } } } func (m *MemoryTrafficMeter) GetSpeed() (uint64, uint64) { - sent, recv := m.Get() - lastSent := atomic.LoadUint64(&m.lastSent) - lastRecv := atomic.LoadUint64(&m.lastRecv) - return sent - lastSent, recv - lastRecv + m.speedLock.Lock() + defer m.speedLock.Unlock() + return m.sendSpeed, m.recvSpeed } type MemoryAuthenticator struct { diff --git a/stat/memory/meomry_test.go b/stat/memory/memory_test.go similarity index 56% rename from stat/memory/meomry_test.go rename to stat/memory/memory_test.go index e41358d..288de59 100644 --- a/stat/memory/meomry_test.go +++ b/stat/memory/memory_test.go @@ -35,7 +35,38 @@ func TestMemoryAuth(t *testing.T) { } }() - for i := 0; i < 100; i++ { + for i := 0; i < 15; i++ { + fmt.Println(traffic.GetSpeed()) + time.Sleep(time.Millisecond * 1000) + } + cancel() +} + +func TestLimitSpeed(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + config := &conf.GlobalConfig{ + Hash: map[string]string{ + "hash": "password", + }, + } + auth, err := NewMemoryAuth(ctx, config) + common.Must(err) + valid, traffic := auth.AuthUser("hash") + if !valid { + t.Fail() + } + traffic.LimitSpeed(5000, 6000) + go func() { + for { + traffic.Count(50, 0) + } + }() + go func() { + for { + traffic.Count(0, 100) + } + }() + for i := 0; i < 15; i++ { fmt.Println(traffic.GetSpeed()) time.Sleep(time.Millisecond * 1000) } diff --git a/stat/stat.go b/stat/stat.go index 0c7e669..d536bce 100644 --- a/stat/stat.go +++ b/stat/stat.go @@ -11,17 +11,17 @@ import ( type TrafficMeter interface { io.Closer Hash() string - Count(sent uint64, recv uint64) + Count(sent int, recv int) Get() (sent uint64, recv uint64) Reset() GetAndReset() (sent uint64, recv uint64) GetSpeed() (sent uint64, recv uint64) - LimitSpeed(sent uint64, recv uint64) + LimitSpeed(sent int, recv int) } type Authenticator interface { io.Closer - AuthUser(hash string) (bool, TrafficMeter) + AuthUser(hash string) (valid bool, meter TrafficMeter) AddUser(hash string) error DelUser(hash string) error ListUsers() []TrafficMeter diff --git a/test/proxy_test.go b/test/proxy_test.go index 9b246bf..0da601c 100644 --- a/test/proxy_test.go +++ b/test/proxy_test.go @@ -20,6 +20,7 @@ import ( _ "github.com/p4gefau1t/trojan-go/log/golog" "github.com/p4gefau1t/trojan-go/proxy/client" "github.com/p4gefau1t/trojan-go/proxy/server" + _ "github.com/p4gefau1t/trojan-go/stat/db" _ "github.com/p4gefau1t/trojan-go/stat/memory" "golang.org/x/net/proxy" "golang.org/x/net/websocket" @@ -182,6 +183,19 @@ func addTCPOption(config *conf.GlobalConfig) *conf.GlobalConfig { return config } +func addMySQLOption(config *conf.GlobalConfig) *conf.GlobalConfig { + config.MySQL = conf.MySQLConfig{ + Enabled: true, + ServerHost: "127.0.0.1", + ServerPort: 3306, + Database: "trojan", + Username: "root", + Password: "password", + CheckRate: 1, + } + return config +} + func RunClient(ctx context.Context, config *conf.GlobalConfig) { c := client.Client{} r, err := c.Build(config) @@ -479,3 +493,11 @@ func TestTCPOptions(t *testing.T) { clientConfig := addTCPOption(getBasicClientConfig()) CheckClientServer(t, clientConfig, serverConfig) } + +func TestMySQL(t *testing.T) { + serverConfig := addMySQLOption(getBasicServerConfig()) + clientConfig := getBasicClientConfig() + clientConfig.Passwords = getPasswords("mysqlpassword") + clientConfig.Hash = getHash("mysqlpassword") + CheckClientServer(t, clientConfig, serverConfig) +}