add api service

api/api.go

@@ -0,0 +1,24 @@
+package api
+
+import (
+	"context"
+	"github.com/p4gefau1t/trojan-go/log"
+	"github.com/p4gefau1t/trojan-go/statistic"
+)
+
+type Handler func(ctx context.Context, auth statistic.Authenticator) error
+
+var handlers = map[string]Handler{}
+
+func RegisterHandler(name string, handler Handler) {
+	handlers[name] = handler
+}
+
+func RunService(ctx context.Context, name string, auth statistic.Authenticator) error {
+	if h, ok := handlers[name]; ok {
+		log.Debug("api handler found", name)
+		return h(ctx, auth)
+	}
+	log.Debug("api handler not found", name)
+	return nil
+}

api/service/api.proto

@@ -0,0 +1,88 @@
+syntax = "proto3";
+package trojan.api;
+option go_package = ".;service";
+
+message Traffic {
+    uint64 upload_traffic = 1;
+    uint64 download_traffic = 2;
+}
+
+message Speed {
+    uint64 upload_speed = 1;
+    uint64 download_speed = 2;
+}
+
+message User {
+    string password = 1;
+    string hash = 2; //optional
+}
+
+message UserStatus {
+    User user = 1;
+    Traffic traffic_total = 2;
+    Speed speed_current = 3;
+    Speed speed_limit = 4;
+    int32 ip_current = 5;
+    int32 ip_limit = 6;
+}
+
+message GetTrafficRequest {
+    User user = 1;
+}
+
+message GetTrafficResponse {
+    bool success = 1;
+    string info = 2;
+    Traffic traffic_total = 3;
+    Speed speed_current = 4;
+}
+
+message ListUsersRequest {
+
+}
+
+message ListUsersResponse {
+    User user = 1;
+    UserStatus status = 2;
+}
+
+message GetUsersRequest {
+    User user = 1;
+}
+
+message GetUsersResponse {
+    bool success = 1;
+    string info = 2;
+    User user = 3;
+    UserStatus status = 4;
+}
+
+message SetUsersRequest {
+    User user = 1;
+    enum Operation {
+        Add = 0;
+        Delete = 1;
+        Modify = 2;
+    }
+    Operation operation = 2;
+    Speed speed_limit = 3;
+    int32 ip_limit = 4;
+}
+
+message SetUsersResponse {
+    bool success = 1;
+    string info = 2;
+}
+
+service TrojanClientService {
+    rpc GetTraffic(GetTrafficRequest) returns(GetTrafficResponse){}
+}
+
+service TrojanServerService {
+    // list all users
+    rpc ListUsers(ListUsersRequest) returns(stream ListUsersResponse){}
+    // obtain specified user's info
+    rpc GetUsers(stream GetUsersRequest) returns(stream GetUsersResponse){}
+    // setup exsisting users' config
+    rpc SetUsers(stream SetUsersRequest) returns(stream SetUsersResponse){}
+}

api/service/client.go

@@ -0,0 +1,83 @@
+package service
+
+import (
+	"context"
+	"fmt"
+	"github.com/p4gefau1t/trojan-go/api"
+	"github.com/p4gefau1t/trojan-go/common"
+	"github.com/p4gefau1t/trojan-go/config"
+	"github.com/p4gefau1t/trojan-go/log"
+	"github.com/p4gefau1t/trojan-go/statistic"
+	"github.com/p4gefau1t/trojan-go/tunnel/trojan"
+	"google.golang.org/grpc"
+	"net"
+)
+
+type ClientAPI struct {
+	TrojanClientServiceServer
+
+	auth          statistic.Authenticator
+	ctx           context.Context
+	uploadSpeed   uint64
+	downloadSpeed uint64
+	lastSent      uint64
+	lastRecv      uint64
+}
+
+func (s *ClientAPI) GetTraffic(ctx context.Context, req *GetTrafficRequest) (*GetTrafficResponse, error) {
+	log.Debug("API: GetTraffic")
+	if req.User == nil {
+		return nil, common.NewError("User is unspecified")
+	}
+	if req.User.Hash == "" {
+		req.User.Hash = common.SHA224String(req.User.Password)
+	}
+	valid, user := s.auth.AuthUser(req.User.Hash)
+	if !valid {
+		return nil, common.NewError("User " + req.User.Hash + " not found")
+	}
+	sent, recv := user.GetTraffic()
+	sentSpeed, recvSpeed := user.GetSpeed()
+	resp := &GetTrafficResponse{
+		Success: true,
+		TrafficTotal: &Traffic{
+			UploadTraffic:   sent,
+			DownloadTraffic: recv,
+		},
+		SpeedCurrent: &Speed{
+			UploadSpeed:   sentSpeed,
+			DownloadSpeed: recvSpeed,
+		},
+	}
+	return resp, nil
+}
+
+func RunClientAPI(ctx context.Context, auth statistic.Authenticator) error {
+	server := grpc.NewServer()
+	service := &ClientAPI{
+		ctx:  ctx,
+		auth: auth,
+	}
+	RegisterTrojanClientServiceServer(server, service)
+	cfg := config.FromContext(ctx, Name).(*Config)
+	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", cfg.APIHost, cfg.APIPort))
+	if err != nil {
+		return err
+	}
+	log.Info("client-side api service is listening on", listener.Addr().String())
+	errChan := make(chan error, 1)
+	go func() {
+		errChan <- server.Serve(listener)
+	}()
+	select {
+	case err := <-errChan:
+		return err
+	case <-ctx.Done():
+		server.Stop()
+		return nil
+	}
+}
+
+func init() {
+	api.RegisterHandler(trojan.Name+"_CLIENT", RunClientAPI)
+}

