refine udp

proxy/shadowsocks_2022/inbound.go

@@ -15,6 +15,7 @@ import (
 	"github.com/xtls/xray-core/core"
 	"github.com/xtls/xray-core/features/policy"
 	"github.com/xtls/xray-core/features/routing"
+	"github.com/xtls/xray-core/transport"
 	"github.com/xtls/xray-core/transport/internet/stat"
 )
 
@@ -146,6 +147,9 @@ func (i *Inbound) processTCP(ctx context.Context, conn net.Conn, dispatcher rout
 }
 
 func (i *Inbound) processUDP(ctx context.Context, conn stat.Connection, dispatcher routing.Dispatcher) error {
+	var sessionItem *ServerUDPSession
+	var link *transport.Link
+
 	reader := buf.NewPacketReader(conn)
 	for {
 		mb, err := reader.ReadMultiBuffer()
@@ -161,19 +165,24 @@ func (i *Inbound) processUDP(ctx context.Context, conn stat.Connection, dispatch
 				continue
 			}
 
-			sessionItem := i.udpCodec.GetSession(decoded.SessionID)
+			if link == nil {
-			if sessionItem.User == nil {
+				sessionItem = i.udpCodec.GetSession(decoded.SessionID)
-				sessionItem.Lock()
 				if sessionItem.User == nil {
-					sessionItem.User = i.user
+					sessionItem.Lock()
+					if sessionItem.User == nil {
+						sessionItem.User = i.user
+					}
+					sessionItem.Unlock()
 				}
-				sessionItem.Unlock()
+				var err error
-			}
+				link, err = sessionItem.EnsureLink(ctx, conn, decoded.Destination, dispatcher, i.policyManager, func(dest net.Destination, payload []byte) ([]byte, error) {
-			link, err := sessionItem.EnsureLink(ctx, conn, decoded.Destination, dispatcher, i.policyManager, func(dest net.Destination, payload []byte) ([]byte, error) {
+					return i.udpCodec.EncodeServerPacket(decoded.SessionID, dest, payload)
-				return i.udpCodec.EncodeServerPacket(decoded.SessionID, dest, payload)
+				})
-			})
+				if err != nil {
-			if err != nil {
+					continue
-				continue
+				}
+			} else {
+				sessionItem.UpdateConn(conn)
 			}
 
 			payloadBuf := buf.New()

proxy/shadowsocks_2022/inbound_multi.go

@@ -23,6 +23,7 @@ import (
 	"github.com/xtls/xray-core/core"
 	"github.com/xtls/xray-core/features/policy"
 	"github.com/xtls/xray-core/features/routing"
+	"github.com/xtls/xray-core/transport"
 	"github.com/xtls/xray-core/transport/internet/stat"
 )
 
@@ -272,6 +273,9 @@ func (i *MultiUserInbound) processTCP(ctx context.Context, conn net.Conn, dispat
 }
 
 func (i *MultiUserInbound) processUDP(ctx context.Context, conn stat.Connection, dispatcher routing.Dispatcher) error {
+	var sessionItem *ServerUDPSession
+	var link *transport.Link
+
 	reader := buf.NewPacketReader(conn)
 	for {
 		mb, err := reader.ReadMultiBuffer()
@@ -296,7 +300,7 @@ func (i *MultiUserInbound) processUDP(ctx context.Context, conn stat.Connection,
 			sessionID := binary.BigEndian.Uint64(rawHeader[:8])
 			packetID := binary.BigEndian.Uint64(rawHeader[8:16])
 
-			sessionItem := i.udpSessions.GetOrCreate(sessionID)
+			sessionItem = i.udpSessions.GetOrCreate(sessionID)
 
 			if !sessionItem.CheckPacketID(packetID) {
 				b.Release()
@@ -336,11 +340,16 @@ func (i *MultiUserInbound) processUDP(ctx context.Context, conn stat.Connection,
 			}
 			sessionItem.Unlock()
 
-			link, err := sessionItem.EnsureLink(ctx, conn, decoded.Destination, dispatcher, i.policyManager, func(replyDest net.Destination, payload []byte) ([]byte, error) {
+			if link == nil {
-				return i.encodeServerUDPPacket(sessionID, userPSK, replyDest, payload)
+				var err error
-			})
+				link, err = sessionItem.EnsureLink(ctx, conn, decoded.Destination, dispatcher, i.policyManager, func(replyDest net.Destination, payload []byte) ([]byte, error) {
-			if err != nil {
+					return i.encodeServerUDPPacket(sessionID, userPSK, replyDest, payload)
-				continue
+				})
+				if err != nil {
+					continue
+				}
+			} else {
+				sessionItem.UpdateConn(conn)
 			}
 
 			pBuf := buf.New()

proxy/shadowsocks_2022/inbound_relay.go

@@ -18,6 +18,7 @@ import (
 	"github.com/xtls/xray-core/core"
 	"github.com/xtls/xray-core/features/policy"
 	"github.com/xtls/xray-core/features/routing"
+	"github.com/xtls/xray-core/transport"
 	"github.com/xtls/xray-core/transport/internet/stat"
 )
 
@@ -203,6 +204,9 @@ func (i *RelayInbound) processTCP(ctx context.Context, conn net.Conn, dispatcher
 }
 
 func (i *RelayInbound) processUDP(ctx context.Context, conn stat.Connection, dispatcher routing.Dispatcher) error {
+	var sessionItem *ServerUDPSession
+	var link *transport.Link
+
 	reader := buf.NewPacketReader(conn)
 	for {
 		mb, err := reader.ReadMultiBuffer()
@@ -239,24 +243,29 @@ func (i *RelayInbound) processUDP(ctx context.Context, conn stat.Connection, dis
 			copy(data[AESBlockSize:2*AESBlockSize], packetHeader[:])
 			b.Advance(int32(AESBlockSize))
 
-			dest := targetDest.destination
+			if link == nil {
-			dest.Network = net.Network_UDP
+				dest := targetDest.destination
+				dest.Network = net.Network_UDP
 
-			sessionItem := i.udpSessions.GetOrCreate(sessionID)
+				sessionItem = i.udpSessions.GetOrCreate(sessionID)
-			if sessionItem.User == nil {
-				sessionItem.Lock()
 				if sessionItem.User == nil {
-					sessionItem.User = &protocol.MemoryUser{
+					sessionItem.Lock()
-						Email: targetDest.email,
+					if sessionItem.User == nil {
-						Level: targetDest.level,
+						sessionItem.User = &protocol.MemoryUser{
+							Email: targetDest.email,
+							Level: targetDest.level,
+						}
 					}
+					sessionItem.Unlock()
 				}
-				sessionItem.Unlock()
+				var err error
-			}
+				link, err = sessionItem.EnsureLink(ctx, conn, dest, dispatcher, i.policyManager, nil)
-			link, err := sessionItem.EnsureLink(ctx, conn, dest, dispatcher, i.policyManager, nil)
+				if err != nil {
-			if err != nil {
+					b.Release()
-				b.Release()
+					continue
-				continue
+				}
+			} else {
+				sessionItem.UpdateConn(conn)
 			}
 
 			_ = link.Writer.WriteMultiBuffer(buf.MultiBuffer{b})

proxy/shadowsocks_2022/shadowsocks_2022.go

@@ -19,9 +19,7 @@ import (
 
 func (s *ServerUDPSession) UpdateConn(conn stat.Connection) {
 	s.currentConn.Store(conn)
-	if s.timer != nil {
+	s.timer.Update()
-		s.timer.Update()
-	}
 }
 
 func (s *ServerUDPSession) WriteToClient(b []byte) error {
@@ -55,19 +53,16 @@ func (s *ServerUDPSession) EnsureLink(
 	policyManager policy.Manager,
 	responseEncoder func(dest net.Destination, payload []byte) ([]byte, error),
 ) (*transport.Link, error) {
-	s.UpdateConn(conn)
-
-	if link := s.link.Load(); link != nil {
-		return link, nil
-	}
-
 	s.Lock()
 	defer s.Unlock()
 
 	if link := s.link.Load(); link != nil {
+		s.UpdateConn(conn)
 		return link, nil
 	}
 
+	s.currentConn.Store(conn)
+
 	sessCtx, cancel := context.WithCancel(ctx)
 	inbound := session.InboundFromContext(sessCtx)
 	if inbound != nil && s.User != nil {
@@ -92,7 +87,6 @@ func (s *ServerUDPSession) EnsureLink(
 		return nil, err
 	}
 
-	s.link.Store(link)
 	sessionPolicy := policyManager.ForLevel(level)
 	s.timer = signal.CancelAfterInactivity(sessCtx, func() {
 		if s.manager != nil {
@@ -101,6 +95,7 @@ func (s *ServerUDPSession) EnsureLink(
 		s.Close()
 		cancel()
 	}, sessionPolicy.Timeouts.ConnectionIdle)
+	s.link.Store(link)
 
 	go handleUDPResponse(s, link, dest, responseEncoder)
 	return link, nil
@@ -116,19 +111,13 @@ func ResetTCPConn(conn net.Conn) {
 }
 
 func handleUDPResponse(s *ServerUDPSession, link *transport.Link, fallbackDest net.Destination, encode func(dest net.Destination, payload []byte) ([]byte, error)) {
-	defer func() {
+	defer s.timer.SetTimeout(0)
-		if s.timer != nil {
-			s.timer.SetTimeout(0)
-		}
-	}()
 	for {
 		resMb, err := link.Reader.ReadMultiBuffer()
 		if err != nil {
 			return
 		}
-		if s.timer != nil {
+		s.timer.Update()
-			s.timer.Update()
-		}
 		for i, rb := range resMb {
 			b := rb.Bytes()
 			if encode != nil {