[proxy/vmess] Fix UDP over TCP fragmentation in VMESS + zero security

common/buf/buffer.go

@@ -13,14 +13,23 @@ const (
 
 var pool = bytespool.GetPool(Size)
 
+// ownership represents the data owner of the buffer.
+type ownership uint8
+
+const (
+	managed    ownership = 0
+	unmanaged  ownership = 1
+	bytespools ownership = 2
+)
+
 // 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
 }
 
 // New creates a Buffer with 0 length and 2K capacity.
@@ -30,12 +39,20 @@ func New() *Buffer {
 	}
 }
 
+// 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),
+		ownership: bytespools,
+	}
+}
+
 // FromBytes creates a Buffer with an existed bytearray
 func FromBytes(data []byte) *Buffer {
 	return &Buffer{
 		v:         data,
 		end:       int32(len(data)),
-		unmanaged: true,
+		ownership: unmanaged,
 	}
 }
 
@@ -49,14 +66,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()
-	pool.Put(p) // nolint: staticcheck
+	switch b.ownership {
+	case managed:
+		pool.Put(p) // nolint: staticcheck
+	case bytespools:
+		bytespool.Free(p) // nolint: staticcheck
+	}
 }
 
 // Clear clears the content of the buffer, results an empty buffer with
@@ -151,6 +173,14 @@ func (b *Buffer) Len() int32 {
 	return b.end - b.start
 }
 
