From 24f688cc2018c92b55dfb2318972c8e5323f7140 Mon Sep 17 00:00:00 2001
From: maskedeken <52683904+maskedeken@users.noreply.github.com>
Date: Thu, 22 Oct 2020 12:01:26 +0800
Subject: [PATCH] Trojan Protocol Handler implements UserManager (#344)

* Trojan Protocol Handler implements UserManager

* Update validator.go

Co-authored-by: RPRX <63339210+rprx@users.noreply.github.com>
---
 infra/conf/trojan.go      |  5 -----
 proxy/trojan/server.go    | 10 ++++++++++
 proxy/trojan/validator.go | 27 +++++++++++++++++++++++++--
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/infra/conf/trojan.go b/infra/conf/trojan.go
index 546e9faa7..0929bb649 100644
--- a/infra/conf/trojan.go
+++ b/infra/conf/trojan.go
@@ -99,11 +99,6 @@ type TrojanServerConfig struct {
 // Build implements Buildable
 func (c *TrojanServerConfig) Build() (proto.Message, error) {
 	config := new(trojan.ServerConfig)
-
-	if len(c.Clients) == 0 {
-		return nil, newError("No trojan user settings.")
-	}
-
 	config.Users = make([]*protocol.User, len(c.Clients))
 	for idx, rawUser := range c.Clients {
 		user := new(protocol.User)
diff --git a/proxy/trojan/server.go b/proxy/trojan/server.go
index d2f9d3db5..c271e0321 100644
--- a/proxy/trojan/server.go
+++ b/proxy/trojan/server.go
@@ -85,6 +85,16 @@ func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) {
 	return server, nil
 }
 
+// AddUser implements proxy.UserManager.AddUser().
+func (s *Server) AddUser(ctx context.Context, u *protocol.MemoryUser) error {
+	return s.validator.Add(u)
+}
+
+// RemoveUser implements proxy.UserManager.RemoveUser().
+func (s *Server) RemoveUser(ctx context.Context, e string) error {
+	return s.validator.Del(e)
+}
+
 // Network implements proxy.Inbound.Network().
 func (s *Server) Network() []net.Network {
 	return []net.Network{net.Network_TCP}
diff --git a/proxy/trojan/validator.go b/proxy/trojan/validator.go
index 1c7926e3a..038069e13 100644
--- a/proxy/trojan/validator.go
+++ b/proxy/trojan/validator.go
@@ -1,6 +1,7 @@
 package trojan
 
 import (
+	"strings"
 	"sync"
 
 	"v2ray.com/core/common/protocol"
@@ -8,13 +9,35 @@ import (
 
 // Validator stores valid trojan users
 type Validator struct {
+	// Considering email's usage here, map + sync.Mutex/RWMutex may have better performance.
+	email sync.Map
 	users sync.Map
 }
 
 // Add a trojan user
 func (v *Validator) Add(u *protocol.MemoryUser) error {
-	user := u.Account.(*MemoryAccount)
-	v.users.Store(hexString(user.Key), u)
+	if u.Email != "" {
+		_, loaded := v.email.LoadOrStore(strings.ToLower(u.Email), u)
+		if loaded {
+			return newError("User ", u.Email, " already exists.")
+		}
+	}
+	v.users.Store(hexString(u.Account.(*MemoryAccount).Key), u)
+	return nil
+}
+
+// Del a trojan user
+func (v *Validator) Del(e string) error {
+	if e == "" {
+		return newError("Email must not be empty.")
+	}
+	le := strings.ToLower(e)
+	u, _ := v.email.Load(le)
+	if u == nil {
+		return newError("User ", e, " not found.")
+	}
+	v.email.Delete(le)
+	v.users.Delete(hexString(u.(*protocol.MemoryUser).Account.(*MemoryAccount).Key))
 	return nil
 }