Refactor: new Shadowsocks validator (#629)

* Refactor: new Shadowsocks validator

* Fix NoneCliper cannot work

* Feat: refine the size of drain

* fix: fix validator after merge 'main'

* fix: UDP user logic

* style: refine code style

proxy/shadowsocks/config.go

@@ -7,8 +7,6 @@ import (
 	"crypto/md5"
 	"crypto/sha1"
 	"io"
-	"reflect"
-	"strconv"
 
 	"golang.org/x/crypto/chacha20poly1305"
 	"golang.org/x/crypto/hkdf"
@@ -28,6 +26,10 @@ type MemoryAccount struct {
 	replayFilter antireplay.GeneralizedReplayFilter
 }
 
+var (
+	ErrIVNotUnique = newError("IV is not unique")
+)
+
 // Equals implements protocol.Account.Equals().
 func (a *MemoryAccount) Equals(another protocol.Account) bool {
 	if account, ok := another.(*MemoryAccount); ok {
@@ -43,24 +45,7 @@ func (a *MemoryAccount) CheckIV(iv []byte) error {
 	if a.replayFilter.Check(iv) {
 		return nil
 	}
-	return newError("IV is not unique")
-}
-
-func (a *MemoryAccount) GetCipherName() string {
-	switch a.Cipher.(type) {
-	case *AEADCipher:
-		switch reflect.ValueOf(a.Cipher.(*AEADCipher).AEADAuthCreator).Pointer() {
-		case reflect.ValueOf(createAesGcm).Pointer():
-			keyBytes := a.Cipher.(*AEADCipher).KeyBytes
-			return "AES_" + strconv.FormatInt(int64(keyBytes*8), 10) + "_GCM"
-		case reflect.ValueOf(createChaCha20Poly1305).Pointer():
-			return "CHACHA20_POLY1305"
-		}
-	case *NoneCipher:
-		return "NONE"
-	}
-
-	return ""
+	return ErrIVNotUnique
 }
 
 func createAesGcm(key []byte) cipher.AEAD {

proxy/shadowsocks/protocol.go

@@ -1,11 +1,7 @@
 package shadowsocks
 
 import (
-	"crypto/cipher"
-	"crypto/hmac"
 	"crypto/rand"
-	"crypto/sha256"
-	"hash/crc32"
 	"io"
 
 	"github.com/xtls/xray-core/common"
@@ -54,10 +50,7 @@ func (r *FullReader) Read(p []byte) (n int, err error) {
 
 // ReadTCPSession reads a Shadowsocks TCP session from the given reader, returns its header and remaining parts.
 func ReadTCPSession(validator *Validator, reader io.Reader) (*protocol.RequestHeader, buf.Reader, error) {
-	hashkdf := hmac.New(sha256.New, []byte("SSBSKDF"))
-
-	behaviorSeed := crc32.ChecksumIEEE(hashkdf.Sum(nil))
-
+	behaviorSeed := validator.GetBehaviorSeed()
 	behaviorRand := dice.NewDeterministicDice(int64(behaviorSeed))
 	BaseDrainSize := behaviorRand.Roll(3266)
 	RandDrainMax := behaviorRand.Roll(64) + 1
@@ -65,23 +58,10 @@ func ReadTCPSession(validator *Validator, reader io.Reader) (*protocol.RequestHe
 	DrainSize := BaseDrainSize + 16 + 38 + RandDrainRolled
 	readSizeRemain := DrainSize
 
-	var r2 buf.Reader
+	var r buf.Reader
 	buffer := buf.New()
 	defer buffer.Release()
 
-	var user *protocol.MemoryUser
-	var ivLen int32
-	var iv []byte
-	var err error
-
-	count := validator.Count()
-	if count == 0 {
-		readSizeRemain -= int(buffer.Len())
-		DrainConnN(reader, readSizeRemain)
-		return nil, nil, newError("invalid user")
-	} else if count > 1 {
-		var aead cipher.AEAD
-
 	if _, err := buffer.ReadFullFrom(reader, 50); err != nil {
 		readSizeRemain -= int(buffer.Len())
 		DrainConnN(reader, readSizeRemain)
@@ -89,48 +69,41 @@ func ReadTCPSession(validator *Validator, reader io.Reader) (*protocol.RequestHe
 	}
 
 	bs := buffer.Bytes()
-		user, aead, _, ivLen, err = validator.Get(bs, protocol.RequestCommandTCP)
+	user, aead, _, ivLen, err := validator.Get(bs, protocol.RequestCommandTCP)
 
-		if user != nil {
-			if ivLen > 0 {
-				iv = append([]byte(nil), bs[:ivLen]...)
-			}
+	switch err {
+	case ErrNotFound:
+		readSizeRemain -= int(buffer.Len())
+		DrainConnN(reader, readSizeRemain)
+		return nil, nil, newError("failed to match an user").Base(err)
+	case ErrIVNotUnique:
+		readSizeRemain -= int(buffer.Len())
+		DrainConnN(reader, readSizeRemain)
+		return nil, nil, newError("failed iv check").Base(err)
+	default:
 		reader = &FullReader{reader, bs[ivLen:]}
+		readSizeRemain -= int(ivLen)
+
+		if aead != nil {
 			auth := &crypto.AEADAuthenticator{
 				AEAD:           aead,
 				NonceGenerator: crypto.GenerateInitialAEADNonce(),
 			}
-			r2 = crypto.NewAuthenticationReader(auth, &crypto.AEADChunkSizeParser{
+			r = crypto.NewAuthenticationReader(auth, &crypto.AEADChunkSizeParser{
 				Auth: auth,
 			}, reader, protocol.TransferTypeStream, nil)
 		} else {
-			readSizeRemain -= int(buffer.Len())
-			DrainConnN(reader, readSizeRemain)
-			return nil, nil, newError("failed to match an user").Base(err)
-		}
-	} else {
-		user, ivLen = validator.GetOnlyUser()
 			account := user.Account.(*MemoryAccount)
-		hashkdf.Write(account.Key)
-		if ivLen > 0 {
-			if _, err := buffer.ReadFullFrom(reader, ivLen); err != nil {
-				readSizeRemain -= int(buffer.Len())
-				DrainConnN(reader, readSizeRemain)
-				return nil, nil, newError("failed to read IV").Base(err)
-			}
-			iv = append([]byte(nil), buffer.BytesTo(ivLen)...)
-		}
-
-		r, err := account.Cipher.NewDecryptionReader(account.Key, iv, reader)
+			iv := append([]byte(nil), buffer.BytesTo(ivLen)...)
+			r, err = account.Cipher.NewDecryptionReader(account.Key, iv, reader)
 			if err != nil {
-			readSizeRemain -= int(buffer.Len())
 				DrainConnN(reader, readSizeRemain)
 				return nil, nil, newError("failed to initialize decoding stream").Base(err).AtError()
 			}
-		r2 = r
+		}
 	}
 
-	br := &buf.BufferedReader{Reader: r2}
+	br := &buf.BufferedReader{Reader: r}
 
 	request := &protocol.RequestHeader{
 		Version: Version,
@@ -138,7 +111,6 @@ func ReadTCPSession(validator *Validator, reader io.Reader) (*protocol.RequestHe
 		Command: protocol.RequestCommandTCP,
 	}
 
-	readSizeRemain -= int(buffer.Len())
 	buffer.Clear()
 
 	addr, port, err := addrParser.ReadAddressPort(buffer, br)
@@ -157,13 +129,6 @@ func ReadTCPSession(validator *Validator, reader io.Reader) (*protocol.RequestHe
 		return nil, nil, newError("invalid remote address.")
 	}
 
-	account := user.Account.(*MemoryAccount)
-	if ivError := account.CheckIV(iv); ivError != nil {
-		readSizeRemain -= int(buffer.Len())
-		DrainConnN(reader, readSizeRemain)
-		return nil, nil, newError("failed iv check").Base(ivError)
-	}
-
 	return request, br, nil
 }
 
@@ -273,36 +238,27 @@ func DecodeUDPPacket(validator *Validator, payload *buf.Buffer) (*protocol.Reque
 		return nil, nil, newError("len(bs) <= 32")
 	}
 
-	var user *protocol.MemoryUser
-	var err error
-
-	count := validator.Count()
-	if count == 0 {
-		return nil, nil, newError("invalid user")
-	} else if count > 1 {
-		var d []byte
-		user, _, d, _, err = validator.Get(bs, protocol.RequestCommandUDP)
-
-		if user != nil {
+	user, _, d, _, err := validator.Get(bs, protocol.RequestCommandUDP)
+	switch err {
+	case ErrIVNotUnique:
+		return nil, nil, newError("failed iv check").Base(err)
+	case ErrNotFound:
+		return nil, nil, newError("failed to match an user").Base(err)
+	default:
+		account := user.Account.(*MemoryAccount)
+		if account.Cipher.IsAEAD() {
 			payload.Clear()
 			payload.Write(d)
 		} else {
-			return nil, nil, newError("failed to decrypt UDP payload").Base(err)
-		}
-	} else {
-		user, _ = validator.GetOnlyUser()
-		account := user.Account.(*MemoryAccount)
-
-		var iv []byte
-		if !account.Cipher.IsAEAD() && account.Cipher.IVSize() > 0 {
-			// Keep track of IV as it gets removed from payload in DecodePacket.
-			iv = make([]byte, account.Cipher.IVSize())
+			if account.Cipher.IVSize() > 0 {
+				iv := make([]byte, account.Cipher.IVSize())
 				copy(iv, payload.BytesTo(account.Cipher.IVSize()))
 			}
 			if err = account.Cipher.DecodePacket(account.Key, payload); err != nil {
 				return nil, nil, newError("failed to decrypt UDP payload").Base(err)
 			}
 		}
+	}
 
 	request := &protocol.RequestHeader{
 		Version: Version,

proxy/shadowsocks/server.go

@@ -115,10 +115,6 @@ func (s *Server) handleUDPPayload(ctx context.Context, conn stat.Connection, dis
 		panic("no inbound metadata")
 	}
 
-	if s.validator.Count() == 1 {
-		inbound.User, _ = s.validator.GetOnlyUser()
-	}
-
 	var dest *net.Destination
 
 	reader := buf.NewPacketReader(conn)

proxy/shadowsocks/validator.go

@@ -2,112 +2,127 @@ package shadowsocks
 
 import (
 	"crypto/cipher"
+	"crypto/hmac"
+	"crypto/sha256"
+	"hash/crc64"
 	"strings"
 	"sync"
 
+	"github.com/xtls/xray-core/common/dice"
 	"github.com/xtls/xray-core/common/protocol"
 )
 
 // Validator stores valid Shadowsocks users.
 type Validator struct {
-	// Considering email's usage here, map + sync.Mutex/RWMutex may have better performance.
-	email sync.Map
-	users sync.Map
+	sync.RWMutex
+	users []*protocol.MemoryUser
+
+	behaviorSeed  uint64
+	behaviorFused bool
 }
 
-// Add a Shadowsocks user, Email must be empty or unique.
+var (
+	ErrNotFound = newError("Not Found")
+)
+
+// Add a Shadowsocks user.
 func (v *Validator) Add(u *protocol.MemoryUser) error {
+	v.Lock()
+	defer v.Unlock()
+
 	account := u.Account.(*MemoryAccount)
+	if !account.Cipher.IsAEAD() && len(v.users) > 0 {
+		return newError("The cipher is not support Single-port Multi-user")
+	}
+	v.users = append(v.users, u)
 
-	if !account.Cipher.IsAEAD() && v.Count() > 0 {
-		return newError("The cipher do not support Single-port Multi-user")
+	if !v.behaviorFused {
+		hashkdf := hmac.New(sha256.New, []byte("SSBSKDF"))
+		hashkdf.Write(account.Key)
+		v.behaviorSeed = crc64.Update(v.behaviorSeed, crc64.MakeTable(crc64.ECMA), hashkdf.Sum(nil))
 	}
 
-	if u.Email != "" {
-		_, loaded := v.email.LoadOrStore(strings.ToLower(u.Email), u)
-		if loaded {
-			return newError("User ", u.Email, " already exists.")
-		}
-	}
-
-	v.users.Store(string(account.Key)+"&"+account.GetCipherName(), u)
 	return nil
 }
 
 // Del a Shadowsocks user with a non-empty Email.
-func (v *Validator) Del(e string) error {
-	if e == "" {
+func (v *Validator) Del(email string) error {
+	if email == "" {
 		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.Lock()
+	defer v.Unlock()
+
+	email = strings.ToLower(email)
+	idx := -1
+	for i, u := range v.users {
+		if strings.EqualFold(u.Email, email) {
+			idx = i
+			break
 		}
-	account := u.(*protocol.MemoryUser).Account.(*MemoryAccount)
-	v.email.Delete(le)
-	v.users.Delete(string(account.Key) + "&" + account.GetCipherName())
+	}
+
+	if idx == -1 {
+		return newError("User ", email, " not found.")
+	}
+	ulen := len(v.users)
+
+	v.users[idx] = v.users[ulen-1]
+	v.users[ulen-1] = nil
+	v.users = v.users[:ulen-1]
+
 	return nil
 }
 
-// Count the number of Shadowsocks users
-func (v *Validator) Count() int {
-	length := 0
-	v.users.Range(func(_, _ interface{}) bool {
-		length++
-
-		return true
-	})
-	return length
-}
-
-// Get a Shadowsocks user and the user's cipher.
+// Get a Shadowsocks user.
 func (v *Validator) Get(bs []byte, command protocol.RequestCommand) (u *protocol.MemoryUser, aead cipher.AEAD, ret []byte, ivLen int32, err error) {
-	var dataSize int
+	v.RLock()
+	defer v.RUnlock()
 
-	switch command {
-	case protocol.RequestCommandTCP:
-		dataSize = 16
-	case protocol.RequestCommandUDP:
-		dataSize = 8192
-	}
-
-	var aeadCipher *AEADCipher
-	subkey := make([]byte, 32)
-	data := make([]byte, dataSize)
-
-	v.users.Range(func(key, user interface{}) bool {
-		account := user.(*protocol.MemoryUser).Account.(*MemoryAccount)
-		aeadCipher = account.Cipher.(*AEADCipher)
+	for _, user := range v.users {
+		if account := user.Account.(*MemoryAccount); account.Cipher.IsAEAD() {
+			aeadCipher := account.Cipher.(*AEADCipher)
 			ivLen = aeadCipher.IVSize()
+			iv := bs[:ivLen]
+			subkey := make([]byte, 32)
 			subkey = subkey[:aeadCipher.KeyBytes]
-		hkdfSHA1(account.Key, bs[:ivLen], subkey)
+			hkdfSHA1(account.Key, iv, subkey)
 			aead = aeadCipher.AEADAuthCreator(subkey)
 
+			var matchErr error
 			switch command {
 			case protocol.RequestCommandTCP:
-			ret, err = aead.Open(data[:0], data[4:16], bs[ivLen:ivLen+18], nil)
+				data := make([]byte, 16)
+				ret, matchErr = aead.Open(data[:0], data[4:16], bs[ivLen:ivLen+18], nil)
 			case protocol.RequestCommandUDP:
-			ret, err = aead.Open(data[:0], data[8180:8192], bs[ivLen:], nil)
+				data := make([]byte, 8192)
+				ret, matchErr = aead.Open(data[:0], data[8180:8192], bs[ivLen:], nil)
 			}
 
-		if err == nil {
-			u = user.(*protocol.MemoryUser)
-			return false
-		}
-		return true
-	})
-
+			if matchErr == nil {
+				u = user
+				err = account.CheckIV(iv)
 				return
 			}
-
-// Get the only user without authentication
-func (v *Validator) GetOnlyUser() (u *protocol.MemoryUser, ivLen int32) {
-	v.users.Range(func(_, user interface{}) bool {
-		u = user.(*protocol.MemoryUser)
-		return false
-	})
-	ivLen = u.Account.(*MemoryAccount).Cipher.IVSize()
-
+		} else {
+			u = user
+			ivLen = user.Account.(*MemoryAccount).Cipher.IVSize()
+			// err = user.Account.(*MemoryAccount).CheckIV(bs[:ivLen]) // The IV size of None Cipher is 0.
 			return
 		}
+	}
+
+	return nil, nil, nil, 0, ErrNotFound
+}
+
+func (v *Validator) GetBehaviorSeed() uint64 {
+	v.Lock()
+	defer v.Unlock()
+
+	v.behaviorFused = true
+	if v.behaviorSeed == 0 {
+		v.behaviorSeed = dice.RollUint64()
+	}
+	return v.behaviorSeed
+}