+// Cap returns the capacity of the buffer content.
+func (b *Buffer) Cap() int32 {
+	if b == nil {
+		return 0
+	}
+	return int32(len(b.v))
+}
+
 // IsEmpty returns true if the buffer is empty.
 func (b *Buffer) IsEmpty() bool {
 	return b.Len() == 0

common/net/packetstream/packetstream.go

@@ -0,0 +1,62 @@
+package packetstream
+
+import (
+	"encoding/binary"
+	"io"
+
+	"github.com/v2fly/v2ray-core/v5/common/buf"
+)
+
+// PacketStreamWriter is a Writer that writes one packet with length info as header into a byte stream.
+type PacketStreamWriter struct {
+	io.Writer
+}
+
+// WriteMultiBuffer implements buf.Writer
+func (w *PacketStreamWriter) WriteMultiBuffer(multiBuffer buf.MultiBuffer) error {
+	defer buf.ReleaseMulti(multiBuffer)
+	for mb := multiBuffer; mb.Len() > 0; {
+		var payload *buf.Buffer
+		mb, payload = buf.SplitFirst(mb)
+		if payload.IsEmpty() {
+			continue
+		}
+
+		var lengthBuf [2]byte
+		binary.BigEndian.PutUint16(lengthBuf[:], uint16(payload.Len()))
+
+		buffer := buf.NewWithSize(2 + payload.Len())
+		defer buffer.Release()
+		if _, err := buffer.Write(lengthBuf[:]); err != nil {
+			return err
+		}
+		if _, err := buffer.Write(payload.Bytes()); err != nil {
+			return err
+		}
+		if _, err := w.Writer.Write(buffer.Bytes()); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// PacketStreamReader is a Reader that reads one complete packet every time from a byte stream.
+// It first reads the packet length header, then read payload of exact length from the reader.
+type PacketStreamReader struct {
+	io.Reader
+}
+
+// ReadMultiBuffer implements buf.Reader.
+func (r *PacketStreamReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
+	var lengthBuf [2]byte
+	if _, err := io.ReadFull(r, lengthBuf[:]); err != nil {
+		return nil, err
+	}
+	length := binary.BigEndian.Uint16(lengthBuf[:])
+
+	payload := buf.NewWithSize(int32(length))
+	if _, err := payload.ReadFullFrom(r, int32(length)); err != nil {
+		return nil, err
+	}
+	return buf.MultiBuffer{payload}, nil
+}

proxy/vmess/encoding/client.go

@@ -21,6 +21,7 @@ import (
 	"github.com/v2fly/v2ray-core/v5/common/crypto"
 	"github.com/v2fly/v2ray-core/v5/common/dice"
 	"github.com/v2fly/v2ray-core/v5/common/drain"
+	"github.com/v2fly/v2ray-core/v5/common/net/packetstream"
 	"github.com/v2fly/v2ray-core/v5/common/protocol"
 	"github.com/v2fly/v2ray-core/v5/common/serial"
 	"github.com/v2fly/v2ray-core/v5/proxy/vmess"
@@ -164,8 +165,11 @@ func (c *ClientSession) EncodeRequestBody(request *protocol.RequestHeader, write
 			}
 			return crypto.NewAuthenticationWriter(auth, sizeParser, writer, protocol.TransferTypePacket, padding), nil
 		}
-
-		return buf.NewWriter(writer), nil
+		if request.Command == protocol.RequestCommandUDP {
+			return &packetstream.PacketStreamWriter{Writer: writer}, nil
+		} else {
+			return buf.NewWriter(writer), nil
+		}
 	case protocol.SecurityType_LEGACY:
 		aesStream := crypto.NewAesEncryptionStream(c.requestBodyKey[:], c.requestBodyIV[:])
 		cryptionWriter := crypto.NewCryptionWriter(aesStream, writer)
@@ -339,8 +343,12 @@ func (c *ClientSession) DecodeResponseBody(request *protocol.RequestHeader, read
 
 			return crypto.NewAuthenticationReader(auth, sizeParser, reader, protocol.TransferTypePacket, padding), nil
 		}
+		if request.Command == protocol.RequestCommandUDP {
+			return buf.NewReader(&packetstream.PacketStreamReader{Reader: reader}), nil
+		} else {
+			return buf.NewReader(reader), nil
+		}
 
-		return buf.NewReader(reader), nil
 	case protocol.SecurityType_LEGACY:
 		if request.Option.Has(protocol.RequestOptionChunkStream) {
 			auth := &crypto.AEADAuthenticator{

proxy/vmess/encoding/server.go

@@ -20,6 +20,7 @@ import (
 	"github.com/v2fly/v2ray-core/v5/common/crypto"
 	"github.com/v2fly/v2ray-core/v5/common/drain"
 	"github.com/v2fly/v2ray-core/v5/common/net"
+	"github.com/v2fly/v2ray-core/v5/common/net/packetstream"
 	"github.com/v2fly/v2ray-core/v5/common/protocol"
 	"github.com/v2fly/v2ray-core/v5/common/task"
 	"github.com/v2fly/v2ray-core/v5/proxy/vmess"
@@ -333,7 +334,11 @@ func (s *ServerSession) DecodeRequestBody(request *protocol.RequestHeader, reade
 			}
 			return crypto.NewAuthenticationReader(auth, sizeParser, reader, protocol.TransferTypePacket, padding), nil
 		}
-		return buf.NewReader(reader), nil
+		if request.Command == protocol.RequestCommandUDP {
+			return &packetstream.PacketStreamReader{Reader: reader}, nil
+		} else {
+			return buf.NewReader(reader), nil
+		}
 
 	case protocol.SecurityType_LEGACY:
 		aesStream := crypto.NewAesDecryptionStream(s.requestBodyKey[:], s.requestBodyIV[:])
@@ -480,7 +485,11 @@ func (s *ServerSession) EncodeResponseBody(request *protocol.RequestHeader, writ
 			}
 			return crypto.NewAuthenticationWriter(auth, sizeParser, writer, protocol.TransferTypePacket, padding), nil
 		}
-		return buf.NewWriter(writer), nil
+		if request.Command == protocol.RequestCommandUDP {
+			return buf.NewWriter(&packetstream.PacketStreamWriter{Writer: writer}), nil
+		} else {
+			return buf.NewWriter(writer), nil
+		}
 
 	case protocol.SecurityType_LEGACY:
 		if request.Option.Has(protocol.RequestOptionChunkStream) {