From 4e826abebf3711335b44fec52cdf3451d1c8b4cc Mon Sep 17 00:00:00 2001
From: patterniha <71074308+patterniha@users.noreply.github.com>
Date: Wed, 23 Jul 2025 14:53:37 +0200
Subject: [PATCH] Chore: Three small fixes (#4922)

---
 app/proxyman/inbound/worker.go | 8 +++++++-
 proxy/shadowsocks/server.go    | 1 +
 proxy/socks/server.go          | 4 +++-
 proxy/trojan/server.go         | 1 +
 proxy/vless/inbound/inbound.go | 5 ++++-
 5 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/app/proxyman/inbound/worker.go b/app/proxyman/inbound/worker.go
index 12e29876..d490777c 100644
--- a/app/proxyman/inbound/worker.go
+++ b/app/proxyman/inbound/worker.go
@@ -161,6 +161,7 @@ type udpConn struct {
 	uplink           stats.Counter
 	downlink         stats.Counter
 	inactive         bool
+	cancel           context.CancelFunc
 }
 
 func (c *udpConn) setInactive() {
@@ -203,6 +204,9 @@ func (c *udpConn) Write(buf []byte) (int, error) {
 }
 
 func (c *udpConn) Close() error {
+	if c.cancel != nil {
+		c.cancel()
+	}
 	common.Must(c.done.Close())
 	common.Must(common.Close(c.writer))
 	return nil
@@ -259,6 +263,7 @@ func (w *udpWorker) getConnection(id connID) (*udpConn, bool) {
 	defer w.Unlock()
 
 	if conn, found := w.activeConn[id]; found && !conn.done.Done() {
+		conn.updateActivity()
 		return conn, true
 	}
 
@@ -306,7 +311,8 @@ func (w *udpWorker) callback(b *buf.Buffer, source net.Destination, originalDest
 		common.Must(w.checker.Start())
 
 		go func() {
-			ctx := w.ctx
+			ctx, cancel := context.WithCancel(w.ctx)
+			conn.cancel = cancel
 			sid := session.NewID()
 			ctx = c.ContextWithID(ctx, sid)
 
diff --git a/proxy/shadowsocks/server.go b/proxy/shadowsocks/server.go
index 356868e4..ec022084 100644
--- a/proxy/shadowsocks/server.go
+++ b/proxy/shadowsocks/server.go
@@ -128,6 +128,7 @@ func (s *Server) handleUDPPayload(ctx context.Context, conn stat.Connection, dis
 
 		conn.Write(data.Bytes())
 	})
+	defer udpServer.RemoveRay()
 
 	inbound := session.InboundFromContext(ctx)
 	var dest *net.Destination
diff --git a/proxy/socks/server.go b/proxy/socks/server.go
index dd6f3953..8a159fad 100644
--- a/proxy/socks/server.go
+++ b/proxy/socks/server.go
@@ -245,13 +245,15 @@ func (s *Server) handleUDPPayload(ctx context.Context, conn stat.Connection, dis
 		udpMessage, err := EncodeUDPPacket(request, payload.Bytes())
 		payload.Release()
 
-		defer udpMessage.Release()
 		if err != nil {
 			errors.LogWarningInner(ctx, err, "failed to write UDP response")
+			return
 		}
+		defer udpMessage.Release()
 
 		conn.Write(udpMessage.Bytes())
 	})
+	defer udpServer.RemoveRay()
 
 	inbound := session.InboundFromContext(ctx)
 	if inbound != nil && inbound.Source.IsValid() {
diff --git a/proxy/trojan/server.go b/proxy/trojan/server.go
index 44662ac3..0ce14408 100644
--- a/proxy/trojan/server.go
+++ b/proxy/trojan/server.go
@@ -259,6 +259,7 @@ func (s *Server) handleUDPPayload(ctx context.Context, clientReader *PacketReade
 			errors.LogWarningInner(ctx, err, "failed to write response")
 		}
 	})
+	defer udpServer.RemoveRay()
 
 	inbound := session.InboundFromContext(ctx)
 	user := inbound.User
diff --git a/proxy/vless/inbound/inbound.go b/proxy/vless/inbound/inbound.go
index 1da2e091..df1f9f3c 100644
--- a/proxy/vless/inbound/inbound.go
+++ b/proxy/vless/inbound/inbound.go
@@ -206,7 +206,10 @@ func (h *Handler) Process(ctx context.Context, network net.Network, connection s
 
 	first := buf.FromBytes(make([]byte, buf.Size))
 	first.Clear()
-	firstLen, _ := first.ReadFrom(connection)
+	firstLen, errR := first.ReadFrom(connection)
+	if errR != nil {
+		return errR
+	}
 	errors.LogInfo(ctx, "firstLen = ", firstLen)
 
 	reader := &buf.BufferedReader{