diff --git a/api/client.go b/api/client.go index a9d9fc5..83908c5 100644 --- a/api/client.go +++ b/api/client.go @@ -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 { diff --git a/api/client_test.go b/api/client_test.go index 9f4405f..33a6348 100644 --- a/api/client_test.go +++ b/api/client_test.go @@ -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() } diff --git a/go.mod b/go.mod index 677724b..895740e 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index a2b6479..88d3c30 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/main.go b/main.go index a47c1a5..9abcd0b 100644 --- a/main.go +++ b/main.go @@ -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" ) diff --git a/protocol/protocol.go b/protocol/protocol.go index 4aa61a5..e664609 100644 --- a/protocol/protocol.go +++ b/protocol/protocol.go @@ -72,10 +72,6 @@ type NeedAuth interface { SetAuth(auth stat.Authenticator) } -type NeedMeter interface { - SetMeter(meter stat.TrafficMeter) -} - type ConnSession interface { io.ReadWriteCloser } diff --git a/protocol/simplesocks/simplesocks.go b/protocol/simplesocks/simplesocks.go index 7820f8d..46429bc 100644 --- a/protocol/simplesocks/simplesocks.go +++ b/protocol/simplesocks/simplesocks.go @@ -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 } diff --git a/protocol/trojan/inbound.go b/protocol/trojan/inbound.go index 3d26867..5716f33 100644 --- a/protocol/trojan/inbound.go +++ b/protocol/trojan/inbound.go @@ -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[:]) diff --git a/protocol/trojan/outbound.go b/protocol/trojan/outbound.go index 14c3e72..d5da1eb 100644 --- a/protocol/trojan/outbound.go +++ b/protocol/trojan/outbound.go @@ -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) diff --git a/proxy/client/app.go b/proxy/client/app.go new file mode 100644 index 0000000..cd1a02f --- /dev/null +++ b/proxy/client/app.go @@ -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 +} diff --git a/proxy/client/client.go b/proxy/client/client.go index d456a77..1acf57b 100644 --- a/proxy/client/client.go +++ b/proxy/client/client.go @@ -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() { diff --git a/proxy/client/forward.go b/proxy/client/forward.go index 5a6151f..0295e90 100644 --- a/proxy/client/forward.go +++ b/proxy/client/forward.go @@ -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() { diff --git a/proxy/client/mux.go b/proxy/client/mux.go index 66b3d7e..a100045 100644 --- a/proxy/client/mux.go +++ b/proxy/client/mux.go @@ -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 diff --git a/proxy/client/nat.go b/proxy/client/nat.go index 01724f7..03d01c6 100644 --- a/proxy/client/nat.go +++ b/proxy/client/nat.go @@ -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() { diff --git a/proxy/server/server.go b/proxy/server/server.go index 1dd0095..19d978d 100644 --- a/proxy/server/server.go +++ b/proxy/server/server.go @@ -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 } diff --git a/stat/db.go b/stat/db.go deleted file mode 100644 index aef6e18..0000000 --- a/stat/db.go +++ /dev/null @@ -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, "a, &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) -} diff --git a/stat/db/db.go b/stat/db/db.go new file mode 100644 index 0000000..463846d --- /dev/null +++ b/stat/db/db.go @@ -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, "a, &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) +} diff --git a/stat/db/db_test.go b/stat/db/db_test.go new file mode 100644 index 0000000..3f396ae --- /dev/null +++ b/stat/db/db_test.go @@ -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() +} diff --git a/stat/db_test.go b/stat/db_test.go deleted file mode 100644 index b44c025..0000000 --- a/stat/db_test.go +++ /dev/null @@ -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")) -} diff --git a/stat/generic.go b/stat/generic.go deleted file mode 100644 index b48e827..0000000 --- a/stat/generic.go +++ /dev/null @@ -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 -} diff --git a/stat/memory.go b/stat/memory.go deleted file mode 100644 index 47a9117..0000000 --- a/stat/memory.go +++ /dev/null @@ -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 -} diff --git a/stat/memory/memory.go b/stat/memory/memory.go new file mode 100644 index 0000000..feff55d --- /dev/null +++ b/stat/memory/memory.go @@ -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) +} diff --git a/stat/memory/meomry_test.go b/stat/memory/meomry_test.go new file mode 100644 index 0000000..76c6532 --- /dev/null +++ b/stat/memory/meomry_test.go @@ -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() +} diff --git a/stat/stat.go b/stat/stat.go index 66ffbda..0c7e669 100644 --- a/stat/stat.go +++ b/stat/stat.go @@ -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) } diff --git a/test/proxy_test.go b/test/proxy_test.go index a563023..9b246bf 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/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) {