refactor: Implement extended IO

transport/vless/conn.go

@@ -1,7 +1,6 @@
 package vless
 
 import (
-	"bytes"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -9,12 +8,16 @@ import (
 	"net"
 
 	"github.com/gofrs/uuid"
+	"github.com/sagernet/sing/common"
+	"github.com/sagernet/sing/common/buf"
+	"github.com/sagernet/sing/common/bufio"
+	"github.com/sagernet/sing/common/network"
 	xtls "github.com/xtls/go"
 	"google.golang.org/protobuf/proto"
 )
 
 type Conn struct {
-	net.Conn
+	network.ExtendedConn
 	dst      *DstAddr
 	id       *uuid.UUID
 	addons   *Addons
@@ -23,57 +26,82 @@ type Conn struct {
 
 func (vc *Conn) Read(b []byte) (int, error) {
 	if vc.received {
-		return vc.Conn.Read(b)
+		return vc.ExtendedConn.Read(b)
 	}
 
 	if err := vc.recvResponse(); err != nil {
 		return 0, err
 	}
 	vc.received = true
-	return vc.Conn.Read(b)
+	return vc.ExtendedConn.Read(b)
 }
 
-func (vc *Conn) sendRequest() error {
-	buf := &bytes.Buffer{}
+func (vc *Conn) ReadBuffer(buffer *buf.Buffer) error {
+	if vc.received {
+		return vc.ExtendedConn.ReadBuffer(buffer)
+	}
 
-	buf.WriteByte(Version)   // protocol version
-	buf.Write(vc.id.Bytes()) // 16 bytes of uuid
+	if err := vc.recvResponse(); err != nil {
+		return err
+	}
+	vc.received = true
+	return vc.ExtendedConn.ReadBuffer(buffer)
+}
 
+func (vc *Conn) sendRequest() (err error) {
+	requestLen := 1  // protocol version
+	requestLen += 16 // UUID
+	requestLen += 1  // addons length
+	var addonsBytes []byte
 	if vc.addons != nil {
-		bytes, err := proto.Marshal(vc.addons)
+		addonsBytes, err = proto.Marshal(vc.addons)
 		if err != nil {
 			return err
 		}
-
-		buf.WriteByte(byte(len(bytes)))
-		buf.Write(bytes)
-	} else {
-		buf.WriteByte(0) // addon data length. 0 means no addon data
 	}
+	requestLen += len(addonsBytes)
+	requestLen += 1 // command
+	if !vc.dst.Mux {
+		requestLen += 2 // port
+		requestLen += 1 // addr type
+		requestLen += len(vc.dst.Addr)
+	}
+	_buffer := buf.StackNewSize(requestLen)
+	defer common.KeepAlive(_buffer)
+	buffer := common.Dup(_buffer)
+	defer buffer.Release()
+
+	common.Must(
+		buffer.WriteByte(Version),                 // protocol version
+		common.Error(buffer.Write(vc.id.Bytes())), // 16 bytes of uuid
+		buffer.WriteByte(byte(len(addonsBytes))),
+		common.Error(buffer.Write(addonsBytes)),
+	)
 
 	if vc.dst.Mux {
-		buf.WriteByte(CommandMux)
+		common.Must(buffer.WriteByte(CommandMux))
 	} else {
 		if vc.dst.UDP {
-			buf.WriteByte(CommandUDP)
+			common.Must(buffer.WriteByte(CommandUDP))
 		} else {
-			buf.WriteByte(CommandTCP)
+			common.Must(buffer.WriteByte(CommandTCP))
 		}
 
-		// Port AddrType Addr
-		binary.Write(buf, binary.BigEndian, vc.dst.Port)
-		buf.WriteByte(vc.dst.AddrType)
-		buf.Write(vc.dst.Addr)
+		binary.BigEndian.PutUint16(buffer.Extend(2), vc.dst.Port)
+		common.Must(
+			buffer.WriteByte(vc.dst.AddrType),
+			common.Error(buffer.Write(vc.dst.Addr)),
+		)
 	}
 
-	_, err := vc.Conn.Write(buf.Bytes())
-	return err
+	_, err = vc.ExtendedConn.Write(buffer.Bytes())
+	return
 }
 
 func (vc *Conn) recvResponse() error {
 	var err error
-	buf := make([]byte, 1)
-	_, err = io.ReadFull(vc.Conn, buf)
+	var buf [1]byte
+	_, err = io.ReadFull(vc.ExtendedConn, buf[:])
 	if err != nil {
 		return err
 	}
@@ -82,25 +110,32 @@ func (vc *Conn) recvResponse() error {
 		return errors.New("unexpected response version")
 	}
 
-	_, err = io.ReadFull(vc.Conn, buf)
+	_, err = io.ReadFull(vc.ExtendedConn, buf[:])
 	if err != nil {
 		return err
 	}
 
 	length := int64(buf[0])
 	if length != 0 { // addon data length > 0
-		io.CopyN(io.Discard, vc.Conn, length) // just discard
+		io.CopyN(io.Discard, vc.ExtendedConn, length) // just discard
 	}
 
 	return nil
 }
 
+func (vc *Conn) Upstream() any {
+	if wrapper, ok := vc.ExtendedConn.(*bufio.ExtendedConnWrapper); ok {
+		return wrapper.Conn
+	}
+	return vc.ExtendedConn
+}
+
 // newConn return a Conn instance
 func newConn(conn net.Conn, client *Client, dst *DstAddr) (*Conn, error) {
 	c := &Conn{
-		Conn: conn,
-		id:   client.uuid,
-		dst:  dst,
+		ExtendedConn: bufio.NewExtendedConn(conn),
+		id:           client.uuid,
+		dst:          dst,
 	}
 
 	if !dst.UDP && client.Addons != nil {