Update network handler usages

Author: nekohasekai

brutal.go

@@ -7,7 +7,7 @@ import (
 	"github.com/sagernet/sing/common"
 	"github.com/sagernet/sing/common/buf"
 	E "github.com/sagernet/sing/common/exceptions"
-	"github.com/sagernet/sing/common/rw"
+	"github.com/sagernet/sing/common/varbin"
 )
 
 const (
@@ -32,7 +32,7 @@ func WriteBrutalResponse(writer io.Writer, receiveBPS uint64, ok bool, message s
 	if ok {
 		common.Must(binary.Write(buffer, binary.BigEndian, receiveBPS))
 	} else {
-		err := rw.WriteVString(buffer, message)
+		err := varbin.Write(buffer, binary.BigEndian, message)
 		if err != nil {
 			return err
 		}
@@ -52,7 +52,7 @@ func ReadBrutalResponse(reader io.Reader) (uint64, error) {
 		return receiveBPS, err
 	} else {
 		var message string
-		message, err = rw.ReadVString(reader)
+		message, err = varbin.ReadValue[string](reader, binary.BigEndian)
 		if err != nil {
 			return 0, err
 		}

go.mod

@@ -4,7 +4,7 @@ go 1.18
 
 require (
 	github.com/hashicorp/yamux v0.1.2
-	github.com/sagernet/sing v0.5.0
+	github.com/sagernet/sing v0.6.0-alpha.3
 	github.com/sagernet/smux v0.0.0-20231208180855-7041f6ea79e7
 	golang.org/x/net v0.31.0
 	golang.org/x/sys v0.27.0

go.sum

@@ -3,8 +3,10 @@ github.com/hashicorp/yamux v0.1.2 h1:XtB8kyFOyHXYVFnwT5C3+Bdo8gArse7j2AQ0DA0Uey8
 github.com/hashicorp/yamux v0.1.2/go.mod h1:C+zze2n6e/7wshOZep2A70/aQU6QBRWJO/G6FT1wIns=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/sagernet/sing v0.2.18/go.mod h1:OL6k2F0vHmEzXz2KW19qQzu172FDgSbUSODylighuVo=
-github.com/sagernet/sing v0.5.0 h1:soo2wVwLcieKWWKIksFNK6CCAojUgAppqQVwyRYGkEM=
-github.com/sagernet/sing v0.5.0/go.mod h1:ARkL0gM13/Iv5VCZmci/NuoOlePoIsW0m7BWfln/Hak=
+github.com/sagernet/sing v0.5.1-0.20241109034027-099899991126 h1:pLMpV9pEAinrS9R1n1JLcbNesCl369RfvyxnYCPrkbw=
+github.com/sagernet/sing v0.5.1-0.20241109034027-099899991126/go.mod h1:ARkL0gM13/Iv5VCZmci/NuoOlePoIsW0m7BWfln/Hak=
+github.com/sagernet/sing v0.6.0-alpha.3 h1:GLp9d6Gbt+Ioeplauuzojz1nY2J6moceVGYIOv/h5gA=
+github.com/sagernet/sing v0.6.0-alpha.3/go.mod h1:ARkL0gM13/Iv5VCZmci/NuoOlePoIsW0m7BWfln/Hak=
 github.com/sagernet/smux v0.0.0-20231208180855-7041f6ea79e7 h1:DImB4lELfQhplLTxeq2z31Fpv8CQqqrUwTbrIRumZqQ=
 github.com/sagernet/smux v0.0.0-20231208180855-7041f6ea79e7/go.mod h1:FP9X2xjT/Az1EsG/orYYoC+5MojWnuI7hrffz8fGwwo=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=

protocol.go

@@ -12,6 +12,7 @@ import (
 	M "github.com/sagernet/sing/common/metadata"
 	N "github.com/sagernet/sing/common/network"
 	"github.com/sagernet/sing/common/rw"
+	"github.com/sagernet/sing/common/varbin"
 )
 
 const (
@@ -41,14 +42,18 @@ type Request struct {
 }
 
 func ReadRequest(reader io.Reader) (*Request, error) {
-	version, err := rw.ReadByte(reader)
+	var (
+		version  byte
+		protocol byte
+	)
+	err := binary.Read(reader, binary.BigEndian, &version)
 	if err != nil {
 		return nil, err
 	}
 	if version < Version0 || version > Version1 {
 		return nil, E.New("unsupported version: ", version)
 	}
-	protocol, err := rw.ReadByte(reader)
+	err = binary.Read(reader, binary.BigEndian, &protocol)
 	if err != nil {
 		return nil, err
 	}
@@ -166,13 +171,12 @@ type StreamResponse struct {
 
 func ReadStreamResponse(reader io.Reader) (*StreamResponse, error) {
 	var response StreamResponse
-	status, err := rw.ReadByte(reader)
+	err := binary.Read(reader, binary.BigEndian, &response.Status)
 	if err != nil {
 		return nil, err
 	}
-	response.Status = status
-	if status == statusError {
-		response.Message, err = rw.ReadVString(reader)
+	if response.Status == statusError {
+		response.Message, err = varbin.ReadValue[string](reader, binary.BigEndian)
 		if err != nil {
 			return nil, err
 		}

server.go

@@ -13,15 +13,24 @@ import (
 	"github.com/sagernet/sing/common/task"
 )
 
+// Deprecated: Use ServiceHandlerEx instead.
+//
+//nolint:staticcheck
 type ServiceHandler interface {
 	N.TCPConnectionHandler
 	N.UDPConnectionHandler
 }
 
+type ServiceHandlerEx interface {
+	N.TCPConnectionHandlerEx
+	N.UDPConnectionHandlerEx
+}
+
 type Service struct {
 	newStreamContext func(context.Context, net.Conn) context.Context
 	logger           logger.ContextLogger
 	handler          ServiceHandler
+	handlerEx        ServiceHandlerEx
 	padding          bool
 	brutal           BrutalOptions
 }
@@ -30,6 +39,7 @@ type ServiceOptions struct {
 	NewStreamContext func(context.Context, net.Conn) context.Context
 	Logger           logger.ContextLogger
 	Handler          ServiceHandler
+	HandlerEx        ServiceHandlerEx
 	Padding          bool
 	Brutal           BrutalOptions
 }
@@ -42,12 +52,26 @@ func NewService(options ServiceOptions) (*Service, error) {
 		newStreamContext: options.NewStreamContext,
 		logger:           options.Logger,
 		handler:          options.Handler,
+		handlerEx:        options.HandlerEx,
 		padding:          options.Padding,
 		brutal:           options.Brutal,
 	}, nil
 }
 
+// Deprecated: Use NewConnectionEx instead.
 func (s *Service) NewConnection(ctx context.Context, conn net.Conn, metadata M.Metadata) error {
+	return s.newConnection(ctx, conn, metadata.Source)
+}
+
+func (s *Service) NewConnectionEx(ctx context.Context, conn net.Conn, source M.Socksaddr, destination M.Socksaddr, onClose N.CloseHandlerFunc) {
+	err := s.newConnection(ctx, conn, source)
+	N.CloseOnHandshakeFailure(conn, onClose, err)
+	if err != nil {
+		s.logger.ErrorContext(ctx, E.Cause(err, "process multiplex connection from ", source))
+	}
+}
+
+func (s *Service) newConnection(ctx context.Context, conn net.Conn, source M.Socksaddr) error {
 	request, err := ReadRequest(conn)
 	if err != nil {
 		return err
@@ -70,9 +94,10 @@ func (s *Service) NewConnection(ctx context.Context, conn net.Conn, metadata M.M
 			}
 			streamCtx := s.newStreamContext(ctx, stream)
 			go func() {
-				hErr := s.newConnection(streamCtx, conn, stream, metadata)
+				hErr := s.newSession(streamCtx, conn, stream, source)
 				if hErr != nil {
-					s.logger.ErrorContext(streamCtx, E.Cause(hErr, "handle connection"))
+					stream.Close()
+					s.logger.ErrorContext(streamCtx, E.Cause(hErr, "process multiplex stream"))
 				}
 			}()
 		}
@@ -83,13 +108,13 @@ func (s *Service) NewConnection(ctx context.Context, conn net.Conn, metadata M.M
 	return group.Run(ctx)
 }
 
-func (s *Service) newConnection(ctx context.Context, sessionConn net.Conn, stream net.Conn, metadata M.Metadata) error {
+func (s *Service) newSession(ctx context.Context, sessionConn net.Conn, stream net.Conn, source M.Socksaddr) error {
 	stream = &wrapStream{stream}
 	request, err := ReadStreamRequest(stream)
 	if err != nil {
 		return E.Cause(err, "read multiplex stream request")
 	}
-	metadata.Destination = request.Destination
+	destination := request.Destination
 	if request.Network == N.NetworkTCP {
 		conn := &serverConn{ExtendedConn: bufio.NewExtendedConn(stream)}
 		if request.Destination.Fqdn == BrutalExchangeDomain {
@@ -127,20 +152,28 @@ func (s *Service) newConnection(ctx context.Context, sessionConn net.Conn, strea
 			}
 			return nil
 		}
-		s.logger.InfoContext(ctx, "inbound multiplex connection to ", metadata.Destination)
-		s.handler.NewConnection(ctx, conn, metadata)
-		stream.Close()
+		s.logger.InfoContext(ctx, "inbound multiplex connection to ", destination)
+		if s.handler != nil {
+			//nolint:staticcheck
+			s.handler.NewConnection(ctx, conn, M.Metadata{Source: source, Destination: destination})
+		} else {
+			s.handlerEx.NewConnectionEx(ctx, conn, source, destination, nil)
+		}
 	} else {
 		var packetConn N.PacketConn
 		if !request.PacketAddr {
-			s.logger.InfoContext(ctx, "inbound multiplex packet connection to ", metadata.Destination)
+			s.logger.InfoContext(ctx, "inbound multiplex packet connection to ", destination)
 			packetConn = &serverPacketConn{ExtendedConn: bufio.NewExtendedConn(stream), destination: request.Destination}
 		} else {
 			s.logger.InfoContext(ctx, "inbound multiplex packet connection")
 			packetConn = &serverPacketAddrConn{ExtendedConn: bufio.NewExtendedConn(stream)}
 		}
-		s.handler.NewPacketConnection(ctx, packetConn, metadata)
-		stream.Close()
+		if s.handler != nil {
+			//nolint:staticcheck
+			s.handler.NewPacketConnection(ctx, packetConn, M.Metadata{Source: source, Destination: destination})
+		} else {
+			s.handlerEx.NewPacketConnectionEx(ctx, packetConn, source, destination, nil)
+		}
 	}
 	return nil
 }

server_conn.go

@@ -10,7 +10,7 @@ import (
 	"github.com/sagernet/sing/common/buf"
 	M "github.com/sagernet/sing/common/metadata"
 	N "github.com/sagernet/sing/common/network"
-	"github.com/sagernet/sing/common/rw"
+	"github.com/sagernet/sing/common/varbin"
 )
 
 type serverConn struct {
@@ -24,11 +24,11 @@ func (c *serverConn) NeedHandshake() bool {
 
 func (c *serverConn) HandshakeFailure(err error) error {
 	errMessage := err.Error()
-	buffer := buf.NewSize(1 + rw.UVariantLen(uint64(len(errMessage))) + len(errMessage))
+	buffer := buf.NewSize(1 + varbin.UvarintLen(uint64(len(errMessage))) + len(errMessage))
 	defer buffer.Release()
 	common.Must(
 		buffer.WriteByte(statusError),
-		rw.WriteVString(buffer, errMessage),
+		varbin.Write(buffer, binary.BigEndian, errMessage),
 	)
 	return common.Error(c.ExtendedConn.Write(buffer.Bytes()))
 }
@@ -88,11 +88,11 @@ func (c *serverPacketConn) NeedHandshake() bool {
 
 func (c *serverPacketConn) HandshakeFailure(err error) error {
 	errMessage := err.Error()
-	buffer := buf.NewSize(1 + rw.UVariantLen(uint64(len(errMessage))) + len(errMessage))
+	buffer := buf.NewSize(1 + varbin.UvarintLen(uint64(len(errMessage))) + len(errMessage))
 	defer buffer.Release()
 	common.Must(
 		buffer.WriteByte(statusError),
-		rw.WriteVString(buffer, errMessage),
+		varbin.Write(buffer, binary.BigEndian, errMessage),
 	)
 	return common.Error(c.ExtendedConn.Write(buffer.Bytes()))
 }
@@ -188,11 +188,11 @@ func (c *serverPacketAddrConn) NeedHandshake() bool {
 
 func (c *serverPacketAddrConn) HandshakeFailure(err error) error {
 	errMessage := err.Error()
-	buffer := buf.NewSize(1 + rw.UVariantLen(uint64(len(errMessage))) + len(errMessage))
+	buffer := buf.NewSize(1 + varbin.UvarintLen(uint64(len(errMessage))) + len(errMessage))
 	defer buffer.Release()
 	common.Must(
 		buffer.WriteByte(statusError),
-		rw.WriteVString(buffer, errMessage),
+		varbin.Write(buffer, binary.BigEndian, errMessage),
 	)
 	return common.Error(c.ExtendedConn.Write(buffer.Bytes()))
 }