🛒 Rework TCP response header processing so it's zero-copy as well

Author: database64128

client/client.go

@@ -124,45 +124,64 @@ type duplexConnAdaptor struct {
 	saltPool          *service.SaltPool
 }
 
-func (dc *duplexConnAdaptor) readHeader() error {
-	err := ss.ParseTCPRespHeader(dc.r, dc.w.Salt(), dc.cipherConfig)
+func (dc *duplexConnAdaptor) readHeader() ([]byte, error) {
+	initPayload, err := ss.ParseTCPRespHeader(dc.r, dc.w.Salt(), dc.cipherConfig)
 	if err != nil {
 		dc.Close()
-		return err
+		return nil, err
 	}
 
 	// 2022 spec: check salt
 	if dc.cipherConfig.IsSpec2022 && !dc.saltPool.Add(*(*[32]byte)(dc.r.Salt())) {
 		io.Copy(io.Discard, dc.r)
 		dc.Close()
-		return ErrRepeatedSalt
+		return nil, ErrRepeatedSalt
 	}
 
-	return nil
+	return initPayload, nil
 }
 
-func (dc *duplexConnAdaptor) Read(b []byte) (int, error) {
+func (dc *duplexConnAdaptor) Read(b []byte) (n int, err error) {
 	if !dc.isHeaderProcessed {
-		err := dc.readHeader()
+		var initPayload []byte
+		initPayload, err = dc.readHeader()
 		if err != nil {
 			return 0, err
 		}
 		dc.isHeaderProcessed = true
+
+		if len(initPayload) > 0 {
+			n = copy(b, initPayload)
+			if n < len(initPayload) {
+				err = io.ErrShortBuffer
+			}
+			return
+		}
 	}
 
 	return dc.r.Read(b)
 }
 
-func (dc *duplexConnAdaptor) WriteTo(w io.Writer) (int64, error) {
+func (dc *duplexConnAdaptor) WriteTo(w io.Writer) (n int64, err error) {
 	if !dc.isHeaderProcessed {
-		err := dc.readHeader()
+		initPayload, err := dc.readHeader()
 		if err != nil {
 			return 0, err
 		}
 		dc.isHeaderProcessed = true
+
+		if len(initPayload) > 0 {
+			wn, err := w.Write(initPayload)
+			n = int64(wn)
+			if err != nil {
+				return n, err
+			}
+		}
 	}
 
-	return io.Copy(w, dc.r)
+	cn, err := io.Copy(w, dc.r)
+	n += cn
+	return
 }
 
 func (dc *duplexConnAdaptor) CloseRead() error {

service/tcp.go

@@ -278,7 +278,7 @@ func (s *tcpService) handleConnection(clientTCPConn tfo.Conn) {
 		}
 
 		ssr := ss.NewShadowsocksReader(clientReader, cipherEntry.Cipher)
-		tgtAddr, initPayload, err := ss.ParseTCPReqHeader(ssr, cipherEntry.Cipher.Config(), ss.HeaderTypeClientStream)
+		tgtAddr, initPayload, err := ss.ParseTCPReqHeader(ssr, cipherEntry.Cipher.Config())
 		if err != nil {
 			logger.Warn("Failed to parse header",
 				zap.Stringer("clientConnLocalAddress", clientLocalAddr),

shadowsocks/header.go

@@ -5,7 +5,6 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
-	"io"
 	"math/rand"
 	"net"
 	"strconv"
@@ -50,7 +49,7 @@ var (
 // For Shadowsocks 2022, the first payload chunk MUST contain
 // either a non-zero-length padding or initial payload, or both (not recommended client behavior but allowed).
 // If the padding length is 0 and there's no initial payload, an error is returned.
-func ParseTCPReqHeader(r Reader, cipherConfig CipherConfig, htype byte) (string, []byte, error) {
+func ParseTCPReqHeader(r Reader, cipherConfig CipherConfig) (string, []byte, error) {
 	// Read first payload chunk.
 	if err := r.EnsureLeftover(); err != nil {
 		return "", nil, err
@@ -65,7 +64,7 @@ func ParseTCPReqHeader(r Reader, cipherConfig CipherConfig, htype byte) (string,
 		}
 
 		// Verify type
-		if b[0] != htype {
+		if b[0] != HeaderTypeClientStream {
 			return "", nil, ErrTypeMismatch
 		}
 
@@ -139,39 +138,42 @@ func WriteTCPReqHeader(dst, socksaddr []byte, addPadding bool, cipherConfig Ciph
 	return
 }
 
-func ParseTCPRespHeader(r io.Reader, clientSalt []byte, cipherConfig CipherConfig) error {
+func ParseTCPRespHeader(r Reader, clientSalt []byte, cipherConfig CipherConfig) ([]byte, error) {
 	if !cipherConfig.IsSpec2022 {
-		return nil
+		return nil, nil
 	}
 
-	b := make([]byte, 1+8+len(clientSalt))
+	// Read first payload chunk.
+	if err := r.EnsureLeftover(); err != nil {
+		return nil, err
+	}
+	b := r.LeftoverZeroCopy()
 
-	// Read response header
-	_, err := io.ReadFull(r, b)
-	if err != nil {
-		return fmt.Errorf("failed to read response header: %w", err)
+	headerLen := 1 + 8 + len(clientSalt)
+	if len(b) < headerLen {
+		return nil, ErrIncompleteHeaderInFirstChunk
 	}
 
 	// Verify type
 	if b[0] != HeaderTypeServerStream {
-		return ErrTypeMismatch
+		return nil, ErrTypeMismatch
 	}
 
 	// Verify timestamp
 	epoch := int64(binary.BigEndian.Uint64(b[1 : 1+8]))
 	nowEpoch := time.Now().Unix()
 	diff := epoch - nowEpoch
 	if diff < -30 || diff > 30 {
-		return ErrBadTimestamp
+		return nil, ErrBadTimestamp
 	}
 
 	// Verify client salt
-	n := bytes.Compare(clientSalt, b[1+8:])
+	n := bytes.Compare(clientSalt, b[1+8:headerLen])
 	if n != 0 {
-		return ErrClientSaltMismatch
+		return nil, ErrClientSaltMismatch
 	}
 
-	return nil
+	return b[headerLen:], nil
 }
 
 func WriteTCPRespHeader(dst, clientSalt []byte, cipherConfig CipherConfig) (n int) {

shadowsocks/packet.go

@@ -27,7 +27,6 @@ import (
 	wgreplay "golang.zx2c4.com/wireguard/replay"
 )
 
-// ErrShortPacket is identical to shadowaead.ErrShortPacket
 var ErrShortPacket = errors.New("short packet")
 
 // Pack encrypts a Shadowsocks-UDP packet and returns a slice containing the encrypted packet.