refine udp

Co-authored-by: ChownYourLife <260039878+ChownYourLife@users.noreply.github.com>

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 sessionItem.User == nil {
-				sessionItem.Lock()
+			if link == nil {
+				sessionItem = i.udpCodec.GetSession(decoded.SessionID)
 				if sessionItem.User == nil {
-					sessionItem.User = i.user
+					sessionItem.Lock()
+					if sessionItem.User == nil {
+						sessionItem.User = i.user
+					}
+					sessionItem.Unlock()
 				}
-				sessionItem.Unlock()
-			}
-			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)
-			})
-			if err != nil {
-				continue
+				var err 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)
+				})
+				if err != nil {
+					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) {
-				return i.encodeServerUDPPacket(sessionID, userPSK, replyDest, payload)
-			})
-			if err != nil {
-				continue
+			if link == nil {
+				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)
+				})
+				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
-			dest.Network = net.Network_UDP
+			if link == nil {
+				dest := targetDest.destination
+				dest.Network = net.Network_UDP
 
-			sessionItem := i.udpSessions.GetOrCreate(sessionID)
-			if sessionItem.User == nil {
-				sessionItem.Lock()
+				sessionItem = i.udpSessions.GetOrCreate(sessionID)
 				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, err := sessionItem.EnsureLink(ctx, conn, dest, dispatcher, i.policyManager, nil)
-			if err != nil {
-				b.Release()
-				continue
+				var err error
+				link, err = sessionItem.EnsureLink(ctx, conn, dest, dispatcher, i.policyManager, nil)
+				if err != nil {
+					b.Release()
+					continue
+				}
+			} else {
+				sessionItem.UpdateConn(conn)
 			}
 
 			_ = link.Writer.WriteMultiBuffer(buf.MultiBuffer{b})

proxy/shadowsocks_2022/outbound.go

@@ -169,8 +169,8 @@ func (o *Outbound) Process(ctx context.Context, link *transport.Link, dialer int
 			return buf.Copy(responseReader, link.Writer, buf.UpdateActivity(timer))
 		}
 
-		responseDoneAndCloseWriter := task.OnSuccess(responseDone, task.Close(link.Writer))
-		if err := task.Run(ctx, requestDone, responseDoneAndCloseWriter); err != nil {
+		responseDonePost := task.OnSuccess(responseDone, task.Close(link.Writer))
+		if err := task.Run(ctx, requestDone, responseDonePost); err != nil {
 			return errors.New("connection ends").Base(err)
 		}
 
@@ -212,8 +212,8 @@ func (o *Outbound) Process(ctx context.Context, link *transport.Link, dialer int
 			return nil
 		}
 
-		responseDoneAndCloseWriter := task.OnSuccess(responseDone, task.Close(link.Writer))
-		if err := task.Run(ctx, requestDone, responseDoneAndCloseWriter); err != nil {
+		responseDonePost := task.OnSuccess(responseDone, task.Close(link.Writer))
+		if err := task.Run(ctx, requestDone, responseDonePost); err != nil {
 			return errors.New("connection ends").Base(err)
 		}
 

proxy/shadowsocks_2022/shadowsocks_2022.go

@@ -18,12 +18,8 @@ import (
 )
 
 func (s *ServerUDPSession) UpdateConn(conn stat.Connection) {
-	if s.currentConn.Load() == nil {
-		s.currentConn.Store(conn)
-	}
-	if s.timer != nil {
-		s.timer.Update()
-	}
+	s.currentConn.Store(conn)
+	s.timer.Update()
 }
 
 func (s *ServerUDPSession) WriteToClient(b []byte) error {
@@ -57,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 {
@@ -94,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 {
@@ -103,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
@@ -118,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() {
-		if s.timer != nil {
-			s.timer.SetTimeout(0)
-		}
-	}()
+	defer 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 {

proxy/shadowsocks_2022/stream.go

@@ -11,6 +11,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/xtls/xray-core/common"
 	"github.com/xtls/xray-core/common/antireplay"
 	"github.com/xtls/xray-core/common/buf"
 	"github.com/xtls/xray-core/common/errors"
@@ -629,6 +630,9 @@ func InitServerStream(conn net.Conn, method *CipherMethod, psk, saltSlice []byte
 func TransportTCP(ctx context.Context, sessionPolicy policy.Session, reader buf.Reader, writer buf.Writer, link *transport.Link) error {
 	ctx, cancel := context.WithCancel(ctx)
 	timer := signal.CancelAfterInactivity(ctx, cancel, sessionPolicy.Timeouts.ConnectionIdle)
+	defer timer.SetTimeout(0)
+	defer common.Interrupt(link.Reader)
+	defer common.Interrupt(link.Writer)
 	ctx = policy.ContextWithBufferPolicy(ctx, sessionPolicy.Buffer)
 
 	requestDone := func() error {
@@ -644,6 +648,6 @@ func TransportTCP(ctx context.Context, sessionPolicy policy.Session, reader buf.
 		return buf.Copy(link.Reader, writer, buf.UpdateActivity(timer))
 	}
 
-	responseDoneAndCloseWriter := task.OnSuccess(responseDone, task.Close(link.Writer))
-	return task.Run(ctx, requestDone, responseDoneAndCloseWriter)
+	requestDonePost := task.OnSuccess(requestDone, task.Close(link.Writer))
+	return task.Run(ctx, requestDonePost, responseDone)
 }