Fix udp match

common/net/find_process_linux.go

@@ -26,20 +26,20 @@ func FindProcess(network, srcIP string, srcPort uint16, destIP string, destPort
 		panic("Unsupported network type for process lookup.")
 	}
 
-	var procFile string
+	var procFiles []string
 
 	switch network {
 	case "tcp":
 		if net.ParseIP(srcIP).To4() != nil {
-			procFile = "/proc/net/tcp"
+			procFiles = []string{"/proc/net/tcp"}
 		} else {
-			procFile = "/proc/net/tcp6"
+			procFiles = []string{"/proc/net/tcp6"}
 		}
 	case "udp":
 		if net.ParseIP(srcIP).To4() != nil {
-			procFile = "/proc/net/udp"
+			procFiles = []string{"/proc/net/udp", "/proc/net/udp6"}
 		} else {
-			procFile = "/proc/net/udp6"
+			procFiles = []string{"/proc/net/udp6", "/proc/net/udp"}
 		}
 	default:
 		panic("Unsupported network type for process lookup.")
@@ -50,12 +50,17 @@ func FindProcess(network, srcIP string, srcPort uint16, destIP string, destPort
 		return 0, "", "", errors.New("failed to format address: ", err)
 	}
 
-	inode, err := findInodeInFile(procFile, targetHexAddr)
+	portFallback := ""
+	if network == "udp" {
+		portFallback = fmt.Sprintf(":%04X", srcPort)
+	}
+
+	inode, err := findInodeInFiles(procFiles, targetHexAddr, portFallback)
 	if err != nil {
-		return 0, "", "", errors.New("could not search in ", procFile).Base(err)
+		return 0, "", "", errors.New("could not search in ", strings.Join(procFiles, " and ")).Base(err)
 	}
 	if inode == "" {
-		return 0, "", "", errors.New("connection for ", srcIP, ":", srcPort, " not found in ", procFile)
+		return 0, "", "", errors.New("connection for ", srcIP, ":", srcPort, " not found in ", strings.Join(procFiles, " or "))
 	}
 
 	pidStr, err := findPidByInode(inode)
@@ -102,31 +107,43 @@ func formatLittleEndianString(addr net.IP, port Port) (string, error) {
 	return fmt.Sprintf("%s:%s", ipHex, portHex), nil
 }
 
-func findInodeInFile(filePath, targetHexAddr string) (string, error) {
-	file, err := os.Open(filePath)
-	if err != nil {
-		return "", err
-	}
-	defer file.Close()
+func findInodeInFiles(filePaths []string, targetHexAddr, portFallback string) (string, error) {
+	onlyMatchedPort := ""
 
-	scanner := bufio.NewScanner(file)
-
-	for scanner.Scan() {
-		line := scanner.Text()
-		fields := strings.Fields(line)
-
-		if len(fields) < 10 {
-			continue
+	for i, filePath := range filePaths {
+		file, err := os.Open(filePath)
+		if err != nil {
+			if i > 0 && os.IsNotExist(err) {
+				continue
+			}
+			return "", err
 		}
+		defer file.Close()
 
-		localAddress := fields[1]
-		if localAddress == targetHexAddr {
-			inode := fields[9]
-			return inode, nil
+		scanner := bufio.NewScanner(file)
+		for scanner.Scan() {
+			fields := strings.Fields(scanner.Text())
+			if len(fields) < 10 {
+				continue
+			}
+
+			localAddress := fields[1]
+			if localAddress == targetHexAddr {
+				return fields[9], nil
+			}
+			if portFallback != "" && strings.HasSuffix(localAddress, portFallback) {
+				wildcardIP := strings.TrimSuffix(localAddress, portFallback)
+				if (len(wildcardIP) == 8 || len(wildcardIP) == 32) && strings.Trim(wildcardIP, "0") == "" {
+					onlyMatchedPort = fields[9]
+				}
+			}
+		}
+		if err := scanner.Err(); err != nil {
+			return "", err
 		}
 	}
 
-	return "", scanner.Err()
+	return onlyMatchedPort, nil
 }
 
 func findPidByInode(inode string) (string, error) {

proxy/shadowsocks_2022/inbound.go

@@ -53,23 +53,21 @@ func NewServer(ctx context.Context, config *ServerConfig) (*Inbound, error) {
 		return nil, err
 	}
 
-	user := &protocol.MemoryUser{
-		Email: config.Email,
-		Level: uint32(config.Level),
-	}
-
-	udpCodec, err := NewUDPServerCodec(method, psk, user, 500*time.Second)
+	udpCodec, err := NewUDPServerCodec(method, psk, 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:          user,
+		networks:   networks,
+		method:     method,
+		psk:        psk,
+		saltFilter: antireplay.NewMapFilter[[32]byte](60),
+		user: &protocol.MemoryUser{
+			Email: config.Email,
+			Level: uint32(config.Level),
+		},
 		udpCodec:      udpCodec,
 		policyManager: v.GetFeature(policy.ManagerType()).(policy.Manager),
 	}, nil
@@ -157,12 +155,20 @@ func (i *Inbound) processUDP(ctx context.Context, conn stat.Connection, dispatch
 		}
 
 		for _, b := range mb {
-			sessionItem, decoded, err := i.udpCodec.DecodePacket(b.Bytes())
+			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,6 +2,8 @@ package shadowsocks_2022
 
 import (
 	"context"
+	"crypto/cipher"
+	"encoding/binary"
 	"strconv"
 	"strings"
 	"sync"
@@ -39,7 +41,8 @@ type MultiUserInbound struct {
 	usersByEmail    *utils.TypedSyncMap[string, *protocol.MemoryUser]
 	userCount       atomic.Int64
 	saltFilter      *antireplay.ReplayFilter[[32]byte]
-	udpCodec        *UDPMultiServerCodec
+	udpSessions     *UDPSessionManager
+	udpMasterCipher cipher.Block
 	policyManager   policy.Manager
 }
 
@@ -65,21 +68,22 @@ func NewMultiServer(ctx context.Context, config *MultiUserServerConfig) (*MultiU
 		return nil, err
 	}
 
-	udpCodec, err := NewUDPMultiServerCodec(method, masterPSK, 500*time.Second)
+	masterBlock, err := method.NewBlock(masterPSK)
 	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),
-		udpCodec:      udpCodec,
-		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),
+		udpSessions:     NewUDPSessionManager(500 * time.Second),
+		udpMasterCipher: masterBlock,
+		policyManager:   v.GetFeature(policy.ManagerType()).(policy.Manager),
 	}
 
 	for idx, user := range config.Users {
@@ -277,26 +281,68 @@ func (i *MultiUserInbound) processUDP(ctx context.Context, conn stat.Connection,
 		}
 
 		for _, b := range mb {
-			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 {
-					return nil, nil, false
-				}
-				return u, u.Account.(*MemoryAccount).Key, true
-			})
-			b.Release()
-			if err != nil || decoded.HeaderType != HeaderTypeClient {
+			// 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)
+				if !ok {
+					b.Release()
+					continue
+				}
+				currentUser = user
+				userPSK = user.Account.(*MemoryAccount).Key
+			}
+
+			decoded, err := sessionItem.DecryptAESPayload(i.method, userPSK, sessionID, packetID, rawHeader[:], packetBytes[32:])
+			b.Release()
+			if err != nil {
+				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.udpCodec.EncodeServerPacket(decoded.SessionID, userPSK, replyDest, payload)
+					return i.encodeServerUDPPacket(sessionID, userPSK, replyDest, payload)
 				})
 				if err != nil {
 					continue
@@ -310,3 +356,7 @@ 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,14 +243,16 @@ 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.User = &protocol.MemoryUser{
-					Email: targetDest.email,
-					Level: targetDest.level,
+				sessionItem.Lock()
+				if sessionItem.User == nil {
+					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      *UDPClientCodec
+	udpCodec      *UDPPacketCodec
 	policyManager policy.Manager
 }
 
@@ -50,9 +50,9 @@ func NewClient(ctx context.Context, config *ClientConfig) (*Outbound, error) {
 	}
 
 	finalPSK := pskList[len(pskList)-1]
-	udpCodec, err := NewUDPClientCodec(method, pskList)
+	udpCodec, err := NewUDPPacketCodec(method, pskList)
 	if err != nil {
-		return nil, errors.New("failed to create udp client codec").Base(err)
+		return nil, errors.New("failed to create udp packet codec").Base(err)
 	}
 
 	v := core.MustFromContext(ctx)

proxy/shadowsocks_2022/packet.go

@@ -7,78 +7,33 @@ 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 UDPClientCodec struct {
+type UDPCodec 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
 }
 
-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),
+type (
+	UDPPacketCodec = UDPCodec
+	UDPServerCodec = UDPCodec
+)
+
+func newUDPCodec(method *CipherMethod, psk []byte) (*UDPCodec, error) {
+	c := &UDPCodec{
+		method: method,
+		psk:    psk,
 	}
 	var err error
 	if method.IsChaCha {
@@ -92,129 +47,46 @@ func NewUDPServerCodec(method *CipherMethod, psk []byte, user *protocol.MemoryUs
 	return c, nil
 }
 
-func (c *UDPServerCodec) Sessions() *UDPSessionManager {
-	return c.sessions
-}
-
-func (c *UDPServerCodec) GetSession(sessionID uint64) *ServerUDPSession {
-	if c.sessions == nil {
-		return 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")
 	}
-	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)
+	finalPSK := pskList[len(pskList)-1]
+	c, err := newUDPCodec(method, finalPSK)
 	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
+	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
 			}
 		}
 	}
-
-	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
+	return c, 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)
+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 {
+	return c.sessions
+}
+
+func (c *UDPCodec) GetSession(sessionID uint64) *ServerUDPSession {
+	if c.sessions == nil {
+		return nil
+	}
+	return c.sessions.GetOrCreate(sessionID)
 }
 
 type DecodedUDPPacket struct {
@@ -321,42 +193,44 @@ func parsePlainUDPPacket(sessionID, packetID uint64, bodyPlain []byte) (DecodedU
 	}, nil
 }
 
-func (c *UDPServerCodec) DecodePacket(data []byte) (*ServerUDPSession, DecodedUDPPacket, error) {
+func (c *UDPCodec) DecodePacket(data []byte) (DecodedUDPPacket, error) {
 	if len(data) < PacketMinimalHeaderSize {
-		return nil, DecodedUDPPacket{}, ErrPacketTooShort
+		return DecodedUDPPacket{}, ErrPacketTooShort
 	}
 
 	if c.method.IsChaCha {
 		if len(data) < PacketNonceSize+AEADTagSize {
-			return nil, DecodedUDPPacket{}, ErrPacketTooShort
+			return DecodedUDPPacket{}, ErrPacketTooShort
 		}
 		nonce := data[:PacketNonceSize]
 		ciphertext := data[PacketNonceSize:]
 		plain, err := c.chachaCipher.Open(nil, nonce, ciphertext, nil)
 		if err != nil {
-			return nil, DecodedUDPPacket{}, errors.New("failed to decrypt chacha udp packet").Base(err)
+			return DecodedUDPPacket{}, errors.New("failed to decrypt chacha udp packet").Base(err)
 		}
 		if len(plain) < 16+1+8+2 {
-			return nil, DecodedUDPPacket{}, ErrPacketTooShort
+			return 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 nil, DecodedUDPPacket{}, err
+			return DecodedUDPPacket{}, err
 		}
 
 		if decoded.HeaderType != HeaderTypeClient {
-			return nil, DecodedUDPPacket{}, ErrBadHeaderType
+			return DecodedUDPPacket{}, ErrBadHeaderType
 		}
 
-		sessionItem := c.sessions.GetOrCreate(sessionID)
-		if !sessionItem.Commit(packetID, nil, c.user, c.psk) {
-			return nil, DecodedUDPPacket{}, ErrPacketIdNotUnique
-		}
-		return sessionItem, decoded, nil
+		sessionItem.AddPacketID(packetID)
+		return decoded, nil
 	}
 
 	// AES mode
@@ -365,48 +239,49 @@ func (c *UDPServerCodec) DecodePacket(data []byte) (*ServerUDPSession, DecodedUD
 	sessionID := binary.BigEndian.Uint64(rawHeader[:8])
 	packetID := binary.BigEndian.Uint64(rawHeader[8:16])
 
-	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()
+	sessionItem := c.sessions.GetOrCreate(sessionID)
+	if !sessionItem.CheckPacketID(packetID) {
+		return DecodedUDPPacket{}, ErrPacketIdNotUnique
 	}
 
+	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(c.psk, rawHeader[:8], c.method.KeySaltLength)
+		bodyKey := DeriveSessionSubKey(psk, rawHeader[:8], method.KeySaltLength)
 		var err error
-		bodyAead, err = c.method.NewAEAD(bodyKey)
+		bodyAead, err = method.NewAEAD(bodyKey)
 		if err != nil {
-			return nil, DecodedUDPPacket{}, err
+			return DecodedUDPPacket{}, err
 		}
+		isNewCipher = true
 	}
 
 	bodyNonce := rawHeader[4:16]
-	bodyCipher := data[16:]
 	bodyPlain, err := bodyAead.Open(nil, bodyNonce, bodyCipher, nil)
 	if err != nil {
-		return nil, DecodedUDPPacket{}, errors.New("failed to decrypt aes udp body").Base(err)
+		return DecodedUDPPacket{}, errors.New("failed to decrypt aes udp body").Base(err)
 	}
 
 	decoded, err := parsePlainUDPPacket(sessionID, packetID, bodyPlain)
 	if err != nil {
-		return nil, DecodedUDPPacket{}, err
-	}
-	if decoded.HeaderType != HeaderTypeClient {
-		return nil, DecodedUDPPacket{}, ErrBadHeaderType
+		return DecodedUDPPacket{}, err
 	}
 
-	if sessionItem == nil {
-		sessionItem = c.sessions.GetOrCreate(sessionID)
+	if decoded.HeaderType != HeaderTypeClient {
+		return DecodedUDPPacket{}, ErrBadHeaderType
 	}
-	if !sessionItem.Commit(packetID, bodyAead, c.user, c.psk) {
-		return nil, DecodedUDPPacket{}, ErrPacketIdNotUnique
+
+	s.AddPacketID(packetID)
+
+	if isNewCipher {
+		s.clientBodyCipher = bodyAead
 	}
-	return sessionItem, decoded, nil
+
+	return decoded, nil
 }
 
 func (s *ServerUDPSession) EnsureServerState(method *CipherMethod, psk []byte) error {
@@ -428,9 +303,6 @@ 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
 	}
 
@@ -515,12 +387,11 @@ func (s *ServerUDPSession) EncodeServerPacket(method *CipherMethod, clientSessio
 	return res, nil
 }
 
-func (c *UDPServerCodec) EncodeServerPacket(clientSessionID uint64, dest net.Destination, payload []byte) ([]byte, error) {
+func (c *UDPCodec) 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
@@ -528,20 +399,21 @@ type serverSessionState struct {
 }
 
 func (st *serverSessionState) check(packetID uint64) bool {
-	st.Lock()
-	defer st.Unlock()
+	if st.window == nil {
+		st.window = new(SlidingWindow)
+	}
 	return st.window.Check(packetID)
 }
 
-func (st *serverSessionState) commit(packetID uint64) bool {
-	st.Lock()
-	defer st.Unlock()
-	return st.window.CheckAndAdd(packetID)
+func (st *serverSessionState) add(packetID uint64) {
+	if st.window == nil {
+		st.window = new(SlidingWindow)
+	}
+	st.window.Add(packetID)
 }
 
 type ClientUDPSession struct {
-	mu               sync.Mutex
-	codec            *UDPClientCodec
+	codec            *UDPCodec
 	clientSessionID  uint64
 	nextPacketID     atomic.Uint64
 	clientBodyCipher cipher.AEAD
@@ -549,7 +421,7 @@ type ClientUDPSession struct {
 	old              atomic.Pointer[serverSessionState]
 }
 
-func (c *UDPClientCodec) NewClientSession() (*ClientUDPSession, error) {
+func (c *UDPCodec) NewClientSession() (*ClientUDPSession, error) {
 	var sessID [8]byte
 	if _, err := io.ReadFull(rand.Reader, sessID[:]); err != nil {
 		return nil, err
@@ -574,7 +446,7 @@ func (c *UDPClientCodec) NewClientSession() (*ClientUDPSession, error) {
 	}, nil
 }
 
-func (s *ClientUDPSession) getOrTransitionServerSession(sessionID uint64, cipher cipher.AEAD, now int64) (*serverSessionState, error) {
+func (s *ClientUDPSession) getServerSession(sessionID uint64, now int64) (*serverSessionState, error) {
 	cur := s.current.Load()
 	if cur != nil && cur.sessionID == sessionID {
 		return cur, nil
@@ -583,42 +455,39 @@ func (s *ClientUDPSession) getOrTransitionServerSession(sessionID uint64, cipher
 	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
 	}
 
-	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
-	}
-
+	// New server session:
 	// 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,
-		window:    new(SlidingWindow),
-		cipher:    cipher,
+		cipher:    bodyAead,
 	}
 	newState.lastSeen.Store(now)
 
 	if cur == nil {
-		s.current.Store(newState)
-		return newState, nil
+		s.current.CompareAndSwap(nil, newState)
+		return s.current.Load(), nil
 	}
 
 	s.old.Store(cur)
@@ -768,6 +637,15 @@ 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
@@ -780,16 +658,7 @@ func (s *ClientUDPSession) DecodePacket(data []byte) (DecodedUDPPacket, error) {
 			return DecodedUDPPacket{}, errors.New("client session ID mismatch")
 		}
 
-		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.add(packetID)
 		st.lastSeen.Store(now)
 
 		return decoded, nil
@@ -802,79 +671,18 @@ func (s *ClientUDPSession) DecodePacket(data []byte) (DecodedUDPPacket, error) {
 	packetID := binary.BigEndian.Uint64(rawHeader[8:16])
 
 	now := time.Now().Unix()
-
-	// 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)
+	st, err := s.getServerSession(sessionID, now)
 	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 := trialAead.Open(nil, bodyNonce, bodyCipher, nil)
+	bodyPlain, err := bodyAead.Open(nil, bodyNonce, bodyCipher, nil)
 	if err != nil {
 		return DecodedUDPPacket{}, errors.New("failed to decrypt aes udp body").Base(err)
 	}
@@ -891,15 +699,7 @@ func (s *ClientUDPSession) DecodePacket(data []byte) (DecodedUDPPacket, error) {
 		return DecodedUDPPacket{}, errors.New("client session ID mismatch")
 	}
 
-	// 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.add(packetID)
 	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, nil, 300*time.Second)
+	destCodec, err := NewUDPServerCodec(method, destKey, 300*time.Second)
 	common.Must(err)
 
 	mu.Lock()
@@ -202,8 +202,7 @@ func TestRelayUDPSessionStabilityAndDispatch(t *testing.T) {
 		t.Fatalf("expected 2 received packets at destination, got %d", len(pkts))
 	}
 
-	dec1, dec2 := DecodedUDPPacket{}, DecodedUDPPacket{}
-	_, dec1, err = destCodec.DecodePacket(pkts[0])
+	dec1, err := destCodec.DecodePacket(pkts[0])
 	if err != nil {
 		t.Fatalf("dest failed to decode packet 1: %v", err)
 	}
@@ -211,7 +210,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,27 +100,19 @@ 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) Commit(packetID uint64, bodyAead cipher.AEAD, user *protocol.MemoryUser, userPSK []byte) bool {
+func (s *ServerUDPSession) AddPacketID(packetID uint64) {
 	s.Lock()
 	defer s.Unlock()
-	if s.User != nil && user != nil && s.User != user {
-		return false
+	if s.Window == nil {
+		s.Window = new(SlidingWindow)
 	}
-	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
+	s.Window.Add(packetID)
 }
 
 type UDPSessionManager struct {
@@ -136,13 +128,6 @@ 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 {
@@ -152,7 +137,6 @@ 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,16 +63,10 @@ func (s *ServerUDPSession) EnsureLink(
 
 	s.currentConn.Store(conn)
 
-	// UDP Session might travel across multiple connections
-	// don't bind its lifetime to a parent context
-	sessCtx, cancel := context.WithCancel(context.WithoutCancel(ctx))
+	sessCtx, cancel := context.WithCancel(ctx)
 	inbound := session.InboundFromContext(sessCtx)
-	if inbound != nil {
-		newInbound := *inbound
-		if s.User != nil {
-			newInbound.User = s.User
-		}
-		sessCtx = session.ContextWithInbound(sessCtx, &newInbound)
+	if inbound != nil && s.User != nil {
+		inbound.User = s.User
 	}
 	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, nil, time.Minute)
+			serverCodec, err := NewUDPServerCodec(method, psk, 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, nil, time.Minute)
+			serverCodec, err := NewUDPServerCodec(method, rawKey, 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() {