github/trojan-go
4
0
Fork 0
mirror of https://github.com/p4gefau1t/trojan-go.git synced 2024-04-21 12:21:34 +00:00
Code Issues Releases Wiki Activity

fix db auth, speed updater, add speed limiter

Browse Source
This commit is contained in:
p4gefau1t
2020-05-01 12:18:46 -04:00
parent c0f46d38b5
commit 6205ddd980
9 changed files with 116 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
go.mod
+1
View File
@@ -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
)
go.sum
+1
View File
@@ -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=
protocol/trojan/inbound.go
+2 -2
View File
@@ -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
}
protocol/trojan/outbound.go
+2 -2
View File
@@ -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
}
stat/db/db.go
+11 -5
View File
@@ -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
stat/memory/memory.go
+42 -15
View File
@@ -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 {
stat/memory/meomry_test.go → stat/memory/memory_test.go
+32 -1
View File
@@ -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)
}
stat/stat.go
+3 -3
View File
@@ -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
test/proxy_test.go
+22
View File
@@ -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)
}