SS2022 inbound: Fix data races and duplicate acceptance when decoding one UDP session concurrently

Packets of one client session can be processed by two inbound workers at
the same time (e.g. the client's old and new source address after a NAT
rebinding). The sliding window was checked and updated in two separate
critical sections, and clientBodyCipher and the session's User were read
without the lock, so the same packet ID could be accepted more than once
and go test -race reported data races.

Re-check the window and publish the cached body cipher in one critical
section after the packet is fully validated, and always initialize the
session's User under the lock.

proxy/shadowsocks_2022/inbound.go

@@ -162,13 +162,11 @@ func (i *Inbound) processUDP(ctx context.Context, conn stat.Connection, dispatch
 			}
 
 			sessionItem := i.udpCodec.GetSession(decoded.SessionID)
+			sessionItem.Lock()
 			if sessionItem.User == nil {
-				sessionItem.Lock()
-				if sessionItem.User == nil {
-					sessionItem.User = i.user
-				}
-				sessionItem.Unlock()
+				sessionItem.User = i.user
 			}
+			sessionItem.Unlock()
 			link := sessionItem.link.Load()
 			if link != nil {
 				sessionItem.UpdateConn(conn)

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/packet.go

@@ -229,7 +229,9 @@ func (c *UDPCodec) DecodePacket(data []byte) (DecodedUDPPacket, error) {
 			return DecodedUDPPacket{}, ErrBadHeaderType
 		}
 
-		sessionItem.AddPacketID(packetID)
+		if !sessionItem.CommitPacketID(packetID) {
+			return DecodedUDPPacket{}, ErrPacketIdNotUnique
+		}
 		return decoded, nil
 	}
 
@@ -248,8 +250,7 @@ func (c *UDPCodec) DecodePacket(data []byte) (DecodedUDPPacket, error) {
 }
 
 func (s *ServerUDPSession) DecryptAESPayload(method *CipherMethod, psk []byte, sessionID, packetID uint64, rawHeader, bodyCipher []byte) (DecodedUDPPacket, error) {
-	bodyAead := s.clientBodyCipher
-	isNewCipher := false
+	bodyAead := s.getClientBodyCipher()
 	if bodyAead == nil {
 		bodyKey := DeriveSessionSubKey(psk, rawHeader[:8], method.KeySaltLength)
 		var err error
@@ -257,7 +258,6 @@ func (s *ServerUDPSession) DecryptAESPayload(method *CipherMethod, psk []byte, s
 		if err != nil {
 			return DecodedUDPPacket{}, err
 		}
-		isNewCipher = true
 	}
 
 	bodyNonce := rawHeader[4:16]
@@ -275,10 +275,8 @@ func (s *ServerUDPSession) DecryptAESPayload(method *CipherMethod, psk []byte, s
 		return DecodedUDPPacket{}, ErrBadHeaderType
 	}
 
-	s.AddPacketID(packetID)
-
-	if isNewCipher {
-		s.clientBodyCipher = bodyAead
+	if !s.commitClientPacket(packetID, bodyAead) {
+		return DecodedUDPPacket{}, ErrPacketIdNotUnique
 	}
 
 	return decoded, nil

proxy/shadowsocks_2022/replay.go

@@ -106,13 +106,37 @@ func (s *ServerUDPSession) CheckPacketID(packetID uint64) bool {
 	return s.Window.Check(packetID)
 }
 
-func (s *ServerUDPSession) AddPacketID(packetID uint64) {
+// CommitPacketID records a packet ID after the packet has been fully validated.
+// It re-checks the window under the lock, so when the same packet is validated
+// concurrently (e.g. by the workers of two client source addresses), only one is accepted.
+func (s *ServerUDPSession) CommitPacketID(packetID uint64) bool {
 	s.Lock()
 	defer s.Unlock()
 	if s.Window == nil {
 		s.Window = new(SlidingWindow)
 	}
-	s.Window.Add(packetID)
+	return s.Window.CheckAndAdd(packetID)
+}
+
+func (s *ServerUDPSession) getClientBodyCipher() cipher.AEAD {
+	s.Lock()
+	defer s.Unlock()
+	return s.clientBodyCipher
+}
+
+func (s *ServerUDPSession) commitClientPacket(packetID uint64, bodyAead cipher.AEAD) bool {
+	s.Lock()
+	defer s.Unlock()
+	if s.Window == nil {
+		s.Window = new(SlidingWindow)
+	}
+	if !s.Window.CheckAndAdd(packetID) {
+		return false
+	}
+	if s.clientBodyCipher == nil {
+		s.clientBodyCipher = bodyAead
+	}
+	return true
 }
 
 type UDPSessionManager struct {

proxy/shadowsocks_2022/shadowsocks_2022_test.go

@@ -496,3 +496,60 @@ func TestClientUDPSessionMultiDestination(t *testing.T) {
 		})
 	}
 }
+
+// Packets of one client session may be decoded concurrently, e.g. by the inbound
+// workers of the client's old and new source address after a NAT rebinding.
+func TestUDPServerConcurrentDecodeSameSession(t *testing.T) {
+	for _, methodName := range []string{MethodAES128GCM, MethodAES256GCM, MethodChaCha20Poly1305} {
+		t.Run(methodName, func(t *testing.T) {
+			method, err := GetCipherMethod(methodName)
+			common.Must(err)
+			rawKey := make([]byte, method.KeySaltLength)
+			_, _ = rand.Read(rawKey)
+
+			clientCodec, err := NewUDPPacketCodec(method, [][]byte{rawKey})
+			common.Must(err)
+			serverCodec, err := NewUDPServerCodec(method, rawKey, time.Minute)
+			common.Must(err)
+			session, err := clientCodec.NewClientSession()
+			common.Must(err)
+
+			dest := net.UDPDestination(net.LocalHostIP, net.Port(443))
+			packets := make([][]byte, 256)
+			for i := range packets {
+				pkt, err := session.EncodePacket(dest, []byte("payload"))
+				common.Must(err)
+				packets[i] = append([]byte(nil), pkt.Bytes()...)
+				pkt.Release()
+			}
+
+			const workers = 4
+			var wg sync.WaitGroup
+			var mu sync.Mutex
+			accepted := make([]int, len(packets))
+			for w := 0; w < workers; w++ {
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					// Every worker sees every packet, so each packet must be accepted exactly once.
+					for i, pkt := range packets {
+						if _, err := serverCodec.DecodePacket(pkt); err == nil {
+							mu.Lock()
+							accepted[i]++
+							mu.Unlock()
+						} else if err != ErrPacketIdNotUnique {
+							t.Errorf("packet %d: unexpected error: %v", i, err)
+						}
+					}
+				}()
+			}
+			wg.Wait()
+
+			for i, n := range accepted {
+				if n != 1 {
+					t.Fatalf("packet %d accepted %d times, want 1", i, n)
+				}
+			}
+		})
+	}
+}