update api and auth

This commit is contained in:
p4gefau1t
2020-05-01 04:34:31 -04:00
parent 37421ac8dd
commit 9c4aa2a388
25 changed files with 563 additions and 593 deletions
+23 -28
View File
@@ -3,8 +3,8 @@ package api
import (
"context"
"net"
"time"
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/conf"
"github.com/p4gefau1t/trojan-go/log"
"github.com/p4gefau1t/trojan-go/stat"
@@ -14,16 +14,23 @@ import (
type ClientAPI struct {
TrojanClientServiceServer
meter stat.TrafficMeter
auth stat.Authenticator
ctx context.Context
uploadSpeed uint64
downloadSpeed uint64
lastSent uint64
lastRecv uint64
ctx context.Context
}
func (s *ClientAPI) GetTraffic(context.Context, *GetTrafficRequest) (*GetTrafficResponse, error) {
sent, recv := s.meter.Query("")
func (s *ClientAPI) GetTraffic(ctx context.Context, req *GetTrafficRequest) (*GetTrafficResponse, error) {
if req.User == nil {
return nil, common.NewError("user is unspecified")
}
valid, meter := s.auth.AuthUser(req.User.Hash)
if !valid {
return nil, common.NewError("user " + req.User.Hash + " not found")
}
sent, recv := meter.Get()
resp := &GetTrafficResponse{
TrafficTotal: &Traffic{
UploadTraffic: sent,
@@ -33,39 +40,27 @@ func (s *ClientAPI) GetTraffic(context.Context, *GetTrafficRequest) (*GetTraffic
return resp, nil
}
func (s *ClientAPI) GetSpeed(context.Context, *GetSpeedRequest) (*GetSpeedResponse, error) {
func (s *ClientAPI) GetSpeed(ctx context.Context, req *GetSpeedRequest) (*GetSpeedResponse, error) {
valid, meter := s.auth.AuthUser(req.User.Hash)
if !valid {
return &GetSpeedResponse{}, nil
}
sent, recv := meter.GetSpeed()
resp := &GetSpeedResponse{
SpeedCurrent: &Speed{
UploadSpeed: s.uploadSpeed,
DownloadSpeed: s.downloadSpeed,
UploadSpeed: sent,
DownloadSpeed: recv,
},
}
return resp, nil
}
func (s *ClientAPI) calcSpeed() {
for {
select {
case <-time.After(time.Second):
// TODO avoid racing
sent, recv := s.meter.Query("")
s.uploadSpeed = sent - s.lastSent
s.downloadSpeed = recv - s.lastRecv
s.lastSent = sent
s.lastRecv = recv
case <-s.ctx.Done():
return
}
}
}
func RunClientAPIService(ctx context.Context, config *conf.GlobalConfig, meter stat.TrafficMeter) error {
func RunClientAPIService(ctx context.Context, config *conf.GlobalConfig, auth stat.Authenticator) error {
server := grpc.NewServer()
service := &ClientAPI{
meter: meter,
ctx: ctx,
ctx: ctx,
auth: auth,
}
go service.calcSpeed()
RegisterTrojanClientServiceServer(server, service)
listener, err := net.Listen("tcp", config.API.APIAddress.String())
if err != nil {
+21 -7
View File
@@ -7,25 +7,39 @@ import (
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/conf"
"github.com/p4gefau1t/trojan-go/stat"
"github.com/p4gefau1t/trojan-go/stat/memory"
"google.golang.org/grpc"
)
func TestClientAPI(t *testing.T) {
meter := &stat.MemoryTrafficMeter{}
go RunClientAPIService(context.Background(), &conf.GlobalConfig{
ctx, cancel := context.WithCancel(context.Background())
auth, err := memory.NewMemoryAuth(ctx, &conf.GlobalConfig{})
common.Must(err)
go RunClientAPIService(ctx, &conf.GlobalConfig{
API: conf.APIConfig{
APIAddress: common.NewAddress("127.0.0.1", 10000, "tcp"),
},
}, meter)
meter.Count("test", 123, 456)
}, auth)
common.Must(auth.AddUser("hash1234"))
valid, meter := auth.AuthUser("hash1234")
if !valid {
t.Fail()
}
meter.Count(1234, 5678)
time.Sleep(time.Second)
conn, err := grpc.Dial("127.0.0.1:10000", grpc.WithInsecure())
common.Must(err)
client := NewTrojanClientServiceClient(conn)
resp, err := client.GetTraffic(context.Background(), &GetTrafficRequest{})
resp, err := client.GetTraffic(ctx, &GetTrafficRequest{User: &User{
Hash: "hash1234",
}})
common.Must(err)
if resp.TrafficTotal.DownloadTraffic != 456 || resp.TrafficTotal.UploadTraffic != 123 {
if resp.TrafficTotal.DownloadTraffic != 5678 || resp.TrafficTotal.UploadTraffic != 1234 {
t.Fail()
}
resp, err = client.GetTraffic(ctx, &GetTrafficRequest{})
if err == nil {
t.Fail()
}
cancel()
}
+3
View File
@@ -7,10 +7,13 @@ require (
github.com/go-acme/lego/v3 v3.5.0
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/posener/h2conn v0.0.0-20180911140238-13e7df33ed15
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
github.com/xtaci/smux v1.5.12
github.com/ziutek/mymysql v1.5.4 // indirect
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
+16
View File
@@ -81,6 +81,7 @@ github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1m
github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
github.com/exoscale/egoscale v0.18.1/go.mod h1:Z7OOdzzTOz1Q1PjQXumlz9Wn/CddH0zSYdCF3rnBKXE=
github.com/fatih/structs v1.1.0/go.mod h1:9NiDSp5zOcgEDl+j00MP/WkGVPOlPRLejGD8Ga6PJ7M=
github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I=
github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
github.com/gdamore/encoding v0.0.0-20151215212835-b23993cbb635/go.mod h1:yrQYJKKDTrHmbYxI7CYi+/hbdiDT2m4Hj+t0ikCjsrQ=
github.com/gdamore/tcell v1.1.0/go.mod h1:tqyG50u7+Ctv1w5VX67kLzKcj9YXR/JSBZQq/+mLl1A=
@@ -92,6 +93,8 @@ github.com/go-errors/errors v1.0.1/go.mod h1:f4zRHt4oKfwPJE5k8C9vpYG+aDHdBFUsgrm
github.com/go-gl/glfw v0.0.0-20190409004039-e6da0acd62b1/go.mod h1:vR7hzQXu2zJy9AVAgeJqvqgH9Q5CA+iKCZ2gyEVpxRU=
github.com/go-gl/glfw/v3.3/glfw v0.0.0-20191125211704-12ad95a8df72/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8=
github.com/go-gl/glfw/v3.3/glfw v0.0.0-20200222043503-6f7a984d4dc4/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8=
github.com/go-gorp/gorp v2.0.0+incompatible h1:dIQPsBtl6/H1MjVseWuWPXa7ET4p6Dve4j3Hg+UjqYw=
github.com/go-gorp/gorp v2.0.0+incompatible/go.mod h1:7IfkAQnO7jfT/9IQ3R9wL1dFhukN6aQxzKTHnkxzA/E=
github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9GBnD5lWE=
github.com/go-logfmt/logfmt v0.4.0/go.mod h1:3RMwSq7FuexP4Kalkev3ejPJsZTpXXBr9+V4qmtdjCk=
@@ -154,6 +157,7 @@ github.com/grpc-ecosystem/grpc-gateway v1.8.5/go.mod h1:vNeuVxBJEsws4ogUvrchl83t
github.com/h2non/parth v0.0.0-20190131123155-b4df798d6542/go.mod h1:Ow0tF8D4Kplbc8s8sSb3V2oUCygFHVp8gC3Dn6U4MNI=
github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
github.com/hpcloud/tail v1.0.0 h1:nfCOvKYfkgYP8hkirhJocXT2+zOD8yUNjXaWfTlyFKI=
github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
github.com/iij/doapi v0.0.0-20190504054126-0bbf12d6d7df/go.mod h1:QMZY7/J/KSQEhKWFeDesPjMj+wCHReeknARU3wqlyN4=
@@ -178,6 +182,8 @@ github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/labbsr0x/bindman-dns-webhook v1.0.2/go.mod h1:p6b+VCXIR8NYKpDr8/dg1HKfQoRHCdcsROXKvmoehKA=
github.com/labbsr0x/goh v1.0.1/go.mod h1:8K2UhVoaWXcCU7Lxoa2omWnC8gyW8px7/lmO61c027w=
github.com/lib/pq v1.0.0 h1:X5PMW56eZitiTeO7tKzZxFCSpbFZJtkMMooicw2us9A=
github.com/lib/pq v1.0.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
github.com/linode/linodego v0.10.0/go.mod h1:cziNP7pbvE3mXIPneHj0oRY8L1WtGEIKlZ8LANE4eXA=
github.com/liquidweb/liquidweb-go v1.6.0/go.mod h1:UDcVnAMDkZxpw4Y7NOHkqoeiGacVLEIG/i5J9cyixzQ=
github.com/lucasb-eyer/go-colorful v0.0.0-20180709185858-c7842319cf3a/go.mod h1:NXg0ArsFk0Y01623LgUqoqcouGDB+PwCCQlrwrG6xJ4=
@@ -186,6 +192,8 @@ github.com/mattn/go-isatty v0.0.3/go.mod h1:M+lRXTBqGeGNdLjl/ufCoiOlB5xdOkqRJdNx
github.com/mattn/go-runewidth v0.0.2/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=
github.com/mattn/go-runewidth v0.0.3/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=
github.com/mattn/go-runewidth v0.0.4/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=
github.com/mattn/go-sqlite3 v2.0.3+incompatible h1:gXHsfypPkaMZrKbD5209QV9jbUTJKjyR5WD3HYQSd+U=
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/miekg/dns v1.1.4/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg=
@@ -208,7 +216,9 @@ github.com/nrdcg/goinwx v0.6.1/go.mod h1:XPiut7enlbEdntAqalBIqcYcTEVhpv/dKWgDCX2
github.com/nrdcg/namesilo v0.2.1/go.mod h1:lwMvfQTyYq+BbjJd30ylEG4GPSS6PII0Tia4rRpRiyw=
github.com/olekukonko/tablewriter v0.0.1/go.mod h1:vsDQFd/mU46D+Z4whnwzcISnGGzXWMclvtLoiIKAKIo=
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/gomega v1.4.3 h1:RE1xgDvH7imwFD45h+u2SgIfERHlS2yNG4DObb5BSKU=
github.com/onsi/gomega v1.4.3/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=
@@ -235,6 +245,8 @@ github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d/go.mod h1:c3At6R
github.com/prometheus/procfs v0.0.0-20190117184657-bf6a532e95b1/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
github.com/prometheus/procfs v0.0.2/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsTZCD3I8kEA=
github.com/prometheus/procfs v0.0.3/go.mod h1:4A/X28fw3Fc593LaREMrKMqOKvUAntwMDaekg4FpcdQ=
github.com/proullon/ramsql v0.0.0-20181213202341-817cee58a244 h1:fdX2U+a2Rmc4BjRYcOKzjYXtYTE4ga1B2lb8i7BlefU=
github.com/proullon/ramsql v0.0.0-20181213202341-817cee58a244/go.mod h1:jG8oAQG0ZPHPyxg5QlMERS31airDC+ZuqiAe8DUvFVo=
github.com/rainycape/memcache v0.0.0-20150622160815-1031fa0ce2f2/go.mod h1:7tZKcyumwBO6qip7RNQ5r77yrssm9bfCowcLEBcU5IA=
github.com/rcrowley/go-metrics v0.0.0-20181016184325-3113b8401b8a/go.mod h1:bCqnVzQkZxMG4s8nGwiZ5l3QUCyqpo9Y+/ZMZ9VjZe4=
github.com/refraction-networking/utls v0.0.0-20190909200633-43c36d3c1f57 h1:SL1K0QAuC1b54KoY1pjPWe6kSlsFHwK9/oC960fKrTY=
@@ -270,6 +282,8 @@ github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415/go.mod h1:
github.com/xeipuuv/gojsonschema v1.2.0/go.mod h1:anYRn/JVcOK2ZgGU+IjEV4nwlhoK5sQluxsYJ78Id3Y=
github.com/xtaci/smux v1.5.12 h1:n9OGjdqQuVZXLh46+L4IR5tR2wvuUFwRABnN/V55bIY=
github.com/xtaci/smux v1.5.12/go.mod h1:OMlQbT5vcgl2gb49mFkYo6SMf+zP3rcjcwQz7ZU7IGY=
github.com/ziutek/mymysql v1.5.4 h1:GB0qdRGsTwQSBVYuVShFBKaXSnSnYYC2d9knnE1LHFs=
github.com/ziutek/mymysql v1.5.4/go.mod h1:LMSpPZ6DbqWFxNCHW77HeMg9I646SAhApZ/wKdgO/C0=
go.opencensus.io v0.20.1/go.mod h1:6WKK9ahsWS3RSO+PY9ZHZUfv2irvY6gN279GOPZjmmk=
go.opencensus.io v0.20.2/go.mod h1:6WKK9ahsWS3RSO+PY9ZHZUfv2irvY6gN279GOPZjmmk=
go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU=
@@ -499,6 +513,7 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8
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/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=
gopkg.in/h2non/gock.v1 v1.0.15/go.mod h1:sX4zAkdYX1TRGJ2JY156cFspQn4yRWn6p9EMdODlynE=
gopkg.in/ini.v1 v1.42.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k=
@@ -508,6 +523,7 @@ gopkg.in/resty.v1 v1.9.1/go.mod h1:vo52Hzryw9PnPHcJfPsBiFW62XhNx5OczbV9y+IMpgc=
gopkg.in/resty.v1 v1.12.0/go.mod h1:mDo4pnntr5jdWRML875a/NmxYqAlA73dVijT2AXvQQo=
gopkg.in/square/go-jose.v2 v2.3.1 h1:SK5KegNXmKmqE342YYN2qPHEnUYeoMiXXl1poUlI+o4=
gopkg.in/square/go-jose.v2 v2.3.1/go.mod h1:M9dMgbHiYLoDGQrXy7OpJDJWiKiU//h+vD76mk0e1AI=
gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ=
gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw=
gopkg.in/yaml.v2 v2.0.0-20170812160011-eb3733d160e7/go.mod h1:JAlM8MvJe8wmxCU4Bli9HhUf9+ttbYbLASfIpnQbh74=
gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+2 -1
View File
@@ -8,7 +8,6 @@ import (
//the following modules are optional
//you can comment some of them if you don't need them
_ "github.com/go-sql-driver/mysql"
_ "github.com/p4gefau1t/trojan-go/cert"
_ "github.com/p4gefau1t/trojan-go/daemon"
_ "github.com/p4gefau1t/trojan-go/easy"
@@ -17,6 +16,8 @@ import (
_ "github.com/p4gefau1t/trojan-go/proxy/relay"
_ "github.com/p4gefau1t/trojan-go/proxy/server"
_ "github.com/p4gefau1t/trojan-go/router/mixed"
_ "github.com/p4gefau1t/trojan-go/stat/db"
_ "github.com/p4gefau1t/trojan-go/stat/memory"
_ "github.com/p4gefau1t/trojan-go/version"
//_ "github.com/p4gefau1t/trojan-go/log/simplelog"
)
-4
View File
@@ -72,10 +72,6 @@ type NeedAuth interface {
SetAuth(auth stat.Authenticator)
}
type NeedMeter interface {
SetMeter(meter stat.TrafficMeter)
}
type ConnSession interface {
io.ReadWriteCloser
}
-13
View File
@@ -8,18 +8,15 @@ import (
"github.com/p4gefau1t/trojan-go/conf"
"github.com/p4gefau1t/trojan-go/log"
"github.com/p4gefau1t/trojan-go/protocol"
"github.com/p4gefau1t/trojan-go/stat"
)
type SimpleSocksConnSession struct {
protocol.ConnSession
protocol.NeedMeter
config *conf.GlobalConfig
request *protocol.Request
rwc io.ReadWriteCloser
passwordHash string
meter stat.TrafficMeter
recv uint64
sent uint64
}
@@ -27,18 +24,12 @@ type SimpleSocksConnSession struct {
func (m *SimpleSocksConnSession) Read(p []byte) (int, error) {
n, err := m.rwc.Read(p)
m.recv += uint64(n)
if m.meter != nil {
m.meter.Count(m.passwordHash, 0, uint64(n))
}
return n, err
}
func (m *SimpleSocksConnSession) Write(p []byte) (int, error) {
n, err := m.rwc.Write(p)
m.sent += uint64(n)
if m.meter != nil {
m.meter.Count(m.passwordHash, uint64(n), 0)
}
return n, err
}
@@ -47,10 +38,6 @@ func (m *SimpleSocksConnSession) Close() error {
return m.rwc.Close()
}
func (m *SimpleSocksConnSession) SetMeter(meter stat.TrafficMeter) {
m.meter = meter
}
func (m *SimpleSocksConnSession) GetRequest() *protocol.Request {
return m.request
}
+5 -9
View File
@@ -15,8 +15,6 @@ import (
type TrojanInboundConnSession struct {
protocol.ConnSession
protocol.NeedAuth
protocol.NeedMeter
rwc io.ReadWriteCloser
ctx context.Context
@@ -32,19 +30,15 @@ type TrojanInboundConnSession struct {
func (i *TrojanInboundConnSession) Write(p []byte) (int, error) {
n, err := i.rwc.Write(p)
if i.meter != nil {
i.meter.Count(i.passwordHash, uint64(n), 0)
}
i.sent += uint64(n)
i.meter.Count(uint64(n), 0)
return n, err
}
func (i *TrojanInboundConnSession) Read(p []byte) (int, error) {
n, err := i.rwc.Read(p)
if i.meter != nil {
i.meter.Count(i.passwordHash, 0, uint64(n))
}
i.recv += uint64(n)
i.meter.Count(0, uint64(n))
return n, err
}
@@ -60,10 +54,12 @@ func (i *TrojanInboundConnSession) parseRequest(r *common.RewindReader) error {
if err != nil || n != 56 {
return common.NewError("failed to read hash").Base(err)
}
if !i.auth.CheckHash(string(userHash[:])) {
valid, meter := i.auth.AuthUser(string(userHash[:]))
if !valid {
return common.NewError("invalid hash:" + string(userHash[:]))
}
i.passwordHash = string(userHash[:])
i.meter = meter
crlf := [2]byte{}
r.Read(crlf[:])
+8 -13
View File
@@ -13,13 +13,13 @@ import (
type TrojanOutboundConnSession struct {
protocol.ConnSession
protocol.NeedMeter
config *conf.GlobalConfig
rwc io.ReadWriteCloser
request *protocol.Request
sent uint64
recv uint64
auth stat.Authenticator
meter stat.TrafficMeter
}
@@ -29,18 +29,14 @@ func (o *TrojanOutboundConnSession) SetMeter(meter stat.TrafficMeter) {
func (o *TrojanOutboundConnSession) Write(p []byte) (int, error) {
n, err := o.rwc.Write(p)
if o.meter != nil {
o.meter.Count("", uint64(n), 0)
}
o.meter.Count(uint64(n), 0)
o.sent += uint64(n)
return n, err
}
func (o *TrojanOutboundConnSession) Read(p []byte) (int, error) {
n, err := o.rwc.Read(p)
if o.meter != nil {
o.meter.Count("", 0, uint64(n))
}
o.meter.Count(0, uint64(n))
o.recv += uint64(n)
return n, err
}
@@ -51,11 +47,9 @@ func (o *TrojanOutboundConnSession) Close() error {
}
func (o *TrojanOutboundConnSession) writeRequest() error {
hash := ""
for k := range o.config.Hash {
hash = k
break
}
user := o.auth.ListUsers()[0]
hash := user.Hash()
o.meter = user
buf := bytes.NewBuffer(make([]byte, 0, 128))
crlf := []byte("\r\n")
buf.Write([]byte(hash))
@@ -67,11 +61,12 @@ func (o *TrojanOutboundConnSession) writeRequest() error {
return err
}
func NewOutboundConnSession(req *protocol.Request, rwc io.ReadWriteCloser, config *conf.GlobalConfig) (protocol.ConnSession, error) {
func NewOutboundConnSession(req *protocol.Request, rwc io.ReadWriteCloser, config *conf.GlobalConfig, auth stat.Authenticator) (protocol.ConnSession, error) {
o := &TrojanOutboundConnSession{
request: req,
config: config,
rwc: rwc,
auth: auth,
}
if err := o.writeRequest(); err != nil {
return nil, common.NewError("failed to write request").Base(err)
+55
View File
@@ -0,0 +1,55 @@
package client
import (
"context"
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/conf"
"github.com/p4gefau1t/trojan-go/log"
"github.com/p4gefau1t/trojan-go/protocol"
"github.com/p4gefau1t/trojan-go/protocol/simplesocks"
"github.com/p4gefau1t/trojan-go/protocol/trojan"
"github.com/p4gefau1t/trojan-go/stat"
)
type AppManager struct {
auth stat.Authenticator
config *conf.GlobalConfig
transport TransportManager
ctx context.Context
}
func (m *AppManager) OpenAppConn(req *protocol.Request) (protocol.ConnSession, error) {
var outboundConn protocol.ConnSession
//transport layer
transport, err := m.transport.DialToServer()
if err != nil {
return nil, common.NewError("failed to init transport layer").Base(err)
}
//application layer
if m.config.Mux.Enabled {
outboundConn, err = simplesocks.NewOutboundConnSession(req, transport)
} else {
outboundConn, err = trojan.NewOutboundConnSession(req, transport, m.config, m.auth)
}
if err != nil {
return nil, common.NewError("fail to start conn session").Base(err)
}
return outboundConn, nil
}
func NewAppManager(ctx context.Context, config *conf.GlobalConfig, auth stat.Authenticator) *AppManager {
c := &AppManager{
ctx: ctx,
config: config,
auth: auth,
}
if config.Mux.Enabled {
log.Info("mux enabled")
c.transport = NewMuxPoolManager(ctx, config, auth)
} else {
c.transport = NewTLSManager(config)
}
return c
}
+26 -45
View File
@@ -5,14 +5,12 @@ import (
"io"
"net"
"github.com/p4gefau1t/trojan-go/api"
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/conf"
"github.com/p4gefau1t/trojan-go/log"
"github.com/p4gefau1t/trojan-go/protocol"
"github.com/p4gefau1t/trojan-go/protocol/direct"
"github.com/p4gefau1t/trojan-go/protocol/http"
"github.com/p4gefau1t/trojan-go/protocol/simplesocks"
"github.com/p4gefau1t/trojan-go/protocol/socks"
"github.com/p4gefau1t/trojan-go/protocol/trojan"
"github.com/p4gefau1t/trojan-go/proxy"
@@ -38,29 +36,10 @@ type Client struct {
cancel context.CancelFunc
associated *common.Notifier
router router.Router
meter stat.TrafficMeter
transport TransportManager
tcpListener net.Listener
udpListener *net.UDPConn
}
func (c *Client) openOutboundConn(req *protocol.Request) (protocol.ConnSession, error) {
var outboundConn protocol.ConnSession
//transport layer
transport, err := c.transport.DialToServer()
if err != nil {
return nil, common.NewError("failed to init transport layer").Base(err)
}
//application layer
if c.config.Mux.Enabled {
outboundConn, err = simplesocks.NewOutboundConnSession(req, transport)
} else {
outboundConn, err = trojan.NewOutboundConnSession(req, transport, c.config)
}
if err != nil {
return nil, common.NewError("fail to start conn session").Base(err)
}
return outboundConn, nil
auth stat.Authenticator
appMan *AppManager
}
func (c *Client) handleSocksConn(conn io.ReadWriteCloser) {
@@ -128,13 +107,12 @@ func (c *Client) handleSocksConn(conn io.ReadWriteCloser) {
log.Info("[block] conn to", req)
return
}
outboundConn, err := c.openOutboundConn(req)
outboundConn, err := c.appMan.OpenAppConn(req)
if err != nil {
log.Error(err)
return
}
defer outboundConn.Close()
outboundConn.(protocol.NeedMeter).SetMeter(c.meter)
proxy.ProxyConn(c.ctx, inboundConn, outboundConn, c.config.BufferSize)
}
@@ -174,14 +152,13 @@ func (c *Client) handleHTTPConn(conn io.ReadWriteCloser) {
return
}
outboundConn, err := c.openOutboundConn(req)
outboundConn, err := c.appMan.OpenAppConn(req)
if err != nil {
log.Error(common.NewError("fail to start conn session").Base(err))
return
}
defer outboundConn.Close()
log.Info("conn tunneling to", req)
outboundConn.(protocol.NeedMeter).SetMeter(c.meter)
proxy.ProxyConn(c.ctx, inboundConn, outboundConn, c.config.BufferSize)
} else { //GET/POST requests
defer inboundPacket.Close()
@@ -219,7 +196,7 @@ func (c *Client) handleHTTPConn(conn io.ReadWriteCloser) {
case <-errChan:
return
case packet := <-packetChan:
outboundConn, err := c.openOutboundConn(packet.request)
outboundConn, err := c.appMan.OpenAppConn(req)
if err != nil {
log.Error(err)
continue
@@ -283,7 +260,7 @@ func (c *Client) listenUDP(errChan chan error) {
},
Command: protocol.Associate,
}
outboundConn, err := c.openOutboundConn(req)
outboundConn, err := c.appMan.OpenAppConn(req)
if err != nil {
log.Error(common.NewError("failed to init udp tunnel").Base(err))
return
@@ -342,9 +319,6 @@ func (c *Client) Run() error {
errChan := make(chan error, 2)
go c.listenUDP(errChan)
go c.listenTCP(errChan)
if c.config.API.Enabled {
go api.RunClientAPIService(c.ctx, c.config, c.meter)
}
select {
case err := <-errChan:
return err
@@ -366,26 +340,33 @@ func (c *Client) Close() error {
}
func (c *Client) Build(config *conf.GlobalConfig) (common.Runnable, error) {
c.ctx, c.cancel = context.WithCancel(context.Background())
c.associated = common.NewNotifier()
c.router = &router.EmptyRouter{}
c.meter = &stat.MemoryTrafficMeter{}
var err error
if config.Mux.Enabled {
log.Info("mux enabled")
c.transport = NewMuxPoolManager(c.ctx, config)
} else {
c.transport = NewTLSManager(config)
ctx, cancel := context.WithCancel(context.Background())
auth, err := stat.NewAuth(ctx, "memory", config)
if err != nil {
cancel()
return nil, err
}
var rtr router.Router = &router.EmptyRouter{}
if config.Router.Enabled {
log.Info("router enabled")
c.router, err = router.NewRouter(&config.Router)
rtr, err = router.NewRouter(&config.Router)
if err != nil {
log.Fatal(common.NewError("invalid router list").Base(err))
}
}
c.config = config
return c, nil
appMan := NewAppManager(ctx, config, auth)
newClient := &Client{
ctx: ctx,
cancel: cancel,
config: config,
router: rtr,
associated: common.NewNotifier(),
auth: auth,
appMan: appMan,
}
return newClient, nil
}
func init() {
+24 -33
View File
@@ -10,9 +10,9 @@ import (
"github.com/p4gefau1t/trojan-go/conf"
"github.com/p4gefau1t/trojan-go/log"
"github.com/p4gefau1t/trojan-go/protocol"
"github.com/p4gefau1t/trojan-go/protocol/simplesocks"
"github.com/p4gefau1t/trojan-go/protocol/trojan"
"github.com/p4gefau1t/trojan-go/proxy"
"github.com/p4gefau1t/trojan-go/stat"
)
type dispatchInfo struct {
@@ -32,26 +32,8 @@ type Forward struct {
outboundPacketTable map[string]protocol.PacketSession
udpListener *net.UDPConn
tcpListener net.Listener
transport TransportManager
}
func (f *Forward) openOutboundConn(req *protocol.Request) (protocol.ConnSession, error) {
var outboundConn protocol.ConnSession
//transport layer
transport, err := f.transport.DialToServer()
if err != nil {
return nil, common.NewError("failed to init transport layer").Base(err)
}
//application layer
if f.config.Mux.Enabled {
outboundConn, err = simplesocks.NewOutboundConnSession(req, transport)
} else {
outboundConn, err = trojan.NewOutboundConnSession(req, transport, f.config)
}
if err != nil {
return nil, common.NewError("fail to start conn session").Base(err)
}
return outboundConn, nil
auth stat.Authenticator
appMan *AppManager
}
func (f *Forward) dispatchServerPacket(addr *net.UDPAddr) {
@@ -109,7 +91,7 @@ func (f *Forward) dispatchClientPacket() {
f.outboundPacketTableLock.Lock()
outboundPacket, found := f.outboundPacketTable[packet.addr.String()]
if !found {
outboundConn, err := f.openOutboundConn(associateReq)
outboundConn, err := f.appMan.OpenAppConn(associateReq)
if err != nil {
log.Error(err)
continue
@@ -172,7 +154,7 @@ func (f *Forward) listenTCP(errChan chan error) {
errChan <- common.NewError("error occured when accepting conn").Base(err)
}
handle := func(inboundConn net.Conn) {
outboundConn, err := f.openOutboundConn(req)
outboundConn, err := f.appMan.OpenAppConn(req)
if err != nil {
log.Error(common.NewError("failed to start outbound session").Base(err))
return
@@ -210,17 +192,26 @@ func (f *Forward) Close() error {
}
func (f *Forward) Build(config *conf.GlobalConfig) (common.Runnable, error) {
f.ctx, f.cancel = context.WithCancel(context.Background())
if config.Mux.Enabled {
log.Info("mux enabled")
f.transport = NewMuxPoolManager(f.ctx, config)
} else {
f.transport = NewTLSManager(config)
ctx, cancel := context.WithCancel(context.Background())
authDriver := "memory"
auth, err := stat.NewAuth(ctx, authDriver, config)
if err != nil {
cancel()
return nil, err
}
f.clientPackets = make(chan *dispatchInfo, 512)
f.outboundPacketTable = make(map[string]protocol.PacketSession)
f.config = config
return f, nil
appMan := NewAppManager(ctx, config, auth)
newForward := &Forward{
ctx: ctx,
cancel: cancel,
config: config,
auth: auth,
appMan: appMan,
clientPackets: make(chan *dispatchInfo, 1024),
outboundPacketTable: make(map[string]protocol.PacketSession),
}
return newForward, nil
}
func init() {
+5 -2
View File
@@ -12,6 +12,7 @@ import (
"github.com/p4gefau1t/trojan-go/log"
"github.com/p4gefau1t/trojan-go/protocol"
"github.com/p4gefau1t/trojan-go/protocol/trojan"
"github.com/p4gefau1t/trojan-go/stat"
"github.com/xtaci/smux"
)
@@ -33,6 +34,7 @@ type MuxManager struct {
sync.Mutex
muxPool map[MuxID]*muxClientInfo
config *conf.GlobalConfig
auth stat.Authenticator
ctx context.Context
transport *TLSManager
}
@@ -53,7 +55,7 @@ func (m *MuxManager) newMuxClient() (*muxClientInfo, error) {
if err != nil {
return nil, common.NewError("failed to dail to remote server").Base(err)
}
conn, err := trojan.NewOutboundConnSession(req, rwc, m.config)
conn, err := trojan.NewOutboundConnSession(req, rwc, m.config, m.auth)
if err != nil {
rwc.Close()
log.Error(common.NewError("failed to dial tls tunnel").Base(err))
@@ -155,12 +157,13 @@ func (m *MuxManager) checkAndCloseIdleMuxClient() {
}
}
func NewMuxPoolManager(ctx context.Context, config *conf.GlobalConfig) *MuxManager {
func NewMuxPoolManager(ctx context.Context, config *conf.GlobalConfig, auth stat.Authenticator) *MuxManager {
m := &MuxManager{
ctx: ctx,
config: config,
muxPool: make(map[MuxID]*muxClientInfo),
transport: NewTLSManager(config),
auth: auth,
}
go m.checkAndCloseIdleMuxClient()
return m
+20 -31
View File
@@ -11,10 +11,10 @@ import (
"github.com/p4gefau1t/trojan-go/conf"
"github.com/p4gefau1t/trojan-go/log"
"github.com/p4gefau1t/trojan-go/protocol"
"github.com/p4gefau1t/trojan-go/protocol/simplesocks"
"github.com/p4gefau1t/trojan-go/protocol/tproxy"
"github.com/p4gefau1t/trojan-go/protocol/trojan"
"github.com/p4gefau1t/trojan-go/proxy"
"github.com/p4gefau1t/trojan-go/stat"
)
type NAT struct {
@@ -26,26 +26,8 @@ type NAT struct {
cancel context.CancelFunc
inboundPacket protocol.PacketSession
listener net.Listener
transport TransportManager
}
func (n *NAT) openOutboundConn(req *protocol.Request) (protocol.ConnSession, error) {
var outboundConn protocol.ConnSession
//transport layer
transport, err := n.transport.DialToServer()
if err != nil {
return nil, common.NewError("failed to init transport layer").Base(err)
}
//application layer
if n.config.Mux.Enabled {
outboundConn, err = simplesocks.NewOutboundConnSession(req, transport)
} else {
outboundConn, err = trojan.NewOutboundConnSession(req, transport, n.config)
}
if err != nil {
return nil, common.NewError("fail to start conn session").Base(err)
}
return outboundConn, nil
auth stat.Authenticator
appMan *AppManager
}
func (n *NAT) handleConn(conn net.Conn) {
@@ -55,7 +37,7 @@ func (n *NAT) handleConn(conn net.Conn) {
return
}
defer inboundConn.Close()
outboundConn, err := n.openOutboundConn(req)
outboundConn, err := n.appMan.OpenAppConn(req)
if err != nil {
log.Error(err)
return
@@ -81,7 +63,7 @@ func (n *NAT) listenUDP(errChan chan error) {
Command: protocol.Associate,
}
for {
outboundConn, err := n.openOutboundConn(req)
outboundConn, err := n.appMan.OpenAppConn(req)
if err != nil {
log.Error(err)
time.Sleep(time.Second)
@@ -150,15 +132,22 @@ func (n *NAT) Close() error {
}
func (n *NAT) Build(config *conf.GlobalConfig) (common.Runnable, error) {
n.ctx, n.cancel = context.WithCancel(context.Background())
n.config = config
if config.Mux.Enabled {
log.Info("mux enabled")
n.transport = NewMuxPoolManager(n.ctx, config)
} else {
n.transport = NewTLSManager(config)
ctx, cancel := context.WithCancel(context.Background())
auth, err := stat.NewAuth(ctx, "memory", config)
if err != nil {
cancel()
return nil, err
}
return n, nil
appMan := NewAppManager(ctx, config, auth)
newForward := &Forward{
ctx: ctx,
cancel: cancel,
config: config,
auth: auth,
appMan: appMan,
}
return newForward, nil
}
func init() {
+20 -27
View File
@@ -26,7 +26,6 @@ type Server struct {
listener net.Listener
auth stat.Authenticator
meter stat.TrafficMeter
config *conf.GlobalConfig
shadow *shadow.ShadowManager
ctx context.Context
@@ -40,7 +39,6 @@ func (s *Server) handleMuxConn(stream *smux.Stream) {
log.Error(common.NewError("cannot start inbound session").Base(err))
return
}
inboundConn.(protocol.NeedMeter).SetMeter(s.meter)
switch req.Command {
case protocol.Connect:
outboundConn, err := direct.NewOutboundConnSession(req)
@@ -85,7 +83,6 @@ func (s *Server) handleConn(conn *tls.Conn) {
go s.handleMuxConn(stream)
}
}
inboundConn.(protocol.NeedMeter).SetMeter(s.meter)
if req.Command == protocol.Associate {
inboundPacket, err := trojan.NewPacketSession(inboundConn)
@@ -117,29 +114,10 @@ func (s *Server) handleConn(conn *tls.Conn) {
}
func (s *Server) Run() error {
var err error
if s.config.MySQL.Enabled {
s.auth, err = stat.NewMixedAuthenticator(s.config)
if err != nil {
return common.NewError("failed to init auth").Base(err)
}
s.meter, err = stat.NewDBTrafficMeter(s.config)
if err != nil {
return common.NewError("failed to init traffic meter").Base(err)
}
} else {
s.auth = &stat.ConfigUserAuthenticator{
Config: s.config,
}
}
defer s.auth.Close()
if s.meter != nil {
defer s.meter.Close()
}
log.Info("server is running at", s.config.LocalAddress)
var listener net.Listener
listener, err = net.Listen("tcp", s.config.LocalAddress.String())
listener, err := net.Listen("tcp", s.config.LocalAddress.String())
if err != nil {
return err
}
@@ -211,10 +189,25 @@ func (s *Server) Close() error {
return nil
}
func (s *Server) Build(config *conf.GlobalConfig) (common.Runnable, error) {
s.config = config
s.ctx, s.cancel = context.WithCancel(context.Background())
s.shadow = shadow.NewShadowManager(s.ctx, s.config)
func (*Server) Build(config *conf.GlobalConfig) (common.Runnable, error) {
ctx, cancel := context.WithCancel(context.Background())
var err error
authDriver := "memory"
if config.MySQL.Enabled {
authDriver = "mysql"
}
auth, err := stat.NewAuth(ctx, authDriver, config)
if err != nil {
cancel()
return nil, err
}
s := &Server{
config: config,
ctx: ctx,
cancel: cancel,
shadow: shadow.NewShadowManager(ctx, config),
auth: auth,
}
return s, nil
}
-224
View File
@@ -1,224 +0,0 @@
package stat
import (
"context"
"database/sql"
"fmt"
"strings"
"sync"
"time"
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/conf"
"github.com/p4gefau1t/trojan-go/log"
)
type trafficInfo struct {
passwordHash string
recv uint64
sent uint64
}
type DBTrafficMeter struct {
TrafficMeter
db *sql.DB
trafficChan chan *trafficInfo
ctx context.Context
cancel context.CancelFunc
updateDuration time.Duration
}
func (c *DBTrafficMeter) Query(passwordHash string) (uint64, uint64) {
// TODO Query method
return 0, 0
}
func (c *DBTrafficMeter) Count(passwordHash string, sent uint64, recv uint64) {
c.trafficChan <- &trafficInfo{
passwordHash: passwordHash,
sent: sent,
recv: recv,
}
}
func (c *DBTrafficMeter) Close() error {
c.cancel()
return c.db.Close()
}
func (c *DBTrafficMeter) dbDaemon() {
for {
beginTime := time.Now()
statBuffer := make(map[string]*trafficInfo)
for {
select {
case u := <-c.trafficChan:
t, found := statBuffer[u.passwordHash]
if !found {
t = &trafficInfo{
passwordHash: u.passwordHash,
}
statBuffer[u.passwordHash] = t
}
t.sent += u.sent
t.recv += u.recv
case <-time.After(c.updateDuration):
break
case <-c.ctx.Done():
return
}
if time.Now().Sub(beginTime) > c.updateDuration {
break
}
}
if len(statBuffer) == 0 {
continue
}
tx, err := c.db.Begin()
if err != nil {
log.Error(common.NewError("cannot begin transactin").Base(err))
continue
}
for _, traffic := range statBuffer {
//swap upload and download for users
s, err := tx.Prepare("UPDATE users SET upload=upload+? WHERE password=?;")
common.Must(err)
_, err = s.Exec(traffic.recv, traffic.passwordHash)
s, err = tx.Prepare("UPDATE users SET download=download+? WHERE password=?;")
common.Must(err)
_, err = s.Exec(traffic.sent, traffic.passwordHash)
if err != nil {
log.Error(common.NewError("failed to update data to tx").Base(err))
break
}
}
err = tx.Commit()
if err != nil {
log.Error(common.NewError("failed to commit tx").Base(err))
} else {
log.Info("buffered data has been written into the database")
}
}
}
func NewDBTrafficMeter(config *conf.GlobalConfig) (TrafficMeter, error) {
db, err := connectDatabase(
"mysql",
config.MySQL.Username,
config.MySQL.Password,
config.MySQL.ServerHost,
config.MySQL.ServerPort,
config.MySQL.Database,
)
if err != nil {
return nil, common.NewError("failed to connect to database server").Base(err)
}
c := &DBTrafficMeter{
db: db,
trafficChan: make(chan *trafficInfo, 1024*8),
ctx: context.Background(),
updateDuration: time.Duration(config.MySQL.CheckRate) * time.Second,
}
go c.dbDaemon()
return c, nil
}
type userInfo struct {
username string
passwordHash string
download uint64
upload uint64
quota uint64
}
type DBAuthenticator struct {
db *sql.DB
validUsers sync.Map
ctx context.Context
cancel context.CancelFunc
updateDuration time.Duration
Authenticator
}
func (a *DBAuthenticator) CheckHash(hash string) bool {
_, ok := a.validUsers.Load(hash)
if !ok {
return false
}
return true
}
func (a *DBAuthenticator) updateDaemon() {
for {
rows, err := a.db.Query("SELECT password,quota,download,upload FROM users")
if err != nil {
log.Error(common.NewError("failed to pull data from the database").Base(err))
time.Sleep(a.updateDuration)
continue
}
newValidUsers := make(map[string]string)
for rows.Next() {
var passwordHash string
var quota, download, upload int64
err := rows.Scan(&passwordHash, &quota, &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 {
newValidUsers[passwordHash] = ""
}
}
//delete those out of quota
a.validUsers.Range(func(key interface{}, val interface{}) bool {
if _, found := newValidUsers[key.(string)]; !found {
a.validUsers.Delete(key)
}
return true
})
for k, v := range newValidUsers {
a.validUsers.Store(k, v)
}
select {
case <-time.After(a.updateDuration):
break
case <-a.ctx.Done():
return
}
}
}
func (a *DBAuthenticator) Close() error {
a.cancel()
return a.db.Close()
}
func NewDBAuthenticator(config *conf.GlobalConfig) (Authenticator, error) {
db, err := connectDatabase(
"mysql",
config.MySQL.Username,
config.MySQL.Password,
config.MySQL.ServerHost,
config.MySQL.ServerPort,
config.MySQL.Database,
)
if err != nil {
return nil, common.NewError("failed to connect to database server").Base(err)
}
ctx, cancel := context.WithCancel(context.Background())
a := &DBAuthenticator{
db: db,
cancel: cancel,
ctx: ctx,
updateDuration: time.Duration(config.MySQL.CheckRate) * time.Second,
}
go a.updateDaemon()
return a, nil
}
func connectDatabase(driverName, username, password, ip string, port int, dbName string) (*sql.DB, error) {
path := strings.Join([]string{username, ":", password, "@tcp(", ip, ":", fmt.Sprintf("%d", port), ")/", dbName, "?charset=utf8"}, "")
return sql.Open(driverName, path)
}
+114
View File
@@ -0,0 +1,114 @@
package db
import (
"context"
"database/sql"
"fmt"
"strings"
"time"
_ "github.com/go-sql-driver/mysql"
"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 *sql.DB
updateDuration time.Duration
ctx context.Context
}
func (a *DBAuth) updater() {
for {
users := a.ListUsers()
tx, err := a.db.Begin()
if err != nil {
log.Error(common.NewError("cannot begin transaction").Base(err))
continue
}
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()
_, err = s.Exec(recv, hash)
s, err = tx.Prepare("UPDATE users SET download=download+? WHERE password=?;")
common.Must(err)
_, err = s.Exec(sent, hash)
if err != nil {
log.Error(common.NewError("failed to update data to tx").Base(err))
break
}
}
err = tx.Commit()
if err != nil {
log.Error(common.NewError("failed to commit tx").Base(err))
}
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 {
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, &quota, &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("db daemon exiting...")
return
}
}
}
func connectDatabase(driverName, username, password, ip string, port int, dbName string) (*sql.DB, error) {
path := strings.Join([]string{username, ":", password, "@tcp(", ip, ":", fmt.Sprintf("%d", port), ")/", dbName, "?charset=utf8"}, "")
return sql.Open(driverName, path)
}
func NewDBAuth(ctx context.Context, config *conf.GlobalConfig) (stat.Authenticator, error) {
db, err := connectDatabase(
"mysql",
config.MySQL.Username,
config.MySQL.Password,
config.MySQL.ServerHost,
config.MySQL.ServerPort,
config.MySQL.Database,
)
if err != nil {
return nil, common.NewError("failed to connect to database server").Base(err)
}
a := &DBAuth{
db: db,
ctx: ctx,
updateDuration: time.Duration(config.MySQL.CheckRate) * time.Second,
}
go a.updater()
return a, nil
}
func init() {
stat.RegisterAuthCreator("mysql", NewDBAuth)
}
+42
View File
@@ -0,0 +1,42 @@
package db
import (
"context"
"database/sql"
"testing"
"time"
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/conf"
"github.com/p4gefau1t/trojan-go/stat/memory"
_ "github.com/proullon/ramsql/driver"
)
func TestDBAuth(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
db, err := sql.Open("ramsql", "TestLoadUserAddresses")
common.Must(err)
common.Must2(db.Exec(`
CREATE TABLE users (
password CHAR(56) NOT NULL,
quota BIGINT NOT NULL DEFAULT 0,
download BIGINT NOT NULL DEFAULT 0,
upload BIGINT NOT NULL DEFAULT 0,
);
`))
common.Must2(db.Exec(`INSERT INTO users (password, quota, download, upload) VALUES ("hashhash", 20000, 0, 0);`))
memoryAuth, err := memory.NewMemoryAuth(ctx, &conf.GlobalConfig{})
auth := &DBAuth{
db: db,
ctx: ctx,
updateDuration: time.Second,
MemoryAuthenticator: memoryAuth.(*memory.MemoryAuthenticator),
}
go auth.updater()
time.Sleep(time.Second * 5)
valid, _ := auth.AuthUser("hashhash")
if !valid {
t.Fail()
}
cancel()
}
-69
View File
@@ -1,69 +0,0 @@
package stat
import (
"context"
"database/sql"
"fmt"
"math/rand"
"strings"
"testing"
"time"
_ "github.com/go-sql-driver/mysql"
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/conf"
)
func TestDBTrafficMeter(t *testing.T) {
userName := "root"
password := "password"
ip := "127.0.0.1"
port := "3306"
dbName := "trojan"
path := strings.Join([]string{userName, ":", password, "@tcp(", ip, ":", port, ")/", dbName, "?charset=utf8"}, "")
db, err := sql.Open("mysql", path)
hash := common.SHA224String("hashhash")
common.Must(err)
defer db.Close()
c := &DBTrafficMeter{
db: db,
trafficChan: make(chan *trafficInfo, 1024),
ctx: context.Background(),
updateDuration: time.Second * 5,
}
simulation := func() {
for i := 0; i < 100; i++ {
c.Count(hash, uint64(rand.Intn(500)), uint64(rand.Intn(500)))
time.Sleep(time.Duration(int64(time.Millisecond) * rand.Int63n(100)))
}
fmt.Println("done")
}
for i := 0; i < 100; i++ {
go simulation()
}
go c.dbDaemon()
time.Sleep(time.Second * 30)
}
func TestDBAuthenticator(t *testing.T) {
userName := "root"
password := "password"
ip := "127.0.0.1"
port := "3306"
dbName := "trojan"
path := strings.Join([]string{userName, ":", password, "@tcp(", ip, ":", port, ")/", dbName, "?charset=utf8"}, "")
db, err := sql.Open("mysql", path)
common.Must(err)
defer db.Close()
config := conf.GlobalConfig{
MySQL: conf.MySQLConfig{
CheckRate: 2,
},
}
a, err := NewDBAuthenticator(&config)
common.Must(err)
time.Sleep(time.Second * 5)
hash := common.SHA224String("hashhash")
fmt.Println(common.SHA224String("hashhash"))
fmt.Println(a.CheckHash(hash), a.CheckHash("jasdlkflfejlqjef"))
}
-56
View File
@@ -1,56 +0,0 @@
package stat
import (
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/conf"
)
type ConfigUserAuthenticator struct {
Authenticator
Config *conf.GlobalConfig
}
func (a *ConfigUserAuthenticator) CheckHash(hash string) bool {
_, found := a.Config.Hash[hash]
return found
}
func (a *ConfigUserAuthenticator) Close() error {
return nil
}
type MixedAuthenticator struct {
dbAuth Authenticator
configAuth Authenticator
Authenticator
}
func (a *MixedAuthenticator) CheckHash(hash string) bool {
if a.configAuth.CheckHash(hash) {
return true
} else if a.dbAuth.CheckHash(hash) {
return true
}
return false
}
func (a *MixedAuthenticator) Close() error {
return a.dbAuth.Close()
}
func NewMixedAuthenticator(config *conf.GlobalConfig) (Authenticator, error) {
if config.MySQL.Enabled {
dbAuth, err := NewDBAuthenticator(config)
common.Must(err)
a := &MixedAuthenticator{
configAuth: &ConfigUserAuthenticator{
Config: config,
},
dbAuth: dbAuth,
}
return a, nil
}
return &ConfigUserAuthenticator{
Config: config,
}, nil
}
-18
View File
@@ -1,18 +0,0 @@
package stat
import "sync/atomic"
type MemoryTrafficMeter struct {
TrafficMeter
sent uint64
recv uint64
}
func (m *MemoryTrafficMeter) Count(passwordHash string, sent, recv uint64) {
atomic.AddUint64(&m.sent, uint64(sent))
atomic.AddUint64(&m.recv, uint64(recv))
}
func (m *MemoryTrafficMeter) Query(passwordHash string) (uint64, uint64) {
return m.sent, m.recv
}
+105
View File
@@ -0,0 +1,105 @@
package memory
import (
"context"
"sync"
"sync/atomic"
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/conf"
"github.com/p4gefau1t/trojan-go/stat"
)
type MemoryTrafficMeter struct {
stat.TrafficMeter
sent uint64
recv uint64
hash string
}
func (m *MemoryTrafficMeter) Count(sent, recv uint64) {
atomic.AddUint64(&m.sent, uint64(sent))
atomic.AddUint64(&m.recv, uint64(recv))
}
func (m *MemoryTrafficMeter) Hash() string {
return m.hash
}
func (m *MemoryTrafficMeter) Get() (uint64, uint64) {
return atomic.LoadUint64(&m.sent), atomic.LoadUint64(&m.recv)
}
func (m *MemoryTrafficMeter) Reset() {
atomic.StoreUint64(&m.sent, 0)
atomic.StoreUint64(&m.recv, 0)
}
func (m *MemoryTrafficMeter) GetAndReset() (uint64, uint64) {
sent := atomic.SwapUint64(&m.sent, 0)
recv := atomic.SwapUint64(&m.recv, 0)
return sent, recv
}
type MemoryAuthenticator struct {
stat.Authenticator
sync.RWMutex
users map[string]*MemoryTrafficMeter
}
func (a *MemoryAuthenticator) AuthUser(hash string) (bool, stat.TrafficMeter) {
a.RLock()
defer a.RUnlock()
if user, found := a.users[hash]; found {
return true, user
}
return false, nil
}
func (a *MemoryAuthenticator) AddUser(hash string) error {
a.Lock()
defer a.Unlock()
if _, found := a.users[hash]; found {
return common.NewError("hash " + hash + " is already exist")
}
a.users[hash] = &MemoryTrafficMeter{
hash: hash,
}
return nil
}
func (a *MemoryAuthenticator) DelUser(hash string) error {
a.Lock()
defer a.Unlock()
_, found := a.users[hash]
if !found {
return common.NewError("hash " + hash + "is not exist")
}
delete(a.users, hash)
return nil
}
func (a *MemoryAuthenticator) ListUsers() []stat.TrafficMeter {
a.RLock()
defer a.RUnlock()
result := make([]stat.TrafficMeter, 0, len(a.users))
for _, m := range a.users {
result = append(result, m)
}
return result
}
func NewMemoryAuth(ctx context.Context, config *conf.GlobalConfig) (stat.Authenticator, error) {
a := &MemoryAuthenticator{
users: make(map[string]*MemoryTrafficMeter),
}
for hash := range config.Hash {
a.AddUser(hash)
}
return a, nil
}
func init() {
stat.RegisterAuthCreator("memory", NewMemoryAuth)
}
+30
View File
@@ -0,0 +1,30 @@
package memory
import (
"context"
"testing"
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/conf"
)
func TestMemoryAuth(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.Count(1234, 5678)
sent, recv := traffic.Get()
if sent != 1234 || recv != 5678 {
t.Fail()
}
cancel()
}
+30 -3
View File
@@ -1,17 +1,44 @@
package stat
import (
"context"
"io"
"github.com/p4gefau1t/trojan-go/common"
"github.com/p4gefau1t/trojan-go/conf"
)
type TrafficMeter interface {
io.Closer
Count(passwordHash string, sent uint64, recv uint64)
Query(passwordHash string) (sent uint64, recv uint64)
Hash() string
Count(sent uint64, recv uint64)
Get() (sent uint64, recv uint64)
Reset()
GetAndReset() (sent uint64, recv uint64)
GetSpeed() (sent uint64, recv uint64)
LimitSpeed(sent uint64, recv uint64)
}
type Authenticator interface {
io.Closer
AuthUser(hash string) (bool, TrafficMeter)
AddUser(hash string) error
DelUser(hash string) error
ListUsers() []TrafficMeter
}
CheckHash(hash string) bool
type AuthCreator func(ctx context.Context, config *conf.GlobalConfig) (Authenticator, error)
var authCreators = map[string]AuthCreator{}
func RegisterAuthCreator(name string, creator AuthCreator) {
authCreators[name] = creator
}
func NewAuth(ctx context.Context, name string, config *conf.GlobalConfig) (Authenticator, error) {
creator, found := authCreators[name]
if !found {
return nil, common.NewError("driver name " + name + " not found")
}
return creator(ctx, config)
}
+14 -10
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/memory"
"golang.org/x/net/proxy"
"golang.org/x/net/websocket"
)
@@ -183,26 +184,29 @@ func addTCPOption(config *conf.GlobalConfig) *conf.GlobalConfig {
func RunClient(ctx context.Context, config *conf.GlobalConfig) {
c := client.Client{}
common.Must2(c.Build(config))
go c.Run()
r, err := c.Build(config)
common.Must(err)
go r.Run()
<-ctx.Done()
c.Close()
r.Close()
}
func RunForward(ctx context.Context, config *conf.GlobalConfig) {
f := client.Forward{}
common.Must2(f.Build(config))
go f.Run()
c := client.Forward{}
r, err := c.Build(config)
common.Must(err)
go r.Run()
<-ctx.Done()
f.Close()
r.Close()
}
func RunServer(ctx context.Context, config *conf.GlobalConfig) {
s := server.Server{}
common.Must2(s.Build(config))
go s.Run()
r, err := s.Build(config)
common.Must(err)
go r.Run()
<-ctx.Done()
s.Close()
r.Close()
}
func CheckClientServer(t *testing.T, clientConfig *conf.GlobalConfig, serverConfig *conf.GlobalConfig) {