fix false success from SOCKS server when Dispatch() fails

Postpone SOCKS server ok reply to a TCP Connect command until Dispatch()
actually creates an upstream Link. Previously it was sent immediately
inside Handshake(). UDP works as before.

Author: basinilya

proxy/socks/protocol.go

@@ -30,6 +30,7 @@ const (
 	authNoMatchingMethod = 0xFF
 
 	statusSuccess       = 0x00
+	statusConnRefused   = 0x05
 	statusCmdNotSupport = 0x07
 )
 
@@ -40,10 +41,12 @@ var addrParser = protocol.NewAddressParser(
 )
 
 type ServerSession struct {
-	config        *ServerConfig
-	address       net.Address
-	port          net.Port
-	clientAddress net.Address
+	config         *ServerConfig
+	address        net.Address
+	port           net.Port
+	clientAddress  net.Address
+	deferLastReply bool
+	flushLastReply func(bool) error
 }
 
 func (s *ServerSession) handshake4(cmd byte, reader io.Reader, writer io.Writer) (*protocol.RequestHeader, error) {
@@ -85,7 +88,13 @@ func (s *ServerSession) handshake4(cmd byte, reader io.Reader, writer io.Writer)
 			Port:    port,
 			Version: socks4Version,
 		}
-		if err := writeSocks4Response(writer, socks4RequestGranted, net.AnyIP, net.Port(0)); err != nil {
+		if err := s.setupLastReply(func(ok bool) error {
+			if ok {
+				return writeSocks4Response(writer, socks4RequestGranted, net.AnyIP, net.Port(0))
+			} else {
+				return writeSocks4Response(writer, socks4RequestRejected, net.AnyIP, net.Port(0))
+			}
+		}); err != nil {
 			return nil, err
 		}
 		return request, nil
@@ -203,13 +212,37 @@ func (s *ServerSession) handshake5(nMethod byte, reader io.Reader, writer io.Wri
 			responseAddress = s.address
 		}
 	}
-	if err := writeSocks5Response(writer, statusSuccess, responseAddress, responsePort); err != nil {
+	if err := s.setupLastReply(func(ok bool) error {
+		if ok {
+			return writeSocks5Response(writer, statusSuccess, responseAddress, responsePort)
+		} else {
+			return writeSocks5Response(writer, statusConnRefused, net.AnyIP, net.Port(0))
+		}
+	}); err != nil {
 		return nil, err
 	}
 
 	return request, nil
 }
 
+// Sets the callback and calls or postpones it based on the boolean field
+func (s *ServerSession) setupLastReply(callback func(bool) error) error {
+	noOpCallback := func(bool) error {
+		return nil
+	}
+
+	// set the field even if we call it now because it will be called again
+	s.flushLastReply = func(ok bool) error {
+		s.flushLastReply = noOpCallback
+		return callback(ok)
+	}
+
+	if s.deferLastReply {
+		return nil
+	}
+	return s.flushLastReply(true)
+}
+
 // Handshake performs a Socks4/4a/5 handshake.
 func (s *ServerSession) Handshake(reader io.Reader, writer io.Writer) (*protocol.RequestHeader, error) {
 	buffer := buf.StackNew()

proxy/socks/server.go

@@ -19,6 +19,7 @@ import (
 	"github.com/v2fly/v2ray-core/v5/features"
 	"github.com/v2fly/v2ray-core/v5/features/policy"
 	"github.com/v2fly/v2ray-core/v5/features/routing"
+	"github.com/v2fly/v2ray-core/v5/transport"
 	"github.com/v2fly/v2ray-core/v5/transport/internet"
 	"github.com/v2fly/v2ray-core/v5/transport/internet/udp"
 )
@@ -90,10 +91,11 @@ func (s *Server) processTCP(ctx context.Context, conn internet.Connection, dispa
 	}
 
 	svrSession := &ServerSession{
-		config:        s.config,
-		address:       inbound.Gateway.Address,
-		port:          inbound.Gateway.Port,
-		clientAddress: inbound.Source.Address,
+		config:         s.config,
+		address:        inbound.Gateway.Address,
+		port:           inbound.Gateway.Port,
+		clientAddress:  inbound.Source.Address,
+		deferLastReply: true,
 	}
 
 	reader := &buf.BufferedReader{Reader: buf.NewReader(conn)}
@@ -129,16 +131,46 @@ func (s *Server) processTCP(ctx context.Context, conn internet.Connection, dispa
 			})
 		}
 
+		dispatcher = &handshakeFinalizingDispatcher{Feature: dispatcher, delegate: dispatcher, serverSession: svrSession}
 		return s.transport(ctx, reader, conn, dest, dispatcher)
 	}
 
+	svrSession.flushLastReply(true)
 	if request.Command == protocol.RequestCommandUDP {
 		return s.handleUDP(conn)
 	}
 
 	return nil
 }
 
+// Wrapper to send final SOCKS reply only after Dispatch() returns
+type handshakeFinalizingDispatcher struct {
+	features.Feature // do not inherit Dispatch() in case its signature changes
+	delegate         routing.Dispatcher
+	serverSession    *ServerSession
+}
+
+func (d *handshakeFinalizingDispatcher) Dispatch(ctx context.Context, dest net.Destination) (*transport.Link, error) {
+	link, err := d.delegate.Dispatch(ctx, dest)
+	if err == nil {
+		closeInDefer := true
+		defer func() {
+			if closeInDefer {
+				common.Interrupt(link.Reader)
+				common.Interrupt(link.Writer)
+			}
+		}()
+		err = d.serverSession.flushLastReply(true)
+		if err != nil {
+			return nil, err
+		}
+		closeInDefer = false
+	} else {
+		d.serverSession.flushLastReply(false)
+	}
+	return link, err
+}
+
 func (*Server) handleUDP(c io.Reader) error {
 	// The TCP connection closes after this method returns. We need to wait until
 	// the client closes it.