Changes before merging

#4952 (comment)

Author: RPRX

proxy/vless/encryption/client.go

@@ -66,7 +66,7 @@ func (i *ClientInstance) Handshake(conn net.Conn) (*CommonConn, error) {
 	if i.NfsPKeys == nil {
 		return nil, errors.New("uninitialized")
 	}
-	c := &CommonConn{Conn: conn}
+	c := NewCommonConn(conn)
 
 	ivAndRealysLength := 16 + i.RelaysLength
 	pfsKeyExchangeLength := 18 + 1184 + 32 + 16

proxy/vless/encryption/common.go

@@ -8,12 +8,19 @@ import (
 	"io"
 	"net"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/xtls/xray-core/common/errors"
 	"lukechampine.com/blake3"
 )
 
+var OutBytesPool = sync.Pool{
+	New: func() any {
+		return make([]byte, 5+8192+16)
+	},
+}
+
 type CommonConn struct {
 	net.Conn
 	Client      *ClientInstance
@@ -22,35 +29,44 @@ type CommonConn struct {
 	GCM         *GCM
 	PeerGCM     *GCM
 	PeerPadding []byte
+	PeerInBytes []byte
 	PeerCache   []byte
 }
 
+func NewCommonConn(conn net.Conn) *CommonConn {
+	return &CommonConn{
+		Conn:        conn,
+		PeerInBytes: make([]byte, 5+17000), // no need to use sync.Pool, because we are always reading
+	}
+}
+
 func (c *CommonConn) Write(b []byte) (int, error) {
 	if len(b) == 0 {
 		return 0, nil
 	}
-	var data []byte
+	outBytes := OutBytesPool.Get().([]byte)
+	defer OutBytesPool.Put(outBytes)
 	for n := 0; n < len(b); {
 		b := b[n:]
 		if len(b) > 8192 {
 			b = b[:8192] // for avoiding another copy() in peer's Read()
 		}
 		n += len(b)
-		data = make([]byte, 5+len(b)+16)
-		EncodeHeader(data, len(b)+16)
+		headerAndData := outBytes[:5+len(b)+16]
+		EncodeHeader(headerAndData, len(b)+16)
 		max := false
 		if bytes.Equal(c.GCM.Nonce[:], MaxNonce) {
 			max = true
 		}
-		c.GCM.Seal(data[:5], nil, b, data[:5])
+		c.GCM.Seal(headerAndData[:5], nil, b, headerAndData[:5])
 		if max {
-			c.GCM = NewGCM(data[5:], c.UnitedKey)
+			c.GCM = NewGCM(headerAndData, c.UnitedKey)
 		}
 		if c.PreWrite != nil {
-			data = append(c.PreWrite, data...)
+			headerAndData = append(c.PreWrite, headerAndData...)
 			c.PreWrite = nil
 		}
-		if _, err := c.Conn.Write(data); err != nil {
+		if _, err := c.Conn.Write(headerAndData); err != nil {
 			return 0, err
 		}
 	}
@@ -80,14 +96,18 @@ func (c *CommonConn) Read(b []byte) (int, error) {
 		}
 		c.PeerPadding = nil
 	}
-	if len(c.PeerCache) != 0 {
+	if len(c.PeerCache) > 0 {
 		n := copy(b, c.PeerCache)
 		c.PeerCache = c.PeerCache[n:]
 		return n, nil
 	}
-	h, l, err := ReadAndDecodeHeader(c.Conn) // l: 17~17000
+	peerHeader := c.PeerInBytes[:5]
+	if _, err := io.ReadFull(c.Conn, peerHeader); err != nil {
+		return 0, err
+	}
+	l, err := DecodeHeader(c.PeerInBytes[:5]) // l: 17~17000
 	if err != nil {
-		if c.Client != nil && strings.HasPrefix(err.Error(), "invalid header: ") { // client's 0-RTT
+		if c.Client != nil && strings.Contains(err.Error(), "invalid header: ") { // client's 0-RTT
 			c.Client.RWLock.Lock()
 			if bytes.HasPrefix(c.UnitedKey, c.Client.PfsKey) {
 				c.Client.Expire = time.Now() // expired
@@ -98,7 +118,7 @@ func (c *CommonConn) Read(b []byte) (int, error) {
 		return 0, err
 	}
 	c.Client = nil
-	peerData := make([]byte, l)
+	peerData := c.PeerInBytes[5 : 5+l]
 	if _, err := io.ReadFull(c.Conn, peerData); err != nil {
 		return 0, err
 	}
@@ -108,9 +128,9 @@ func (c *CommonConn) Read(b []byte) (int, error) {
 	}
 	var newGCM *GCM
 	if bytes.Equal(c.PeerGCM.Nonce[:], MaxNonce) {
-		newGCM = NewGCM(peerData, c.UnitedKey)
+		newGCM = NewGCM(c.PeerInBytes[:5+l], c.UnitedKey)
 	}
-	_, err = c.PeerGCM.Open(dst[:0], nil, peerData, h)
+	_, err = c.PeerGCM.Open(dst[:0], nil, peerData, peerHeader)
 	if newGCM != nil {
 		c.PeerGCM = newGCM
 	}
@@ -190,23 +210,3 @@ func DecodeHeader(h []byte) (l int, err error) {
 	}
 	return
 }
-
-func ReadAndDecodeHeader(conn net.Conn) (h []byte, l int, err error) {
-	h = make([]byte, 5)
-	if _, err = io.ReadFull(conn, h); err != nil {
-		return
-	}
-	l, err = DecodeHeader(h)
-	return
-}
-
-func ReadAndDiscardPaddings(conn net.Conn) (h []byte, l int, err error) {
-	for {
-		if h, l, err = ReadAndDecodeHeader(conn); err != nil {
-			return
-		}
-		if _, err = io.ReadFull(conn, make([]byte, l)); err != nil {
-			return
-		}
-	}
-}

proxy/vless/encryption/server.go

@@ -102,7 +102,7 @@ func (i *ServerInstance) Handshake(conn net.Conn) (*CommonConn, error) {
 	if i.NfsSKeys == nil {
 		return nil, errors.New("uninitialized")
 	}
-	c := &CommonConn{Conn: conn}
+	c := NewCommonConn(conn)
 
 	ivAndRelays := make([]byte, 16+i.RelaysLength)
 	if _, err := io.ReadFull(conn, ivAndRelays); err != nil {