Refactor

The relay session context was derived from the inbound connection of the
source address that created the session. After a client changes its source
address (NAT rebinding), the UDP worker of the old address is cleaned up after
~2 minutes of inactivity and cancels that context, which tears down the
dispatched outbound flow; the next packet then creates a brand-new session and
outbound flow (new local port on the server, reset replay window).

Sessions are routed by client session ID (SIP022 3.2.4), so detach the session
context from the originating connection's cancellation and let the session's
own inactivity timer govern its lifetime.

proxy/shadowsocks_2022/inbound.go

@@ -53,21 +53,23 @@ func NewServer(ctx context.Context, config *ServerConfig) (*Inbound, error) {
 		return nil, err
 	}
 
-	udpCodec, err := NewUDPServerCodec(method, psk, 500*time.Second)
+	user := &protocol.MemoryUser{
+		Email: config.Email,
+		Level: uint32(config.Level),
+	}
+
+	udpCodec, err := NewUDPServerCodec(method, psk, user, 500*time.Second)
 	if err != nil {
 		return nil, err
 	}
 
 	v := core.MustFromContext(ctx)
 	return &Inbound{
-		networks:   networks,
-		method:     method,
-		psk:        psk,
-		saltFilter: antireplay.NewMapFilter[[32]byte](60),
-		user: &protocol.MemoryUser{
-			Email: config.Email,
-			Level: uint32(config.Level),
-		},
+		networks:      networks,
+		method:        method,
+		psk:           psk,
+		saltFilter:    antireplay.NewMapFilter[[32]byte](60),
+		user:          user,
 		udpCodec:      udpCodec,
 		policyManager: v.GetFeature(policy.ManagerType()).(policy.Manager),
 	}, nil
