reduce memory usage of vmess internal objects

common/protocol/user_validator.go

@@ -1,7 +0,0 @@
-package protocol
-
-type UserValidator interface {
-	Add(user *User) error
-	Get(timeHash []byte) (*User, Timestamp, bool)
-	Remove(email string) bool
-}

proxy/vmess/encoding/client.go

@@ -31,13 +31,13 @@ func hashTimestamp(t protocol.Timestamp) []byte {
 
 // ClientSession stores connection session info for VMess client.
 type ClientSession struct {
-	requestBodyKey  []byte
-	requestBodyIV   []byte
-	responseHeader  byte
-	responseBodyKey []byte
-	responseBodyIV  []byte
-	responseReader  io.Reader
 	idHash          protocol.IDHash
+	requestBodyKey  [16]byte
+	requestBodyIV   [16]byte
+	responseBodyKey [16]byte
+	responseBodyIV  [16]byte
+	responseReader  io.Reader
+	responseHeader  byte
 }
 
 // NewClientSession creates a new ClientSession.
@@ -46,13 +46,11 @@ func NewClientSession(idHash protocol.IDHash) *ClientSession {
 	common.Must2(rand.Read(randomBytes))
 
 	session := &ClientSession{}
-	session.requestBodyKey = randomBytes[:16]
-	session.requestBodyIV = randomBytes[16:32]
+	copy(session.requestBodyKey[:], randomBytes[:16])
+	copy(session.requestBodyIV[:], randomBytes[16:32])
 	session.responseHeader = randomBytes[32]
-	responseBodyKey := md5.Sum(session.requestBodyKey)
-	responseBodyIV := md5.Sum(session.requestBodyIV)
-	session.responseBodyKey = responseBodyKey[:]
-	session.responseBodyIV = responseBodyIV[:]
+	session.responseBodyKey = md5.Sum(session.requestBodyKey[:])
+	session.responseBodyIV = md5.Sum(session.requestBodyIV[:])
 	session.idHash = idHash
 
 	return session
@@ -72,8 +70,8 @@ func (c *ClientSession) EncodeRequestHeader(header *protocol.RequestHeader, writ
 	defer buffer.Release()
 
 	buffer.AppendBytes(Version)
-	buffer.Append(c.requestBodyIV)
-	buffer.Append(c.requestBodyKey)
+	buffer.Append(c.requestBodyIV[:])
+	buffer.Append(c.requestBodyKey[:])
 	buffer.AppendBytes(c.responseHeader, byte(header.Option))
 
 	padingLen := dice.Roll(16)
@@ -110,7 +108,7 @@ func (c *ClientSession) EncodeRequestHeader(header *protocol.RequestHeader, writ
 func (c *ClientSession) EncodeRequestBody(request *protocol.RequestHeader, writer io.Writer) buf.Writer {
 	var sizeParser crypto.ChunkSizeEncoder = crypto.PlainChunkSizeParser{}
 	if request.Option.Has(protocol.RequestOptionChunkMasking) {
-		sizeParser = NewShakeSizeParser(c.requestBodyIV)
+		sizeParser = NewShakeSizeParser(c.requestBodyIV[:])
 	}
 	switch request.Security {
 	case protocol.SecurityType_NONE:
@@ -128,7 +126,7 @@ func (c *ClientSession) EncodeRequestBody(request *protocol.RequestHeader, write
 
 		return buf.NewWriter(writer)
 	case protocol.SecurityType_LEGACY:
-		aesStream := crypto.NewAesEncryptionStream(c.requestBodyKey, c.requestBodyIV)
+		aesStream := crypto.NewAesEncryptionStream(c.requestBodyKey[:], c.requestBodyIV[:])
 		cryptionWriter := crypto.NewCryptionWriter(aesStream, writer)
 		if request.Option.Has(protocol.RequestOptionChunkStream) {
 			auth := &crypto.AEADAuthenticator{
@@ -141,21 +139,21 @@ func (c *ClientSession) EncodeRequestBody(request *protocol.RequestHeader, write
 
 		return buf.NewWriter(cryptionWriter)
 	case protocol.SecurityType_AES128_GCM:
-		block, _ := aes.NewCipher(c.requestBodyKey)
+		block, _ := aes.NewCipher(c.requestBodyKey[:])
 		aead, _ := cipher.NewGCM(block)
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
-			NonceGenerator:          GenerateChunkNonce(c.requestBodyIV, uint32(aead.NonceSize())),
+			NonceGenerator:          GenerateChunkNonce(c.requestBodyIV[:], uint32(aead.NonceSize())),
 			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType())
 	case protocol.SecurityType_CHACHA20_POLY1305:
-		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(c.requestBodyKey))
+		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(c.requestBodyKey[:]))
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
-			NonceGenerator:          GenerateChunkNonce(c.requestBodyIV, uint32(aead.NonceSize())),
+			NonceGenerator:          GenerateChunkNonce(c.requestBodyIV[:], uint32(aead.NonceSize())),
 			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType())
@@ -165,7 +163,7 @@ func (c *ClientSession) EncodeRequestBody(request *protocol.RequestHeader, write
 }
 
 func (c *ClientSession) DecodeResponseHeader(reader io.Reader) (*protocol.ResponseHeader, error) {
-	aesStream := crypto.NewAesDecryptionStream(c.responseBodyKey, c.responseBodyIV)
+	aesStream := crypto.NewAesDecryptionStream(c.responseBodyKey[:], c.responseBodyIV[:])
 	c.responseReader = crypto.NewCryptionReader(aesStream, reader)
 
 	buffer := buf.New()
@@ -202,7 +200,7 @@ func (c *ClientSession) DecodeResponseHeader(reader io.Reader) (*protocol.Respon
 func (c *ClientSession) DecodeResponseBody(request *protocol.RequestHeader, reader io.Reader) buf.Reader {
 	var sizeParser crypto.ChunkSizeDecoder = crypto.PlainChunkSizeParser{}
 	if request.Option.Has(protocol.RequestOptionChunkMasking) {
-		sizeParser = NewShakeSizeParser(c.responseBodyIV)
+		sizeParser = NewShakeSizeParser(c.responseBodyIV[:])
 	}
 	switch request.Security {
 	case protocol.SecurityType_NONE:
@@ -233,21 +231,21 @@ func (c *ClientSession) DecodeResponseBody(request *protocol.RequestHeader, read
 
 		return buf.NewReader(c.responseReader)
 	case protocol.SecurityType_AES128_GCM:
-		block, _ := aes.NewCipher(c.responseBodyKey)
+		block, _ := aes.NewCipher(c.responseBodyKey[:])
 		aead, _ := cipher.NewGCM(block)
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
-			NonceGenerator:          GenerateChunkNonce(c.responseBodyIV, uint32(aead.NonceSize())),
+			NonceGenerator:          GenerateChunkNonce(c.responseBodyIV[:], uint32(aead.NonceSize())),
 			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType())
 	case protocol.SecurityType_CHACHA20_POLY1305:
-		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(c.responseBodyKey))
+		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(c.responseBodyKey[:]))
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
-			NonceGenerator:          GenerateChunkNonce(c.responseBodyIV, uint32(aead.NonceSize())),
+			NonceGenerator:          GenerateChunkNonce(c.responseBodyIV[:], uint32(aead.NonceSize())),
 			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType())

proxy/vmess/encoding/server.go

@@ -85,19 +85,19 @@ func (h *SessionHistory) removeExpiredEntries() {
 }
 
 type ServerSession struct {
-	userValidator   protocol.UserValidator
+	userValidator   *vmess.TimedUserValidator
 	sessionHistory  *SessionHistory
-	requestBodyKey  []byte
-	requestBodyIV   []byte
-	responseBodyKey []byte
-	responseBodyIV  []byte
-	responseHeader  byte
+	requestBodyKey  [16]byte
+	requestBodyIV   [16]byte
+	responseBodyKey [16]byte
+	responseBodyIV  [16]byte
 	responseWriter  io.Writer
+	responseHeader  byte
 }
 
 // NewServerSession creates a new ServerSession, using the given UserValidator.
 // The ServerSession instance doesn't take ownership of the validator.
-func NewServerSession(validator protocol.UserValidator, sessionHistory *SessionHistory) *ServerSession {
+func NewServerSession(validator *vmess.TimedUserValidator, sessionHistory *SessionHistory) *ServerSession {
 	return &ServerSession{
 		userValidator:  validator,
 		sessionHistory: sessionHistory,
@@ -150,12 +150,12 @@ func (s *ServerSession) DecodeRequestHeader(reader io.Reader) (*protocol.Request
 		Version: buffer.Byte(0),
 	}
 
-	s.requestBodyIV = append([]byte(nil), buffer.BytesRange(1, 17)...)   // 16 bytes
-	s.requestBodyKey = append([]byte(nil), buffer.BytesRange(17, 33)...) // 16 bytes
+	copy(s.requestBodyIV[:], buffer.BytesRange(1, 17))   // 16 bytes
+	copy(s.requestBodyKey[:], buffer.BytesRange(17, 33)) // 16 bytes
 	var sid sessionId
 	copy(sid.user[:], vmessAccount.ID.Bytes())
-	copy(sid.key[:], s.requestBodyKey)
-	copy(sid.nonce[:], s.requestBodyIV)
+	sid.key = s.requestBodyKey
+	sid.nonce = s.requestBodyIV
 	if !s.sessionHistory.addIfNotExits(sid) {
 		return nil, newError("duplicated session id, possibly under replay attack")
 	}
@@ -227,7 +227,7 @@ func (s *ServerSession) DecodeRequestHeader(reader io.Reader) (*protocol.Request
 func (s *ServerSession) DecodeRequestBody(request *protocol.RequestHeader, reader io.Reader) buf.Reader {
 	var sizeParser crypto.ChunkSizeDecoder = crypto.PlainChunkSizeParser{}
 	if request.Option.Has(protocol.RequestOptionChunkMasking) {
-		sizeParser = NewShakeSizeParser(s.requestBodyIV)
+		sizeParser = NewShakeSizeParser(s.requestBodyIV[:])
 	}
 	switch request.Security {
 	case protocol.SecurityType_NONE:
@@ -246,7 +246,7 @@ func (s *ServerSession) DecodeRequestBody(request *protocol.RequestHeader, reade
 
 		return buf.NewReader(reader)
 	case protocol.SecurityType_LEGACY:
-		aesStream := crypto.NewAesDecryptionStream(s.requestBodyKey, s.requestBodyIV)
+		aesStream := crypto.NewAesDecryptionStream(s.requestBodyKey[:], s.requestBodyIV[:])
 		cryptionReader := crypto.NewCryptionReader(aesStream, reader)
 		if request.Option.Has(protocol.RequestOptionChunkStream) {
 			auth := &crypto.AEADAuthenticator{
@@ -259,21 +259,21 @@ func (s *ServerSession) DecodeRequestBody(request *protocol.RequestHeader, reade
 
 		return buf.NewReader(cryptionReader)
 	case protocol.SecurityType_AES128_GCM:
-		block, _ := aes.NewCipher(s.requestBodyKey)
+		block, _ := aes.NewCipher(s.requestBodyKey[:])
 		aead, _ := cipher.NewGCM(block)
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
-			NonceGenerator:          GenerateChunkNonce(s.requestBodyIV, uint32(aead.NonceSize())),
+			NonceGenerator:          GenerateChunkNonce(s.requestBodyIV[:], uint32(aead.NonceSize())),
 			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType())
 	case protocol.SecurityType_CHACHA20_POLY1305:
-		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(s.requestBodyKey))
+		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(s.requestBodyKey[:]))
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
-			NonceGenerator:          GenerateChunkNonce(s.requestBodyIV, uint32(aead.NonceSize())),
+			NonceGenerator:          GenerateChunkNonce(s.requestBodyIV[:], uint32(aead.NonceSize())),
 			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType())
@@ -283,12 +283,12 @@ func (s *ServerSession) DecodeRequestBody(request *protocol.RequestHeader, reade
 }
 
 func (s *ServerSession) EncodeResponseHeader(header *protocol.ResponseHeader, writer io.Writer) {
-	responseBodyKey := md5.Sum(s.requestBodyKey)
-	responseBodyIV := md5.Sum(s.requestBodyIV)
-	s.responseBodyKey = responseBodyKey[:]
-	s.responseBodyIV = responseBodyIV[:]
+	responseBodyKey := md5.Sum(s.requestBodyKey[:])
+	responseBodyIV := md5.Sum(s.requestBodyIV[:])
+	s.responseBodyKey = responseBodyKey
+	s.responseBodyIV = responseBodyIV
 
-	aesStream := crypto.NewAesEncryptionStream(s.responseBodyKey, s.responseBodyIV)
+	aesStream := crypto.NewAesEncryptionStream(s.responseBodyKey[:], s.responseBodyIV[:])
 	encryptionWriter := crypto.NewCryptionWriter(aesStream, writer)
 	s.responseWriter = encryptionWriter
 
@@ -302,7 +302,7 @@ func (s *ServerSession) EncodeResponseHeader(header *protocol.ResponseHeader, wr
 func (s *ServerSession) EncodeResponseBody(request *protocol.RequestHeader, writer io.Writer) buf.Writer {
 	var sizeParser crypto.ChunkSizeEncoder = crypto.PlainChunkSizeParser{}
 	if request.Option.Has(protocol.RequestOptionChunkMasking) {
-		sizeParser = NewShakeSizeParser(s.responseBodyIV)
+		sizeParser = NewShakeSizeParser(s.responseBodyIV[:])
 	}
 	switch request.Security {
 	case protocol.SecurityType_NONE:
@@ -332,21 +332,21 @@ func (s *ServerSession) EncodeResponseBody(request *protocol.RequestHeader, writ
 
 		return buf.NewWriter(s.responseWriter)
 	case protocol.SecurityType_AES128_GCM:
-		block, _ := aes.NewCipher(s.responseBodyKey)
+		block, _ := aes.NewCipher(s.responseBodyKey[:])
 		aead, _ := cipher.NewGCM(block)
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
-			NonceGenerator:          GenerateChunkNonce(s.responseBodyIV, uint32(aead.NonceSize())),
+			NonceGenerator:          GenerateChunkNonce(s.responseBodyIV[:], uint32(aead.NonceSize())),
 			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType())
 	case protocol.SecurityType_CHACHA20_POLY1305:
-		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(s.responseBodyKey))
+		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(s.responseBodyKey[:]))
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
-			NonceGenerator:          GenerateChunkNonce(s.responseBodyIV, uint32(aead.NonceSize())),
+			NonceGenerator:          GenerateChunkNonce(s.responseBodyIV[:], uint32(aead.NonceSize())),
 			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType())

proxy/vmess/inbound/inbound.go

@@ -97,7 +97,7 @@ func (v *userByEmail) Remove(email string) bool {
 type Handler struct {
 	policyManager         core.PolicyManager
 	inboundHandlerManager core.InboundHandlerManager
-	clients               protocol.UserValidator
+	clients               *vmess.TimedUserValidator
 	usersByEmail          *userByEmail
 	detours               *DetourConfig
 	sessionHistory        *encoding.SessionHistory

proxy/vmess/vmess.go

@@ -42,7 +42,7 @@ type indexTimePair struct {
 	timeInc uint32
 }
 
-func NewTimedUserValidator(hasher protocol.IDHash) protocol.UserValidator {
+func NewTimedUserValidator(hasher protocol.IDHash) *TimedUserValidator {
 	tuv := &TimedUserValidator{
 		users:    make([]*user, 0, 16),
 		userHash: make(map[[16]byte]indexTimePair, 1024),