diff --git a/proxy/shadowsocks_2022/inbound.go b/proxy/shadowsocks_2022/inbound.go index 713c05d79..af9be3836 100644 --- a/proxy/shadowsocks_2022/inbound.go +++ b/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) diff --git a/proxy/shadowsocks_2022/inbound_relay.go b/proxy/shadowsocks_2022/inbound_relay.go index f92ce2257..d12d838c2 100644 --- a/proxy/shadowsocks_2022/inbound_relay.go +++ b/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) diff --git a/proxy/shadowsocks_2022/packet.go b/proxy/shadowsocks_2022/packet.go index 310d3d7bd..76f713f30 100644 --- a/proxy/shadowsocks_2022/packet.go +++ b/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 diff --git a/proxy/shadowsocks_2022/replay.go b/proxy/shadowsocks_2022/replay.go index 056acbb2f..45ffb227a 100644 --- a/proxy/shadowsocks_2022/replay.go +++ b/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 { diff --git a/proxy/shadowsocks_2022/shadowsocks_2022_test.go b/proxy/shadowsocks_2022/shadowsocks_2022_test.go index 382696efb..b7a600194 100644 --- a/proxy/shadowsocks_2022/shadowsocks_2022_test.go +++ b/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) + } + } + }) + } +}