QUIC sniffer: Full support for handling multiple initial packets

app/dispatcher/default.go

@@ -33,23 +33,21 @@ type cachedReader struct {
 	cache  buf.MultiBuffer
 }
 
-func (r *cachedReader) Cache(b *buf.Buffer) {
-	mb, _ := r.reader.ReadMultiBufferTimeout(time.Millisecond * 100)
+func (r *cachedReader) Cache(b *buf.Buffer, deadline time.Duration) error {
+	mb, err := r.reader.ReadMultiBufferTimeout(deadline)
+	if err != nil {
+		return err
+	}
 	r.Lock()
 	if !mb.IsEmpty() {
 		r.cache, _ = buf.MergeMulti(r.cache, mb)
 	}
-	cacheLen := r.cache.Len()
-	if cacheLen <= b.Cap() {
-		b.Clear()
-	} else {
-		b.Release()
-		*b = *buf.NewWithSize(cacheLen)
-	}
-	rawBytes := b.Extend(cacheLen)
+	b.Clear()
+	rawBytes := b.Extend(b.Cap())
 	n := r.cache.Copy(rawBytes)
 	b.Resize(0, int32(n))
 	r.Unlock()
+	return nil
 }
 
 func (r *cachedReader) readInternal() buf.MultiBuffer {
@@ -355,7 +353,7 @@ func (d *DefaultDispatcher) DispatchLink(ctx context.Context, destination net.De
 }
 
 func sniffer(ctx context.Context, cReader *cachedReader, metadataOnly bool, network net.Network) (SniffResult, error) {
-	payload := buf.New()
+	payload := buf.NewWithSize(32767)
 	defer payload.Release()
 
 	sniffer := NewSniffer(ctx)
@@ -367,26 +365,33 @@ func sniffer(ctx context.Context, cReader *cachedReader, metadataOnly bool, netw
 	}
 
 	contentResult, contentErr := func() (SniffResult, error) {
+		cacheDeadline := 200 * time.Millisecond
 		totalAttempt := 0
 		for {
 			select {
 			case <-ctx.Done():
 				return nil, ctx.Err()
 			default:
-				totalAttempt++
-				if totalAttempt > 2 {
-					return nil, errSniffingTimeout
-				}
+				cachingStartingTimeStamp := time.Now()
+				cacheErr := cReader.Cache(payload, cacheDeadline)
+				cachingTimeElapsed := time.Since(cachingStartingTimeStamp)
+				cacheDeadline -= cachingTimeElapsed
 
-				cReader.Cache(payload)
 				if !payload.IsEmpty() {
 					result, err := sniffer.Sniff(ctx, payload.Bytes(), network)
-					if err != common.ErrNoClue {
+					switch err {
+					case common.ErrNoClue: // No Clue: protocol not matches, and sniffer cannot determine whether there will be a match or not
+						totalAttempt++
+					case protocol.ErrProtoNeedMoreData: // Protocol Need More Data: protocol matches, but need more data to complete sniffing
+						if cacheErr != nil { // Cache error (e.g. timeout) counts for failed attempt
+							totalAttempt++
+						}
+					default:
 						return result, err
 					}
 				}
-				if payload.IsFull() {
-					return nil, errUnknownContent
+				if totalAttempt >= 2 || cacheDeadline <= 0 {
+					return nil, errSniffingTimeout
 				}
 			}
 		}

app/dispatcher/sniffer.go

@@ -6,6 +6,7 @@ import (
 	"github.com/xtls/xray-core/common"
 	"github.com/xtls/xray-core/common/errors"
 	"github.com/xtls/xray-core/common/net"
+	"github.com/xtls/xray-core/common/protocol"
 	"github.com/xtls/xray-core/common/protocol/bittorrent"
 	"github.com/xtls/xray-core/common/protocol/http"
 	"github.com/xtls/xray-core/common/protocol/quic"
@@ -58,14 +59,17 @@ var errUnknownContent = errors.New("unknown content")
 func (s *Sniffer) Sniff(c context.Context, payload []byte, network net.Network) (SniffResult, error) {
 	var pendingSniffer []protocolSnifferWithMetadata
 	for _, si := range s.sniffer {
-		s := si.protocolSniffer
+		protocolSniffer := si.protocolSniffer
 		if si.metadataSniffer || si.network != network {
 			continue
 		}
-		result, err := s(c, payload)
+		result, err := protocolSniffer(c, payload)
 		if err == common.ErrNoClue {
 			pendingSniffer = append(pendingSniffer, si)
 			continue
+		} else if err == protocol.ErrProtoNeedMoreData { // Sniffer protocol matched, but need more data to complete sniffing
+			s.sniffer = []protocolSnifferWithMetadata{si}
+			return nil, err
 		}
 
 		if err == nil && result != nil {

common/buf/buffer.go

@@ -15,14 +15,23 @@ const (
 
 var pool = bytespool.GetPool(Size)
 
+// ownership represents the data owner of the buffer.
+type ownership uint8
+
+const (
+	managed ownership = iota
+	unmanaged
+	bytespools
+)
+
 // Buffer is a recyclable allocation of a byte array. Buffer.Release() recycles
 // the buffer into an internal buffer pool, in order to recreate a buffer more
 // quickly.
 type Buffer struct {
 	v         []byte
 	start     int32
 	end       int32
-	unmanaged bool
+	ownership ownership
 	UDP       *net.Destination
 }
 
@@ -62,7 +71,7 @@ func FromBytes(b []byte) *Buffer {
 	return &Buffer{
 		v:         b,
 		end:       int32(len(b)),
-		unmanaged: true,
+		ownership: unmanaged,
 	}
 }
 
@@ -83,16 +92,19 @@ func StackNew() Buffer {
 
 // Release recycles the buffer into an internal buffer pool.
 func (b *Buffer) Release() {
-	if b == nil || b.v == nil || b.unmanaged {
+	if b == nil || b.v == nil || b.ownership == unmanaged {
 		return
 	}
 
 	p := b.v
 	b.v = nil
 	b.Clear()
 
-	if cap(p) == Size {
-		pool.Put(p)
+	switch b.ownership {
+	case managed:
+		pool.Put(p) // nolint: staticcheck
+	case bytespools:
+		bytespool.Free(p) // nolint: staticcheck
 	}
 	b.UDP = nil
 }
@@ -104,6 +116,16 @@ func (b *Buffer) Clear() {
 	b.end = 0
 }
 
+// Wipe sets all underlying bytes of the buffer to 0x0, results an empty buffer with
+// Len() = 0
+func (b *Buffer) Wipe() {
+	for i := range cap(b.v) {
+		b.v[i] = 0x0
+	}
+	b.start = 0
+	b.end = 0
+}
+
 // Byte returns the bytes at index.
 func (b *Buffer) Byte(index int32) byte {
 	return b.v[b.start+index]
@@ -218,7 +240,8 @@ func (b *Buffer) Cap() int32 {
 // NewWithSize creates a Buffer with 0 length and capacity with at least the given size.
 func NewWithSize(size int32) *Buffer {
 	return &Buffer{
-		v: bytespool.Alloc(size),
+		v:         bytespool.Alloc(size),
+		ownership: bytespools,
 	}
 }
 

common/protocol/protocol.go

@@ -1 +1,7 @@
 package protocol // import "github.com/xtls/xray-core/common/protocol"
+
+import (
+	"errors"
+)
+
+var ErrProtoNeedMoreData = errors.New("protocol matches, but need more data to complete sniffing")

common/protocol/quic/sniff.go

@@ -11,8 +11,8 @@ import (
 	"github.com/quic-go/quic-go/quicvarint"
 	"github.com/xtls/xray-core/common"
 	"github.com/xtls/xray-core/common/buf"
-	"github.com/xtls/xray-core/common/bytespool"
 	"github.com/xtls/xray-core/common/errors"
+	"github.com/xtls/xray-core/common/protocol"
 	ptls "github.com/xtls/xray-core/common/protocol/tls"
 	"golang.org/x/crypto/hkdf"
 )
@@ -48,6 +48,9 @@ var (
 )
 
 func SniffQUIC(b []byte) (resultReturn *SniffHeader, errorReturn error) {
+	if len(b) == 0 {
+		return nil, common.ErrNoClue
+	}
 	// In extremely rare cases, this sniffer may cause slice error
 	// and we set recover() here to prevent crash.
 	// TODO: Thoroughly fix this panic
@@ -61,8 +64,11 @@ func SniffQUIC(b []byte) (resultReturn *SniffHeader, errorReturn error) {
 
 	// Crypto data separated across packets
 	cryptoLen := 0
-	cryptoData := bytespool.Alloc(int32(len(b)))
-	defer bytespool.Free(cryptoData)
+	cryptoDataBuf := buf.NewWithSize(32767)
+	cryptoDataBuf.Wipe()
+	defer cryptoDataBuf.Release()
+	cache := buf.New()
+	defer cache.Release()
 
 	// Parse QUIC packets
 	for len(b) > 0 {
@@ -105,13 +111,15 @@ func SniffQUIC(b []byte) (resultReturn *SniffHeader, errorReturn error) {
 			return nil, errNotQuic
 		}
 
-		tokenLen, err := quicvarint.Read(buffer)
-		if err != nil || tokenLen > uint64(len(b)) {
-			return nil, errNotQuic
-		}
+		if isQuicInitial { // Only initial packets have token, see https://datatracker.ietf.org/doc/html/rfc9000#section-17.2.2
+			tokenLen, err := quicvarint.Read(buffer)
+			if err != nil || tokenLen > uint64(len(b)) {
+				return nil, errNotQuic
+			}
 
-		if _, err = buffer.ReadBytes(int32(tokenLen)); err != nil {
-			return nil, errNotQuic
+			if _, err = buffer.ReadBytes(int32(tokenLen)); err != nil {
+				return nil, errNotQuic
+			}
 		}
 
 		packetLen, err := quicvarint.Read(buffer)
@@ -130,9 +138,6 @@ func SniffQUIC(b []byte) (resultReturn *SniffHeader, errorReturn error) {
 			continue
 		}
 
-		origPNBytes := make([]byte, 4)
-		copy(origPNBytes, b[hdrLen:hdrLen+4])
-
 		var salt []byte
 		if versionNumber == version1 {
 			salt = quicSalt
@@ -147,44 +152,34 @@ func SniffQUIC(b []byte) (resultReturn *SniffHeader, errorReturn error) {
 			return nil, err
 		}
 
-		cache := buf.New()
-		defer cache.Release()
-
+		cache.Clear()
 		mask := cache.Extend(int32(block.BlockSize()))
 		block.Encrypt(mask, b[hdrLen+4:hdrLen+4+16])
 		b[0] ^= mask[0] & 0xf
-		for i := range b[hdrLen : hdrLen+4] {
+		packetNumberLength := int(b[0]&0x3 + 1)
+		for i := range packetNumberLength {
 			b[hdrLen+i] ^= mask[i+1]
 		}
-		packetNumberLength := b[0]&0x3 + 1
-		if packetNumberLength != 1 {
-			return nil, errNotQuicInitial
-		}
-		var packetNumber uint32
-		{
-			n, err := buffer.ReadByte()
-			if err != nil {
-				return nil, err
-			}
-			packetNumber = uint32(n)
-		}
-
-		extHdrLen := hdrLen + int(packetNumberLength)
-		copy(b[extHdrLen:hdrLen+4], origPNBytes[packetNumberLength:])
-		data := b[extHdrLen : int(packetLen)+hdrLen]
 
 		key := hkdfExpandLabel(crypto.SHA256, secret, []byte{}, "quic key", 16)
 		iv := hkdfExpandLabel(crypto.SHA256, secret, []byte{}, "quic iv", 12)
 		cipher := AEADAESGCMTLS13(key, iv)
+
 		nonce := cache.Extend(int32(cipher.NonceSize()))
-		binary.BigEndian.PutUint64(nonce[len(nonce)-8:], uint64(packetNumber))
+		_, err = buffer.Read(nonce[len(nonce)-packetNumberLength:])
+		if err != nil {
+			return nil, err
+		}
+
+		extHdrLen := hdrLen + packetNumberLength
+		data := b[extHdrLen : int(packetLen)+hdrLen]
 		decrypted, err := cipher.Open(b[extHdrLen:extHdrLen], nonce, data, b[:extHdrLen])
 		if err != nil {
 			return nil, err
 		}
 		buffer = buf.FromBytes(decrypted)
-		for i := 0; !buffer.IsEmpty(); i++ {
-			frameType := byte(0x0) // Default to PADDING frame
+		for !buffer.IsEmpty() {
+			frameType, _ := buffer.ReadByte()
 			for frameType == 0x0 && !buffer.IsEmpty() {
 				frameType, _ = buffer.ReadByte()
 			}
@@ -235,14 +230,14 @@ func SniffQUIC(b []byte) (resultReturn *SniffHeader, errorReturn error) {
 				}
 				if cryptoLen < int(offset+length) {
 					cryptoLen = int(offset + length)
-					if len(cryptoData) < cryptoLen {
-						newCryptoData := bytespool.Alloc(int32(cryptoLen))
-						copy(newCryptoData, cryptoData)
-						bytespool.Free(cryptoData)
-						cryptoData = newCryptoData
+					if cryptoDataBuf.Cap() < int32(cryptoLen) {
+						return nil, io.ErrShortBuffer
+					}
+					if cryptoDataBuf.Len() != int32(cryptoLen) {
+						cryptoDataBuf.Extend(int32(cryptoLen) - cryptoDataBuf.Len())
 					}
 				}
-				if _, err := buffer.Read(cryptoData[offset : offset+length]); err != nil { // Field: Crypto Data
+				if _, err := buffer.Read(cryptoDataBuf.BytesRange(int32(offset), int32(offset+length))); err != nil { // Field: Crypto Data
 					return nil, io.ErrUnexpectedEOF
 				}
 			case 0x1c: // CONNECTION_CLOSE frame, only 0x1c is permitted in initial packet
@@ -267,7 +262,7 @@ func SniffQUIC(b []byte) (resultReturn *SniffHeader, errorReturn error) {
 		}
 
 		tlsHdr := &ptls.SniffHeader{}
-		err = ptls.ReadClientHello(cryptoData[:cryptoLen], tlsHdr)
+		err = ptls.ReadClientHello(cryptoDataBuf.BytesRange(0, int32(cryptoLen)), tlsHdr)
 		if err != nil {
 			// The crypto data may have not been fully recovered in current packets,
 			// So we continue to sniff rest packets.
@@ -276,7 +271,8 @@ func SniffQUIC(b []byte) (resultReturn *SniffHeader, errorReturn error) {
 		}
 		return &SniffHeader{domain: tlsHdr.Domain()}, nil
 	}
-	return nil, common.ErrNoClue
+	// All payload is parsed as valid QUIC packets, but we need more packets for crypto data to read client hello.
+	return nil, protocol.ErrProtoNeedMoreData
 }
 
 func hkdfExpandLabel(hash crypto.Hash, secret, context []byte, label string, length int) []byte {