🦙 taint: add unsafe fallback

Nothing is free. This feature came with the price of 2 more allocations of up to 16+11+16 bytes when accepting a new Shadowsocks 2022 connection.

When unsafe fallback is enabled, the server is tainted and a warning is printed at startup.

service/server.go

@@ -42,6 +42,9 @@ type ServerConfig struct {
 	RejectPolicy  string   `json:"rejectPolicy"`
 	cipherConfig  *ss2022.CipherConfig
 	uPSKMap       map[[ss2022.IdentityHeaderLength]byte]*ss2022.CipherConfig
+
+	// Taint
+	UnsafeFallbackAddress *conn.Addr `json:"unsafeFallbackAddress"`
 }
 
 // TCPRelay creates a TCP relay service from the ServerConfig.
@@ -89,9 +92,16 @@ func (sc *ServerConfig) TCPRelay(router *router.Router, logger *zap.Logger) (*TC
 		return nil, err
 	}
 
+	if sc.UnsafeFallbackAddress != nil {
+		logger.Warn("Unsafe fallback taints the server",
+			zap.String("server", sc.Name),
+			zap.Stringer("fallbackAddress", sc.UnsafeFallbackAddress),
+		)
+	}
+
 	waitForInitialPayload := !server.NativeInitialPayload() && !sc.DisableInitialPayloadWait
 
-	return NewTCPRelay(sc.Name, sc.Listen, sc.ListenerFwmark, sc.ListenerTFO, waitForInitialPayload, server, connCloser, router, logger), nil
+	return NewTCPRelay(sc.Name, sc.Listen, sc.ListenerFwmark, sc.ListenerTFO, waitForInitialPayload, server, connCloser, sc.UnsafeFallbackAddress, router, logger), nil
 }
 
 // UDPRelay creates a UDP relay service from the ServerConfig.

service/tcp.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/database64128/shadowsocks-go/conn"
+	"github.com/database64128/shadowsocks-go/direct"
 	"github.com/database64128/shadowsocks-go/router"
 	"github.com/database64128/shadowsocks-go/zerocopy"
 	"github.com/database64128/tfo-go/v2"
@@ -35,19 +36,21 @@ type TCPRelay struct {
 	waitForInitialPayload bool
 	server                zerocopy.TCPServer
 	connCloser            zerocopy.TCPConnCloser
+	fallbackAddress       *conn.Addr
 	router                *router.Router
 	logger                *zap.Logger
 	listener              *net.TCPListener
 }
 
-func NewTCPRelay(serverName, listenAddress string, listenerFwmark int, listenerTFO, waitForInitialPayload bool, server zerocopy.TCPServer, connCloser zerocopy.TCPConnCloser, router *router.Router, logger *zap.Logger) *TCPRelay {
+func NewTCPRelay(serverName, listenAddress string, listenerFwmark int, listenerTFO, waitForInitialPayload bool, server zerocopy.TCPServer, connCloser zerocopy.TCPConnCloser, fallbackAddress *conn.Addr, router *router.Router, logger *zap.Logger) *TCPRelay {
 	return &TCPRelay{
 		serverName:            serverName,
 		listenAddress:         listenAddress,
 		listenConfig:          conn.NewListenConfig(listenerTFO, listenerFwmark),
 		waitForInitialPayload: waitForInitialPayload,
 		server:                server,
 		connCloser:            connCloser,
+		fallbackAddress:       fallbackAddress,
 		router:                router,
 		logger:                logger,
 	}
@@ -124,8 +127,13 @@ func (s *TCPRelay) handleConn(clientConn *net.TCPConn) {
 			zap.Error(err),
 		)
 
-		s.connCloser(clientConn, s.serverName, s.listenAddress, clientAddress, s.logger)
-		return
+		if s.fallbackAddress == nil || len(payload) == 0 {
+			s.connCloser(clientConn, s.serverName, s.listenAddress, clientAddress, s.logger)
+			return
+		}
+
+		clientRW = direct.NewDirectStreamReadWriter(clientConn)
+		targetAddr = *s.fallbackAddress
 	}
 
 	// Route.

ss2022/stream.go

@@ -46,6 +46,7 @@ func NewShadowStreamServerReadWriter(rw zerocopy.DirectReadWriteCloser, cipherCo
 		return
 	}
 	if n < bufferLen {
+		payload = b[:n]
 		err = &HeaderError[int]{ErrFirstRead, bufferLen, n}
 		return
 	}
@@ -54,13 +55,14 @@ func NewShadowStreamServerReadWriter(rw zerocopy.DirectReadWriteCloser, cipherCo
 	ciphertext := b[saltLen+identityHeaderLen:]
 
 	if identityHeaderLen != 0 {
+		var uPSKHash [IdentityHeaderLength]byte
 		identityHeader := b[saltLen : saltLen+identityHeaderLen]
 		identityHeaderCipher := cipherConfig.NewTCPIdentityHeaderServerCipher(salt)
-		identityHeaderCipher.Decrypt(identityHeader, identityHeader)
+		identityHeaderCipher.Decrypt(uPSKHash[:], identityHeader)
 
-		uPSKHash := *(*[IdentityHeaderLength]byte)(identityHeader)
 		userCipherConfig, ok := uPSKMap[uPSKHash]
 		if !ok {
+			payload = b[:n]
 			err = ErrIdentityHeaderUserPSKNotFound
 			return
 		}
@@ -71,8 +73,9 @@ func NewShadowStreamServerReadWriter(rw zerocopy.DirectReadWriteCloser, cipherCo
 	shadowStreamCipher := cipherConfig.NewShadowStreamCipher(salt)
 
 	// AEAD open.
-	plaintext, err := shadowStreamCipher.DecryptInPlace(ciphertext)
+	plaintext, err := shadowStreamCipher.DecryptTo(nil, ciphertext)
 	if err != nil {
+		payload = b[:n]
 		return
 	}
 

zerocopy/tcp.go

@@ -40,11 +40,13 @@ type TCPClient interface {
 type TCPServer interface {
 	InitialPayloader
 
-	// Accept takes a newly-accepted TCP connection and wraps it into a
-	// protocol stream server.
+	// Accept takes a newly-accepted TCP connection and wraps it into a protocol stream server.
 	//
 	// If the returned error is ErrAcceptDoneNoRelay, the connection has been handled by this method.
 	// Two-way relay is not needed.
+	//
+	// If accept fails, the returned payload must be either nil/empty or the data that has been read
+	// from the connection.
 	Accept(tc *net.TCPConn) (rw ReadWriter, targetAddr conn.Addr, payload []byte, err error)
 
 	// DefaultTCPConnCloser returns the default function to handle the closing