udp

proxy/shadowsocks_2022/inbound.go

@@ -169,11 +169,17 @@ func (i *Inbound) processUDP(ctx context.Context, conn stat.Connection, dispatch
 				}
 				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
+			link := sessionItem.link.Load()
+			if link != nil {
+				sessionItem.UpdateConn(conn)
+			} else {
+				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
+				}
 			}
 
 			payloadBuf := buf.New()

proxy/shadowsocks_2022/inbound_multi.go

@@ -336,11 +336,17 @@ 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
+			link := sessionItem.link.Load()
+			if link != nil {
+				sessionItem.UpdateConn(conn)
+			} else {
+				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
+				}
 			}
 
 			pBuf := buf.New()

proxy/shadowsocks_2022/inbound_relay.go

@@ -253,10 +253,16 @@ func (i *RelayInbound) processUDP(ctx context.Context, conn stat.Connection, dis
 				}
 				sessionItem.Unlock()
 			}
-			link, err := sessionItem.EnsureLink(ctx, conn, dest, dispatcher, i.policyManager, nil)
-			if err != nil {
-				b.Release()
-				continue
+			link := sessionItem.link.Load()
+			if link != nil {
+				sessionItem.UpdateConn(conn)
+			} else {
+				var err error
+				link, err = sessionItem.EnsureLink(ctx, conn, dest, dispatcher, i.policyManager, nil)
+				if err != nil {
+					b.Release()
+					continue
+				}
 			}
 
 			_ = 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,7 @@ 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)
 	ctx = policy.ContextWithBufferPolicy(ctx, sessionPolicy.Buffer)
 
 	requestDone := func() error {
@@ -644,6 +646,12 @@ 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)
+	requestDoneAndCloseWriter := task.OnSuccess(requestDone, task.Close(link.Writer))
+	if err := task.Run(ctx, requestDoneAndCloseWriter, responseDone); err != nil {
+		common.Interrupt(link.Reader)
+		common.Interrupt(link.Writer)
+		return errors.New("connection ends").Base(err)
+	}
+
+	return nil
 }