@@ -155,20 +157,12 @@ func (i *Inbound) processUDP(ctx context.Context, conn stat.Connection, dispatch
 		}
 
 		for _, b := range mb {
-			decoded, err := i.udpCodec.DecodePacket(b.Bytes())
+			sessionItem, decoded, err := i.udpCodec.DecodePacket(b.Bytes())
 			b.Release()
 			if err != nil || decoded.HeaderType != HeaderTypeClient {
 				continue
 			}
 
-			sessionItem := i.udpCodec.GetSession(decoded.SessionID)
-			if sessionItem.User == nil {
-				sessionItem.Lock()
-				if sessionItem.User == nil {
-					sessionItem.User = i.user
-				}
-				sessionItem.Unlock()
-			}
 			link := sessionItem.link.Load()
 			if link != nil {
 				sessionItem.UpdateConn(conn)

proxy/shadowsocks_2022/inbound_multi.go

@@ -2,8 +2,6 @@ package shadowsocks_2022
 
 import (
 	"context"
-	"crypto/cipher"
-	"encoding/binary"
 	"strconv"
 	"strings"
 	"sync"
@@ -41,8 +39,7 @@ type MultiUserInbound struct {
 	usersByEmail    *utils.TypedSyncMap[string, *protocol.MemoryUser]
 	userCount       atomic.Int64
 	saltFilter      *antireplay.ReplayFilter[[32]byte]
-	udpSessions     *UDPSessionManager
-	udpMasterCipher cipher.Block
+	udpCodec        *UDPMultiServerCodec
 	policyManager   policy.Manager
 }
 
@@ -68,22 +65,21 @@ func NewMultiServer(ctx context.Context, config *MultiUserServerConfig) (*MultiU
 		return nil, err
 	}
 
-	masterBlock, err := method.NewBlock(masterPSK)
+	udpCodec, err := NewUDPMultiServerCodec(method, masterPSK, 500*time.Second)
 	if err != nil {
 		return nil, err
 	}
 
 	v := core.MustFromContext(ctx)
 	i := &MultiUserInbound{
-		networks:        networks,
-		method:          method,
-		masterPSK:       masterPSK,
-		usersByHash:     utils.NewTypedSyncMap[[AESBlockSize]byte, *protocol.MemoryUser](),
-		usersByEmail:    utils.NewTypedSyncMap[string, *protocol.MemoryUser](),
-		saltFilter:      antireplay.NewMapFilter[[32]byte](60),
-		udpSessions:     NewUDPSessionManager(500 * time.Second),
-		udpMasterCipher: masterBlock,
-		policyManager:   v.GetFeature(policy.ManagerType()).(policy.Manager),
+		networks:      networks,
+		method:        method,
+		masterPSK:     masterPSK,
+		usersByHash:   utils.NewTypedSyncMap[[AESBlockSize]byte, *protocol.MemoryUser](),
+		usersByEmail:  utils.NewTypedSyncMap[string, *protocol.MemoryUser](),
+		saltFilter:    antireplay.NewMapFilter[[32]byte](60),
+		udpCodec:      udpCodec,
+		policyManager: v.GetFeature(policy.ManagerType()).(policy.Manager),
 	}
 
 	for idx, user := range config.Users {
@@ -281,68 +277,26 @@ func (i *MultiUserInbound) processUDP(ctx context.Context, conn stat.Connection,
 		}
 
 		for _, b := range mb {
-			// In multi-user UDP:
-			// Packet header is 16 bytes: Encrypted(SessionID + PacketID)
-			// Followed by 16 bytes EIH
-			packetBytes := b.Bytes()
-			if len(packetBytes) < 32+1+8+2 {
-				b.Release()
-				continue
-			}
-
-			var rawHeader [16]byte
-			i.udpMasterCipher.Decrypt(rawHeader[:], packetBytes[:16])
-
-			sessionID := binary.BigEndian.Uint64(rawHeader[:8])
-			packetID := binary.BigEndian.Uint64(rawHeader[8:16])
-
-			sessionItem := i.udpSessions.GetOrCreate(sessionID)
-
-			if !sessionItem.CheckPacketID(packetID) {
-				b.Release()
-				continue
-			}
-
-			var userPSK []byte
-			var currentUser *protocol.MemoryUser
-			sessionItem.Lock()
-			currentUser = sessionItem.User
-			userPSK = sessionItem.UserPSK
-			sessionItem.Unlock()
-
-			if currentUser == nil {
-				// Decrypt EIH
-				decryptedHash := DecryptUDPEIH(i.udpMasterCipher, rawHeader[:], packetBytes[16:32])
-
-				user, ok := i.usersByHash.Load(decryptedHash)
+			sessionItem, decoded, user, err := i.udpCodec.DecodePacket(b.Bytes(), func(hash [AESBlockSize]byte) (*protocol.MemoryUser, []byte, bool) {
+				u, ok := i.usersByHash.Load(hash)
 				if !ok {
-					b.Release()
-					continue
+					return nil, nil, false
 				}
-				currentUser = user
-				userPSK = user.Account.(*MemoryAccount).Key
-			}
-
-			decoded, err := sessionItem.DecryptAESPayload(i.method, userPSK, sessionID, packetID, rawHeader[:], packetBytes[32:])
+				return u, u.Account.(*MemoryAccount).Key, true
+			})
 			b.Release()
-			if err != nil {
+			if err != nil || decoded.HeaderType != HeaderTypeClient {
 				continue
 			}
 
-			sessionItem.Lock()
-			if sessionItem.User == nil {
-				sessionItem.User = currentUser
-				sessionItem.UserPSK = userPSK
-			}
-			sessionItem.Unlock()
-
 			link := sessionItem.link.Load()
 			if link != nil {
 				sessionItem.UpdateConn(conn)
 			} else {
+				userPSK := user.Account.(*MemoryAccount).Key
 				var err error
 				link, err = sessionItem.EnsureLink(ctx, conn, decoded.Destination, dispatcher, i.policyManager, func(replyDest net.Destination, payload []byte) ([]byte, error) {
-					return i.encodeServerUDPPacket(sessionID, userPSK, replyDest, payload)
+					return i.udpCodec.EncodeServerPacket(decoded.SessionID, userPSK, replyDest, payload)
 				})
 				if err != nil {
 					continue
@@ -356,7 +310,3 @@ func (i *MultiUserInbound) processUDP(ctx context.Context, conn stat.Connection,
 		}
 	}
 }
-
-func (i *MultiUserInbound) encodeServerUDPPacket(clientSessionID uint64, userPSK []byte, dest net.Destination, payload []byte) ([]byte, error) {
-	return i.udpSessions.EncodeServerPacket(i.method, userPSK, clientSessionID, dest, payload)
-}

proxy/shadowsocks_2022/inbound_relay.go

@@ -243,16 +243,14 @@ func (i *RelayInbound) processUDP(ctx context.Context, conn stat.Connection, dis
 			dest.Network = net.Network_UDP
 
 			sessionItem := i.udpSessions.GetOrCreate(sessionID)
+			sessionItem.Lock()
 			if sessionItem.User == nil {
-				sessionItem.Lock()
-				if sessionItem.User == nil {
-					sessionItem.User = &protocol.MemoryUser{
-						Email: targetDest.email,
-						Level: targetDest.level,
-					}
+				sessionItem.User = &protocol.MemoryUser{
+					Email: targetDest.email,
+					Level: targetDest.level,
 				}
-				sessionItem.Unlock()
 			}
+			sessionItem.Unlock()
 			link := sessionItem.link.Load()
 			if link != nil {
 				sessionItem.UpdateConn(conn)

proxy/shadowsocks_2022/outbound.go

@@ -30,7 +30,7 @@ type Outbound struct {
 	method        *CipherMethod
 	pskList       [][]byte
 	finalPSK      []byte
-	udpCodec      *UDPPacketCodec
+	udpCodec      *UDPClientCodec
 	policyManager policy.Manager
 }
 
@@ -50,9 +50,9 @@ func NewClient(ctx context.Context, config *ClientConfig) (*Outbound, error) {
 	}
 
 	finalPSK := pskList[len(pskList)-1]
-	udpCodec, err := NewUDPPacketCodec(method, pskList)
+	udpCodec, err := NewUDPClientCodec(method, pskList)
 	if err != nil {
-		return nil, errors.New("failed to create udp packet codec").Base(err)
+		return nil, errors.New("failed to create udp client codec").Base(err)
 	}
 
 	v := core.MustFromContext(ctx)

proxy/shadowsocks_2022/packet.go

@@ -7,33 +7,78 @@ import (
 	"io"
 	"math"
 	mrand "math/rand/v2"
+	"sync"
 	"sync/atomic"
 	"time"
 
 	"github.com/xtls/xray-core/common/buf"
 	"github.com/xtls/xray-core/common/errors"
 	"github.com/xtls/xray-core/common/net"
+	"github.com/xtls/xray-core/common/protocol"
 )
 
-type UDPCodec struct {
+type UDPClientCodec struct {
 	method       *CipherMethod
 	pskList      [][]byte
 	psk          []byte
 	blockCipher  cipher.Block
 	blockCiphers []cipher.Block
 	chachaCipher cipher.AEAD
+}
+
+func NewUDPClientCodec(method *CipherMethod, pskList [][]byte) (*UDPClientCodec, error) {
+	if len(pskList) == 0 {
+		return nil, errors.New("empty psk list")
+	}
+	if method.IsChaCha && len(pskList) > 1 {
+		return nil, errors.New("multi-key is not supported for chacha20-poly1305")
+	}
+	finalPSK := pskList[len(pskList)-1]
+	c := &UDPClientCodec{
+		method:  method,
+		pskList: pskList,
+		psk:     finalPSK,
+	}
+	var err error
+	if method.IsChaCha {
+		c.chachaCipher, err = method.NewUDPCipher(finalPSK)
+	} else {
+		c.blockCipher, err = method.NewBlock(finalPSK)
+	}
+	if err != nil {
+		return nil, err
+	}
+	if len(pskList) > 1 {
+		c.blockCiphers = make([]cipher.Block, len(pskList))
+		for i, psk := range pskList {
+			c.blockCiphers[i], err = method.NewBlock(psk)
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
+	return c, nil
+}
+
+func NewUDPPacketCodec(method *CipherMethod, pskList [][]byte) (*UDPClientCodec, error) {
+	return NewUDPClientCodec(method, pskList)
+}
+
+type UDPServerCodec struct {
+	method       *CipherMethod
+	psk          []byte
+	user         *protocol.MemoryUser
+	blockCipher  cipher.Block
+	chachaCipher cipher.AEAD
 	sessions     *UDPSessionManager
 }
 
-type (
-	UDPPacketCodec = UDPCodec
-	UDPServerCodec = UDPCodec
-)
-
-func newUDPCodec(method *CipherMethod, psk []byte) (*UDPCodec, error) {
-	c := &UDPCodec{
-		method: method,
-		psk:    psk,
+func NewUDPServerCodec(method *CipherMethod, psk []byte, user *protocol.MemoryUser, sessionTimeout time.Duration) (*UDPServerCodec, error) {
+	c := &UDPServerCodec{
+		method:   method,
+		psk:      psk,
+		user:     user,
+		sessions: NewUDPSessionManager(sessionTimeout),
 	}
 	var err error
 	if method.IsChaCha {
@@ -47,46 +92,129 @@ func newUDPCodec(method *CipherMethod, psk []byte) (*UDPCodec, error) {
 	return c, nil
 }
 
-func NewUDPPacketCodec(method *CipherMethod, pskList [][]byte) (*UDPCodec, error) {
-	if method.IsChaCha && len(pskList) > 1 {
-		return nil, errors.New("multi-key is not supported for chacha20-poly1305")
-	}
-	finalPSK := pskList[len(pskList)-1]
-	c, err := newUDPCodec(method, finalPSK)
-	if err != nil {
-		return nil, err
-	}
-	c.pskList = pskList
-	if len(pskList) > 1 {
-		c.blockCiphers = make([]cipher.Block, len(pskList))
-		for i, psk := range pskList {
-			c.blockCiphers[i], err = method.NewBlock(psk)
-			if err != nil {
-				return nil, err
-			}
-		}
-	}
-	return c, nil
-}
-
-func NewUDPServerCodec(method *CipherMethod, psk []byte, sessionTimeout time.Duration) (*UDPCodec, error) {
-	c, err := newUDPCodec(method, psk)
-	if err != nil {
-		return nil, err
-	}
-	c.sessions = NewUDPSessionManager(sessionTimeout)
-	return c, nil
-}
-
-func (c *UDPCodec) Sessions() *UDPSessionManager {
+func (c *UDPServerCodec) Sessions() *UDPSessionManager {
 	return c.sessions
 }
 
-func (c *UDPCodec) GetSession(sessionID uint64) *ServerUDPSession {
+func (c *UDPServerCodec) GetSession(sessionID uint64) *ServerUDPSession {
 	if c.sessions == nil {
 		return nil
 	}
-	return c.sessions.GetOrCreate(sessionID)
+	return c.sessions.Get(sessionID)
+}
+
+type UDPMultiServerCodec struct {
+	method      *CipherMethod
+	masterPSK   []byte
+	masterBlock cipher.Block
+	sessions    *UDPSessionManager
+}
+
+func NewUDPMultiServerCodec(method *CipherMethod, masterPSK []byte, sessionTimeout time.Duration) (*UDPMultiServerCodec, error) {
+	if method.IsChaCha {
+		return nil, errors.New("multi-user is not supported for chacha20-poly1305")
+	}
+	masterBlock, err := method.NewBlock(masterPSK)
+	if err != nil {
+		return nil, err
+	}
+	return &UDPMultiServerCodec{
+		method:      method,
+		masterPSK:   masterPSK,
+		masterBlock: masterBlock,
+		sessions:    NewUDPSessionManager(sessionTimeout),
+	}, nil
+}
+
+func (c *UDPMultiServerCodec) Sessions() *UDPSessionManager {
+	return c.sessions
+}
+
+func (c *UDPMultiServerCodec) GetSession(sessionID uint64) *ServerUDPSession {
+	if c.sessions == nil {
+		return nil
+	}
+	return c.sessions.Get(sessionID)
+}
+
+type UserLookupFunc func(hash [AESBlockSize]byte) (*protocol.MemoryUser, []byte, bool)
+
+func (c *UDPMultiServerCodec) DecodePacket(data []byte, userLookup UserLookupFunc) (*ServerUDPSession, DecodedUDPPacket, *protocol.MemoryUser, error) {
+	if len(data) < 2*AESBlockSize+AEADTagSize+RequestHeaderFixedChunkLength {
+		return nil, DecodedUDPPacket{}, nil, ErrPacketTooShort
+	}
+
+	var rawHeader [16]byte
+	c.masterBlock.Decrypt(rawHeader[:], data[:16])
+
+	sessionID := binary.BigEndian.Uint64(rawHeader[:8])
+	packetID := binary.BigEndian.Uint64(rawHeader[8:16])
+
+	sessionItem := c.sessions.Get(sessionID)
+	var currentUser *protocol.MemoryUser
+	var userPSK []byte
+	var bodyAead cipher.AEAD
+
+	if sessionItem != nil {
+		sessionItem.Lock()
+		currentUser = sessionItem.User
+		userPSK = sessionItem.UserPSK
+		bodyAead = sessionItem.clientBodyCipher
+		sessionItem.Unlock()
+
+		if currentUser != nil {
+			if !sessionItem.CheckPacketID(packetID) {
+				return nil, DecodedUDPPacket{}, nil, ErrPacketIdNotUnique
+			}
+		}
+	}
+
+	if currentUser == nil {
+		decryptedHash := DecryptUDPEIH(c.masterBlock, rawHeader[:], data[16:32])
+		user, psk, ok := userLookup(decryptedHash)
+		if !ok {
+			return nil, DecodedUDPPacket{}, nil, errors.New("user not found")
+		}
+		currentUser = user
+		userPSK = psk
+	}
+
+	if bodyAead == nil {
+		bodyKey := DeriveSessionSubKey(userPSK, rawHeader[:8], c.method.KeySaltLength)
+		var err error
+		bodyAead, err = c.method.NewAEAD(bodyKey)
+		if err != nil {
+			return nil, DecodedUDPPacket{}, nil, err
+		}
+	}
+
+	bodyNonce := rawHeader[4:16]
+	bodyCipher := data[32:]
+	bodyPlain, err := bodyAead.Open(nil, bodyNonce, bodyCipher, nil)
+	if err != nil {
+		return nil, DecodedUDPPacket{}, nil, errors.New("failed to decrypt aes udp body").Base(err)
+	}
+
+	decoded, err := parsePlainUDPPacket(sessionID, packetID, bodyPlain)
+	if err != nil {
+		return nil, DecodedUDPPacket{}, nil, err
+	}
+	if decoded.HeaderType != HeaderTypeClient {
+		return nil, DecodedUDPPacket{}, nil, ErrBadHeaderType
+	}
+
+	if sessionItem == nil {
+		sessionItem = c.sessions.GetOrCreate(sessionID)
+	}
+	if !sessionItem.Commit(packetID, bodyAead, currentUser, userPSK) {
+		return nil, DecodedUDPPacket{}, nil, ErrPacketIdNotUnique
+	}
+
+	return sessionItem, decoded, currentUser, nil
+}
+
+func (c *UDPMultiServerCodec) EncodeServerPacket(clientSessionID uint64, userPSK []byte, dest net.Destination, payload []byte) ([]byte, error) {
+	return c.sessions.EncodeServerPacket(c.method, userPSK, clientSessionID, dest, payload)
 }
 
 type DecodedUDPPacket struct {
@@ -193,44 +321,42 @@ func parsePlainUDPPacket(sessionID, packetID uint64, bodyPlain []byte) (DecodedU
 	}, nil
 }
 
-func (c *UDPCodec) DecodePacket(data []byte) (DecodedUDPPacket, error) {
+func (c *UDPServerCodec) DecodePacket(data []byte) (*ServerUDPSession, DecodedUDPPacket, error) {
 	if len(data) < PacketMinimalHeaderSize {
-		return DecodedUDPPacket{}, ErrPacketTooShort
+		return nil, DecodedUDPPacket{}, ErrPacketTooShort
 	}
 
 	if c.method.IsChaCha {
 		if len(data) < PacketNonceSize+AEADTagSize {
-			return DecodedUDPPacket{}, ErrPacketTooShort
+			return nil, DecodedUDPPacket{}, ErrPacketTooShort
 		}
 		nonce := data[:PacketNonceSize]
 		ciphertext := data[PacketNonceSize:]
 		plain, err := c.chachaCipher.Open(nil, nonce, ciphertext, nil)
 		if err != nil {
-			return DecodedUDPPacket{}, errors.New("failed to decrypt chacha udp packet").Base(err)
+			return nil, DecodedUDPPacket{}, errors.New("failed to decrypt chacha udp packet").Base(err)
 		}
 		if len(plain) < 16+1+8+2 {
-			return DecodedUDPPacket{}, ErrPacketTooShort
+			return nil, DecodedUDPPacket{}, ErrPacketTooShort
 		}
 
 		sessionID := binary.BigEndian.Uint64(plain[:8])
 		packetID := binary.BigEndian.Uint64(plain[8:16])
 
-		sessionItem := c.sessions.GetOrCreate(sessionID)
-		if !sessionItem.CheckPacketID(packetID) {
-			return DecodedUDPPacket{}, ErrPacketIdNotUnique
-		}
-
 		decoded, err := parsePlainUDPPacket(sessionID, packetID, plain[16:])
 		if err != nil {
-			return DecodedUDPPacket{}, err
+			return nil, DecodedUDPPacket{}, err
 		}
 
 		if decoded.HeaderType != HeaderTypeClient {
-			return DecodedUDPPacket{}, ErrBadHeaderType
+			return nil, DecodedUDPPacket{}, ErrBadHeaderType
 		}
 
-		sessionItem.AddPacketID(packetID)
-		return decoded, nil
+		sessionItem := c.sessions.GetOrCreate(sessionID)
+		if !sessionItem.Commit(packetID, nil, c.user, c.psk) {
+			return nil, DecodedUDPPacket{}, ErrPacketIdNotUnique
+		}
+		return sessionItem, decoded, nil
 	}
 
 	// AES mode
@@ -239,49 +365,48 @@ func (c *UDPCodec) DecodePacket(data []byte) (DecodedUDPPacket, error) {
 	sessionID := binary.BigEndian.Uint64(rawHeader[:8])
 	packetID := binary.BigEndian.Uint64(rawHeader[8:16])
 
-	sessionItem := c.sessions.GetOrCreate(sessionID)
-	if !sessionItem.CheckPacketID(packetID) {
-		return DecodedUDPPacket{}, ErrPacketIdNotUnique
+	sessionItem := c.sessions.Get(sessionID)
+	var bodyAead cipher.AEAD
+	if sessionItem != nil {
+		if !sessionItem.CheckPacketID(packetID) {
+			return nil, DecodedUDPPacket{}, ErrPacketIdNotUnique
+		}
+		sessionItem.Lock()
+		bodyAead = sessionItem.clientBodyCipher
+		sessionItem.Unlock()
 	}
 
-	return sessionItem.DecryptAESPayload(c.method, c.psk, sessionID, packetID, rawHeader[:], data[16:])
-}
-
-func (s *ServerUDPSession) DecryptAESPayload(method *CipherMethod, psk []byte, sessionID, packetID uint64, rawHeader, bodyCipher []byte) (DecodedUDPPacket, error) {
-	bodyAead := s.clientBodyCipher
-	isNewCipher := false
 	if bodyAead == nil {
-		bodyKey := DeriveSessionSubKey(psk, rawHeader[:8], method.KeySaltLength)
+		bodyKey := DeriveSessionSubKey(c.psk, rawHeader[:8], c.method.KeySaltLength)
 		var err error
-		bodyAead, err = method.NewAEAD(bodyKey)
+		bodyAead, err = c.method.NewAEAD(bodyKey)
 		if err != nil {
-			return DecodedUDPPacket{}, err
+			return nil, DecodedUDPPacket{}, err
 		}
-		isNewCipher = true
 	}
 
 	bodyNonce := rawHeader[4:16]
+	bodyCipher := data[16:]
 	bodyPlain, err := bodyAead.Open(nil, bodyNonce, bodyCipher, nil)
 	if err != nil {
-		return DecodedUDPPacket{}, errors.New("failed to decrypt aes udp body").Base(err)
+		return nil, DecodedUDPPacket{}, errors.New("failed to decrypt aes udp body").Base(err)
 	}
 
 	decoded, err := parsePlainUDPPacket(sessionID, packetID, bodyPlain)
 	if err != nil {
-		return DecodedUDPPacket{}, err
+		return nil, DecodedUDPPacket{}, err
 	}
-
 	if decoded.HeaderType != HeaderTypeClient {
-		return DecodedUDPPacket{}, ErrBadHeaderType
+		return nil, DecodedUDPPacket{}, ErrBadHeaderType
 	}
 
-	s.AddPacketID(packetID)
-
-	if isNewCipher {
-		s.clientBodyCipher = bodyAead
+	if sessionItem == nil {
+		sessionItem = c.sessions.GetOrCreate(sessionID)
 	}
-
-	return decoded, nil
+	if !sessionItem.Commit(packetID, bodyAead, c.user, c.psk) {
+		return nil, DecodedUDPPacket{}, ErrPacketIdNotUnique
+	}
+	return sessionItem, decoded, nil
 }
 
 func (s *ServerUDPSession) EnsureServerState(method *CipherMethod, psk []byte) error {
@@ -303,6 +428,9 @@ func (s *ServerUDPSession) EnsureServerState(method *CipherMethod, psk []byte) e
 	if method.IsChaCha {
 		var err error
 		s.serverChaCha, err = method.NewUDPCipher(psk)
+		if err != nil {
+			s.ServerSessionID = 0
+		}
 		return err
 	}
 
@@ -387,11 +515,12 @@ func (s *ServerUDPSession) EncodeServerPacket(method *CipherMethod, clientSessio
 	return res, nil
 }
 
-func (c *UDPCodec) EncodeServerPacket(clientSessionID uint64, dest net.Destination, payload []byte) ([]byte, error) {
+func (c *UDPServerCodec) EncodeServerPacket(clientSessionID uint64, dest net.Destination, payload []byte) ([]byte, error) {
 	return c.sessions.EncodeServerPacket(c.method, c.psk, clientSessionID, dest, payload)
 }
 
 type serverSessionState struct {
+	sync.Mutex
 	sessionID uint64
 	window    *SlidingWindow
 	cipher    cipher.AEAD
@@ -399,21 +528,20 @@ type serverSessionState struct {
 }
 
 func (st *serverSessionState) check(packetID uint64) bool {
-	if st.window == nil {
-		st.window = new(SlidingWindow)
-	}
+	st.Lock()
+	defer st.Unlock()
 	return st.window.Check(packetID)
 }
 
-func (st *serverSessionState) add(packetID uint64) {
-	if st.window == nil {
-		st.window = new(SlidingWindow)
-	}
-	st.window.Add(packetID)
+func (st *serverSessionState) commit(packetID uint64) bool {
+	st.Lock()
+	defer st.Unlock()
+	return st.window.CheckAndAdd(packetID)
 }
 
 type ClientUDPSession struct {
-	codec            *UDPCodec
+	mu               sync.Mutex
+	codec            *UDPClientCodec
 	clientSessionID  uint64
 	nextPacketID     atomic.Uint64
 	clientBodyCipher cipher.AEAD
@@ -421,7 +549,7 @@ type ClientUDPSession struct {
 	old              atomic.Pointer[serverSessionState]
 }
 
-func (c *UDPCodec) NewClientSession() (*ClientUDPSession, error) {
+func (c *UDPClientCodec) NewClientSession() (*ClientUDPSession, error) {
 	var sessID [8]byte
 	if _, err := io.ReadFull(rand.Reader, sessID[:]); err != nil {
 		return nil, err
@@ -446,7 +574,7 @@ func (c *UDPCodec) NewClientSession() (*ClientUDPSession, error) {
 	}, nil
 }
 
-func (s *ClientUDPSession) getServerSession(sessionID uint64, now int64) (*serverSessionState, error) {
+func (s *ClientUDPSession) getOrTransitionServerSession(sessionID uint64, cipher cipher.AEAD, now int64) (*serverSessionState, error) {
 	cur := s.current.Load()
 	if cur != nil && cur.sessionID == sessionID {
 		return cur, nil
@@ -455,39 +583,42 @@ func (s *ClientUDPSession) getServerSession(sessionID uint64, now int64) (*serve
 	old := s.old.Load()
 	if old != nil && old.sessionID == sessionID {
 		if now-old.lastSeen.Load() > 60 {
-			s.old.CompareAndSwap(old, nil)
 			return nil, errors.New("old server session expired")
 		}
 		return old, nil
 	}
 
-	// New server session:
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	cur = s.current.Load()
+	if cur != nil && cur.sessionID == sessionID {
+		return cur, nil
+	}
+
+	old = s.old.Load()
+	if old != nil && old.sessionID == sessionID {
+		if now-old.lastSeen.Load() > 60 {
+			return nil, errors.New("old server session expired")
+		}
+		return old, nil
+	}
+
 	// Spec §3.2.4: reject newer server sessions when the last packet received from the old session is less than 1 minute old.
 	if old != nil && now-old.lastSeen.Load() < 60 {
 		return nil, errors.New("newer server session rejected: old session is less than 1 minute old")
 	}
 
-	var bodyAead cipher.AEAD
-	if !s.codec.method.IsChaCha {
-		var sessBytes [8]byte
-		binary.BigEndian.PutUint64(sessBytes[:], sessionID)
-		bodyKey := DeriveSessionSubKey(s.codec.psk, sessBytes[:], s.codec.method.KeySaltLength)
-		var err error
-		bodyAead, err = s.codec.method.NewAEAD(bodyKey)
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	newState := &serverSessionState{
 		sessionID: sessionID,
-		cipher:    bodyAead,
+		window:    new(SlidingWindow),
+		cipher:    cipher,
 	}
 	newState.lastSeen.Store(now)
 
 	if cur == nil {
-		s.current.CompareAndSwap(nil, newState)
-		return s.current.Load(), nil
+		s.current.Store(newState)
+		return newState, nil
 	}
 
 	s.old.Store(cur)
@@ -637,15 +768,6 @@ func (s *ClientUDPSession) DecodePacket(data []byte) (DecodedUDPPacket, error) {
 		sessionID := binary.BigEndian.Uint64(plain[:8])
 		packetID := binary.BigEndian.Uint64(plain[8:16])
 
-		now := time.Now().Unix()
-		st, err := s.getServerSession(sessionID, now)
-		if err != nil {
-			return DecodedUDPPacket{}, err
-		}
-		if !st.check(packetID) {
-			return DecodedUDPPacket{}, ErrPacketIdNotUnique
-		}
-
 		decoded, err := parsePlainUDPPacket(sessionID, packetID, plain[16:])
 		if err != nil {
 			return DecodedUDPPacket{}, err
@@ -658,7 +780,16 @@ func (s *ClientUDPSession) DecodePacket(data []byte) (DecodedUDPPacket, error) {
 			return DecodedUDPPacket{}, errors.New("client session ID mismatch")
 		}
 
-		st.add(packetID)
+		now := time.Now().Unix()
+		st, err := s.getOrTransitionServerSession(sessionID, nil, now)
+		if err != nil {
+			return DecodedUDPPacket{}, err
+		}
+
+		// Atomically commit after all validations succeed
+		if !st.commit(packetID) {
+			return DecodedUDPPacket{}, ErrPacketIdNotUnique
+		}
 		st.lastSeen.Store(now)
 
 		return decoded, nil
@@ -671,18 +802,79 @@ func (s *ClientUDPSession) DecodePacket(data []byte) (DecodedUDPPacket, error) {
 	packetID := binary.BigEndian.Uint64(rawHeader[8:16])
 
 	now := time.Now().Unix()
-	st, err := s.getServerSession(sessionID, now)
+
+	// 1. Fast path for known active session
+	cur := s.current.Load()
+	var st *serverSessionState
+	var bodyAead cipher.AEAD
+
+	if cur != nil && cur.sessionID == sessionID {
+		st = cur
+		bodyAead = cur.cipher
+	} else {
+		old := s.old.Load()
+		if old != nil && old.sessionID == sessionID {
+			if now-old.lastSeen.Load() > 60 {
+				return DecodedUDPPacket{}, errors.New("old server session expired")
+			}
+			st = old
+			bodyAead = old.cipher
+		}
+	}
+
+	if st != nil {
+		if !st.check(packetID) {
+			return DecodedUDPPacket{}, ErrPacketIdNotUnique
+		}
+		bodyNonce := rawHeader[4:16]
+		bodyCipher := data[16:]
+		bodyPlain, err := bodyAead.Open(nil, bodyNonce, bodyCipher, nil)
+		if err != nil {
+			return DecodedUDPPacket{}, errors.New("failed to decrypt aes udp body").Base(err)
+		}
+
+		decoded, err := parsePlainUDPPacket(sessionID, packetID, bodyPlain)
+		if err != nil {
+			return DecodedUDPPacket{}, err
+		}
+
+		if decoded.HeaderType != HeaderTypeServer {
+			return DecodedUDPPacket{}, ErrBadHeaderType
+		}
+		if decoded.ClientSessionID != s.clientSessionID {
+			return DecodedUDPPacket{}, errors.New("client session ID mismatch")
+		}
+
+		if !st.commit(packetID) {
+			return DecodedUDPPacket{}, ErrPacketIdNotUnique
+		}
+		st.lastSeen.Store(now)
+
+		return decoded, nil
+	}
+
+	// 2. Unknown session: Trial path without mutating s.current or s.old yet!
+	// Re-check cur in case another worker concurrently transitioned to this sessionID
+	cur = s.current.Load()
+	old := s.old.Load()
+	if (cur == nil || cur.sessionID != sessionID) && (old == nil || old.sessionID != sessionID) {
+		// Spec §3.2.4: reject newer server sessions when the last packet received from the old session is less than 1 minute old.
+		if old != nil && now-old.lastSeen.Load() < 60 {
+			return DecodedUDPPacket{}, errors.New("newer server session rejected: old session is less than 1 minute old")
+		}
+	}
+
+	var sessBytes [8]byte
+	binary.BigEndian.PutUint64(sessBytes[:], sessionID)
+	bodyKey := DeriveSessionSubKey(s.codec.psk, sessBytes[:], s.codec.method.KeySaltLength)
+	trialAead, err := s.codec.method.NewAEAD(bodyKey)
 	if err != nil {
 		return DecodedUDPPacket{}, err
 	}
-	if !st.check(packetID) {
-		return DecodedUDPPacket{}, ErrPacketIdNotUnique
-	}
-	bodyAead := st.cipher
 
 	bodyNonce := rawHeader[4:16]
 	bodyCipher := data[16:]
-	bodyPlain, err := bodyAead.Open(nil, bodyNonce, bodyCipher, nil)
+	bodyPlain, err := trialAead.Open(nil, bodyNonce, bodyCipher, nil)
 	if err != nil {
 		return DecodedUDPPacket{}, errors.New("failed to decrypt aes udp body").Base(err)
 	}
@@ -699,7 +891,15 @@ func (s *ClientUDPSession) DecodePacket(data []byte) (DecodedUDPPacket, error) {
 		return DecodedUDPPacket{}, errors.New("client session ID mismatch")
 	}
 
-	st.add(packetID)
+	// Payload authenticated! Safely commit and transition session under lock
+	st, err = s.getOrTransitionServerSession(sessionID, trialAead, now)
+	if err != nil {
+		return DecodedUDPPacket{}, err
+	}
+
+	if !st.commit(packetID) {
+		return DecodedUDPPacket{}, ErrPacketIdNotUnique
+	}
 	st.lastSeen.Store(now)
 
 	return decoded, nil

proxy/shadowsocks_2022/relay_test.go

@@ -191,7 +191,7 @@ func TestRelayUDPSessionStabilityAndDispatch(t *testing.T) {
 	// Verify downstream destination can decode both packets
 	method, err := GetCipherMethod(MethodAES128GCM)
 	common.Must(err)
-	destCodec, err := NewUDPServerCodec(method, destKey, 300*time.Second)
+	destCodec, err := NewUDPServerCodec(method, destKey, nil, 300*time.Second)
 	common.Must(err)
 
 	mu.Lock()
@@ -202,7 +202,8 @@ func TestRelayUDPSessionStabilityAndDispatch(t *testing.T) {
 		t.Fatalf("expected 2 received packets at destination, got %d", len(pkts))
 	}
 
-	dec1, err := destCodec.DecodePacket(pkts[0])
+	dec1, dec2 := DecodedUDPPacket{}, DecodedUDPPacket{}
+	_, dec1, err = destCodec.DecodePacket(pkts[0])
 	if err != nil {
 		t.Fatalf("dest failed to decode packet 1: %v", err)
 	}
@@ -210,7 +211,7 @@ func TestRelayUDPSessionStabilityAndDispatch(t *testing.T) {
 		t.Fatalf("dec1 mismatch: sess=%x, pktID=%d, payload=%s", dec1.SessionID, dec1.PacketID, string(dec1.Payload))
 	}
 
-	dec2, err := destCodec.DecodePacket(pkts[1])
+	_, dec2, err = destCodec.DecodePacket(pkts[1])
 	if err != nil {
 		t.Fatalf("dest failed to decode packet 2: %v", err)
 	}

proxy/shadowsocks_2022/replay.go

@@ -100,19 +100,27 @@ type ServerUDPSession struct {
 func (s *ServerUDPSession) CheckPacketID(packetID uint64) bool {
 	s.Lock()
 	defer s.Unlock()
-	if s.Window == nil {
-		s.Window = new(SlidingWindow)
-	}
 	return s.Window.Check(packetID)
 }
 
-func (s *ServerUDPSession) AddPacketID(packetID uint64) {
+func (s *ServerUDPSession) Commit(packetID uint64, bodyAead cipher.AEAD, user *protocol.MemoryUser, userPSK []byte) bool {
 	s.Lock()
 	defer s.Unlock()
-	if s.Window == nil {
-		s.Window = new(SlidingWindow)
+	if s.User != nil && user != nil && s.User != user {
+		return false
 	}
-	s.Window.Add(packetID)
+	if !s.Window.CheckAndAdd(packetID) {
+		return false
+	}
+	if s.clientBodyCipher == nil && bodyAead != nil {
+		s.clientBodyCipher = bodyAead
+	}
+	if s.User == nil && user != nil {
+		s.User = user
+		s.UserPSK = userPSK
+	}
+	s.LastActive.Store(time.Now().Unix())
+	return true
 }
 
 type UDPSessionManager struct {
@@ -128,6 +136,13 @@ func NewUDPSessionManager(timeout time.Duration) *UDPSessionManager {
 	}
 }
 
+func (m *UDPSessionManager) Get(sessionID uint64) *ServerUDPSession {
+	if s, ok := m.sessions.Load(sessionID); ok {
+		return s
+	}
+	return nil
+}
+
 func (m *UDPSessionManager) GetOrCreate(sessionID uint64) *ServerUDPSession {
 	now := time.Now().Unix()
 	if s, ok := m.sessions.Load(sessionID); ok {
@@ -137,6 +152,7 @@ func (m *UDPSessionManager) GetOrCreate(sessionID uint64) *ServerUDPSession {
 
 	s := &ServerUDPSession{
 		SessionID: sessionID,
+		Window:    new(SlidingWindow),
 		manager:   m,
 	}
 	s.LastActive.Store(now)

proxy/shadowsocks_2022/shadowsocks_2022.go

@@ -63,10 +63,16 @@ func (s *ServerUDPSession) EnsureLink(
 
 	s.currentConn.Store(conn)
 
-	sessCtx, cancel := context.WithCancel(ctx)
+	// UDP Session might travel across multiple connections
+	// don't bind its lifetime to a parent context
+	sessCtx, cancel := context.WithCancel(context.WithoutCancel(ctx))
 	inbound := session.InboundFromContext(sessCtx)
-	if inbound != nil && s.User != nil {
-		inbound.User = s.User
+	if inbound != nil {
+		newInbound := *inbound
+		if s.User != nil {
+			newInbound.User = s.User
+		}
+		sessCtx = session.ContextWithInbound(sessCtx, &newInbound)
 	}
 	var email string
 	var level uint32

proxy/shadowsocks_2022/shadowsocks_2022_test.go

@@ -265,7 +265,7 @@ func TestUDPCodec(t *testing.T) {
 
 			clientCodec, err := NewUDPPacketCodec(method, [][]byte{psk})
 			common.Must(err)
-			serverCodec, err := NewUDPServerCodec(method, psk, time.Minute)
+			serverCodec, err := NewUDPServerCodec(method, psk, nil, time.Minute)
 			common.Must(err)
 
 			session, err := clientCodec.NewClientSession()
@@ -277,7 +277,7 @@ func TestUDPCodec(t *testing.T) {
 			rawCopy := make([]byte, pktBuf.Len())
 			copy(rawCopy, pktBuf.Bytes())
 
-			decoded, err := serverCodec.DecodePacket(pktBuf.Bytes())
+			_, decoded, err := serverCodec.DecodePacket(pktBuf.Bytes())
 			common.Must(err)
 
 			if decoded.HeaderType != HeaderTypeClient {
@@ -291,7 +291,7 @@ func TestUDPCodec(t *testing.T) {
 			}
 
 			// Replay same packet wire bytes should fail with ErrPacketIdNotUnique
-			_, err = serverCodec.DecodePacket(rawCopy)
+			_, _, err = serverCodec.DecodePacket(rawCopy)
 			if err != ErrPacketIdNotUnique {
 				t.Fatalf("expected ErrPacketIdNotUnique on replay, got: %v", err)
 			}
@@ -427,7 +427,7 @@ func TestClientUDPSessionMultiDestination(t *testing.T) {
 
 			clientCodec, err := NewUDPPacketCodec(method, [][]byte{rawKey})
 			common.Must(err)
-			serverCodec, err := NewUDPServerCodec(method, rawKey, time.Minute)
+			serverCodec, err := NewUDPServerCodec(method, rawKey, nil, time.Minute)
 			common.Must(err)
 
 			session, err := clientCodec.NewClientSession()
@@ -448,9 +448,9 @@ func TestClientUDPSessionMultiDestination(t *testing.T) {
 			defer pkt2.Release()
 
 			// Server decodes both
-			dec1, err := serverCodec.DecodePacket(pkt1.Bytes())
+			_, dec1, err := serverCodec.DecodePacket(pkt1.Bytes())
 			common.Must(err)
-			dec2, err := serverCodec.DecodePacket(pkt2.Bytes())
+			_, dec2, err := serverCodec.DecodePacket(pkt2.Bytes())
 			common.Must(err)
 
 			if dec1.SessionID != session.ClientSessionID() || dec2.SessionID != session.ClientSessionID() {