api/service/client_test.go

@@ -0,0 +1,53 @@
+package service
+
+import (
+	"context"
+	"fmt"
+	"github.com/p4gefau1t/trojan-go/common"
+	"github.com/p4gefau1t/trojan-go/config"
+	"github.com/p4gefau1t/trojan-go/statistic/memory"
+	"google.golang.org/grpc"
+	"testing"
+	"time"
+)
+
+func TestClientAPI(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	ctx = config.WithConfig(ctx, memory.Name,
+		&memory.Config{
+			Passwords: []string{"useless"},
+		})
+	port := common.PickPort("tcp", "127.0.0.1")
+	ctx = config.WithConfig(ctx, Name, &Config{
+		APIConfig{
+			Enabled: true,
+			APIHost: "127.0.0.1",
+			APIPort: port,
+		},
+	})
+	auth, err := memory.NewAuthenticator(ctx)
+	common.Must(err)
+	go RunClientAPI(ctx, auth)
+	common.Must(auth.AddUser("hash1234"))
+	valid, user := auth.AuthUser("hash1234")
+	if !valid {
+		t.Fail()
+	}
+	user.AddTraffic(1234, 5678)
+	time.Sleep(time.Second)
+	conn, err := grpc.Dial(fmt.Sprintf("127.0.0.1:%d", port), grpc.WithInsecure())
+	common.Must(err)
+	client := NewTrojanClientServiceClient(conn)
+	resp, err := client.GetTraffic(ctx, &GetTrafficRequest{User: &User{
+		Hash: "hash1234",
+	}})
+	common.Must(err)
+	if resp.TrafficTotal.DownloadTraffic != 5678 || resp.TrafficTotal.UploadTraffic != 1234 {
+		t.Fail()
+	}
+	resp, err = client.GetTraffic(ctx, &GetTrafficRequest{})
+	if err == nil {
+		t.Fail()
+	}
+	cancel()
+}

api/service/config.go

@@ -0,0 +1,21 @@
+package service
+
+import "github.com/p4gefau1t/trojan-go/config"
+
+const Name = "API_SERVICE"
+
+type APIConfig struct {
+	Enabled bool   `json:"enabled" yaml:"enabled"`
+	APIHost string `json:"api_addr" yaml:"api-addr"`
+	APIPort int    `json:"api_port" yaml:"api-port"`
+}
+
+type Config struct {
+	APIConfig `json,yaml:"api"`
+}
+
+func init() {
+	config.RegisterConfigCreator(Name, func() interface{} {
+		return new(Config)
+	})
+}

api/service/gen.sh

@@ -0,0 +1 @@
+protoc ./api.proto --go_out=plugins=grpc:.

api/service/server.go

@@ -0,0 +1,195 @@
+package service
+
+import (
+	"context"
+	"fmt"
+	"github.com/p4gefau1t/trojan-go/api"
+	"github.com/p4gefau1t/trojan-go/common"
+	"github.com/p4gefau1t/trojan-go/config"
+	"github.com/p4gefau1t/trojan-go/log"
+	"github.com/p4gefau1t/trojan-go/statistic"
+	"github.com/p4gefau1t/trojan-go/tunnel/trojan"
+	"google.golang.org/grpc"
+	"io"
+	"net"
+)
+
+type ServerAPI struct {
+	TrojanServerServiceServer
+	auth statistic.Authenticator
+}
+
+func (s *ServerAPI) GetUsers(stream TrojanServerService_GetUsersServer) error {
+	log.Debug("API: GetUsers")
+	for {
+		req, err := stream.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		if req.User == nil {
+			return common.NewError("User is unspecified")
+		}
+		if req.User.Hash == "" {
+			req.User.Hash = common.SHA224String(req.User.Password)
+		}
+		valid, user := s.auth.AuthUser(req.User.Hash)
+		if !valid {
+			stream.Send(&GetUsersResponse{
+				Success: false,
+				Info:    "Invalid user: " + req.User.Hash,
+			})
+			continue
+		}
+		downloadTraffic, uploadTraffic := user.GetTraffic()
+		downloadSpeed, uploadSpeed := user.GetSpeed()
+		downloadSpeedLimit, uploadSpeedLimit := user.GetSpeedLimit()
+		ipLimit := user.GetIPLimit()
+		ipCurrent := user.GetIP()
+		err = stream.Send(&GetUsersResponse{
+			Success: true,
+			Status: &UserStatus{
+				User: req.User,
+				TrafficTotal: &Traffic{
+					UploadTraffic:   uploadTraffic,
+					DownloadTraffic: downloadTraffic,
+				},
+				SpeedCurrent: &Speed{
+					DownloadSpeed: downloadSpeed,
+					UploadSpeed:   uploadSpeed,
+				},
+				SpeedLimit: &Speed{
+					DownloadSpeed: uint64(downloadSpeedLimit),
+					UploadSpeed:   uint64(uploadSpeedLimit),
+				},
+				IpCurrent: int32(ipCurrent),
+				IpLimit:   int32(ipLimit),
+			},
+		})
+		if err != nil {
+			return err
+		}
+	}
+}
+
+func (s *ServerAPI) SetUsers(stream TrojanServerService_SetUsersServer) error {
+	log.Debug("API: SetUsers")
+	for {
+		req, err := stream.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		if req.User == nil {
+			return common.NewError("User is unspecified")
+		}
+		if req.User.Hash == "" {
+			req.User.Hash = common.SHA224String(req.User.Password)
+		}
+		switch req.Operation {
+		case SetUsersRequest_Add:
+			err = s.auth.AddUser(req.User.Hash)
+			if req.SpeedLimit != nil {
+				valid, user := s.auth.AuthUser(req.User.Hash)
+				if !valid {
+					return common.NewError("Failed to add new user")
+				}
+				user.SetSpeedLimit(int(req.SpeedLimit.DownloadSpeed), int(req.SpeedLimit.UploadSpeed))
+			}
+		case SetUsersRequest_Delete:
+			err = s.auth.DelUser(req.User.Hash)
+		case SetUsersRequest_Modify:
+			valid, user := s.auth.AuthUser(req.User.Hash)
+			if !valid {
+				err = common.NewError("Invalid user " + req.User.Hash)
+			} else {
+				if req.SpeedLimit.DownloadSpeed > 0 || req.SpeedLimit.UploadSpeed > 0 {
+					user.SetSpeedLimit(int(req.SpeedLimit.DownloadSpeed), int(req.SpeedLimit.UploadSpeed))
+				}
+				if req.IpLimit > 0 {
+					user.SetIPLimit(int(req.IpLimit))
+				}
+			}
+		}
+		if err != nil {
+			stream.Send(&SetUsersResponse{
+				Success: false,
+				Info:    err.Error(),
+			})
+			continue
+		}
+		stream.Send(&SetUsersResponse{
+			Success: true,
+		})
+	}
+}
+
+func (s *ServerAPI) ListUsers(req *ListUsersRequest, stream TrojanServerService_ListUsersServer) error {
+	log.Debug("API: ListUsers")
+	users := s.auth.ListUsers()
+	for _, user := range users {
+		downloadTraffic, uploadTraffic := user.GetTraffic()
+		downloadSpeed, uploadSpeed := user.GetSpeed()
+		downloadSpeedLimit, uploadSpeedLimit := user.GetSpeedLimit()
+		ipLimit := user.GetIPLimit()
+		ipCurrent := user.GetIP()
+		err := stream.Send(&ListUsersResponse{
+			User: &User{
+				Hash: user.Hash(),
+			},
+			Status: &UserStatus{
+				TrafficTotal: &Traffic{
+					DownloadTraffic: downloadTraffic,
+					UploadTraffic:   uploadTraffic,
+				},
+				SpeedCurrent: &Speed{
+					DownloadSpeed: downloadSpeed,
+					UploadSpeed:   uploadSpeed,
+				},
+				SpeedLimit: &Speed{
+					DownloadSpeed: uint64(downloadSpeedLimit),
+					UploadSpeed:   uint64(uploadSpeedLimit),
+				},
+				IpLimit:   int32(ipLimit),
+				IpCurrent: int32(ipCurrent),
+			},
+		})
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func RunServerAPI(ctx context.Context, auth statistic.Authenticator) error {
+	server := grpc.NewServer()
+	service := &ServerAPI{
+		auth: auth,
+	}
+	RegisterTrojanServerServiceServer(server, service)
+	cfg := config.FromContext(ctx, Name).(*Config)
+	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", cfg.APIHost, cfg.APIPort))
+	if err != nil {
+		return err
+	}
+	log.Info("server-side api service is listening on", listener.Addr().String())
+	errChan := make(chan error, 1)
+	go func() {
+		errChan <- server.Serve(listener)
+	}()
+	select {
+	case err := <-errChan:
+		return err
+	case <-ctx.Done():
+		server.Stop()
+		return nil
+	}
+}
+
+func init() {
+	api.RegisterHandler(trojan.Name+"_SERVER", RunServerAPI)
+}

api/service/server_test.go

@@ -0,0 +1,132 @@
+package service
+
+import (
+	"context"
+	"fmt"
+	"github.com/p4gefau1t/trojan-go/common"
+	"github.com/p4gefau1t/trojan-go/config"
+	"github.com/p4gefau1t/trojan-go/statistic/memory"
+	"google.golang.org/grpc"
+	"testing"
+	"time"
+)
+
+func TestServerAPI(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	ctx = config.WithConfig(ctx, memory.Name,
+		&memory.Config{
+			Passwords: []string{},
+		})
+	port := common.PickPort("tcp", "127.0.0.1")
+	ctx = config.WithConfig(ctx, Name, &Config{
+		APIConfig{
+			Enabled: true,
+			APIHost: "127.0.0.1",
+			APIPort: port,
+		},
+	})
+	auth, err := memory.NewAuthenticator(ctx)
+	common.Must(err)
+	go RunServerAPI(ctx, auth)
+	common.Must(auth.AddUser("hash1234"))
+	_, user := auth.AuthUser("hash1234")
+	conn, err := grpc.Dial(fmt.Sprintf("127.0.0.1:%d", port), grpc.WithInsecure())
+	common.Must(err)
+	server := NewTrojanServerServiceClient(conn)
+	stream1, err := server.ListUsers(ctx, &ListUsersRequest{})
+	common.Must(err)
+	for {
+		resp, err := stream1.Recv()
+		if err != nil {
+			break
+		}
+		fmt.Println(resp.User.Hash)
+		if resp.User.Hash != "hash1234" {
+			t.Fail()
+		}
+		fmt.Println(resp.Status.SpeedCurrent)
+		fmt.Println(resp.Status.SpeedLimit)
+	}
+	stream1.CloseSend()
+	user.AddTraffic(1234, 5678)
+	time.Sleep(time.Millisecond * 1000)
+	stream2, err := server.GetUsers(ctx)
+	common.Must(err)
+	stream2.Send(&GetUsersRequest{
+		User: &User{
+			Hash: "hash1234",
+		},
+	})
+	resp2, err := stream2.Recv()
+	common.Must(err)
+	if resp2.Status.TrafficTotal.DownloadTraffic != 1234 || resp2.Status.TrafficTotal.UploadTraffic != 5678 {
+		t.Fail()
+	}
+	if resp2.Status.SpeedCurrent.DownloadSpeed != 1234 || resp2.Status.TrafficTotal.UploadTraffic != 5678 {
+		t.Fail()
+	}
+
+	stream3, err := server.SetUsers(ctx)
+	stream3.Send(&SetUsersRequest{
+		User: &User{
+			Hash: "hash1234",
+		},
+		Operation: SetUsersRequest_Delete,
+	})
+	resp3, err := stream3.Recv()
+	if err != nil || !resp3.Success {
+		t.Fail()
+	}
+	valid, _ := auth.AuthUser("hash1234")
+	if valid {
+		t.Fail()
+	}
+	stream3.Send(&SetUsersRequest{
+		User: &User{
+			Hash: "newhash",
+		},
+		Operation: SetUsersRequest_Add,
+	})
+	resp3, err = stream3.Recv()
+	if err != nil || !resp3.Success {
+		t.Fail()
+	}
+	valid, user = auth.AuthUser("newhash")
+	if !valid {
+		t.Fail()
+	}
+	stream3.Send(&SetUsersRequest{
+		User: &User{
+			Hash: "newhash",
+		},
+		Operation: SetUsersRequest_Modify,
+		SpeedLimit: &Speed{
+			DownloadSpeed: 5000,
+			UploadSpeed:   3000,
+		},
+	})
+	go func() {
+		for {
+			user.AddTraffic(200, 0)
+		}
+	}()
+	go func() {
+		for {
+			user.AddTraffic(0, 300)
+		}
+	}()
+	time.Sleep(time.Second * 3)
+	for i := 0; i < 3; i++ {
+		stream2.Send(&GetUsersRequest{
+			User: &User{
+				Hash: "newhash",
+			},
+		})
+		resp2, err = stream2.Recv()
+		fmt.Println(resp2.Status.SpeedCurrent)
+		fmt.Println(resp2.Status.SpeedLimit)
+		time.Sleep(time.Second)
+	}
+	stream2.CloseSend()
+	cancel()
+}