From e404695a0de07a338161829d87262c76d4827a03 Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Thu, 11 May 2023 15:34:28 +0800
Subject: [PATCH] fix: mux's udp should add write lock

---
 adapter/outbound/singmux.go |  2 +-
 adapter/outbound/vless.go   | 13 +++++++------
 adapter/outbound/vmess.go   | 13 +------------
 common/net/packet.go        | 20 ++++++++++++++++++++
 4 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/adapter/outbound/singmux.go b/adapter/outbound/singmux.go
index acfdfe99..555a0ecb 100644
--- a/adapter/outbound/singmux.go
+++ b/adapter/outbound/singmux.go
@@ -92,7 +92,7 @@ func (s *SingMux) ListenPacketContext(ctx context.Context, metadata *C.Metadata,
 	if pc == nil {
 		return nil, E.New("packetConn is nil")
 	}
-	return newPacketConn(CN.NewRefPacketConn(pc, s), s.ProxyAdapter), nil
+	return newPacketConn(CN.NewRefPacketConn(CN.NewThreadSafePacketConn(pc), s), s.ProxyAdapter), nil
 }
 
 func (s *SingMux) SupportUDP() bool {
diff --git a/adapter/outbound/vless.go b/adapter/outbound/vless.go
index d7db27d5..82ecf927 100644
--- a/adapter/outbound/vless.go
+++ b/adapter/outbound/vless.go
@@ -13,6 +13,7 @@ import (
 	"sync"
 
 	"github.com/Dreamacro/clash/common/convert"
+	N "github.com/Dreamacro/clash/common/net"
 	"github.com/Dreamacro/clash/component/dialer"
 	"github.com/Dreamacro/clash/component/proxydialer"
 	"github.com/Dreamacro/clash/component/resolver"
@@ -372,15 +373,15 @@ func (v *Vless) ListenPacketOnStreamConn(ctx context.Context, c net.Conn, metada
 	}
 
 	if v.option.XUDP {
-		return newPacketConn(&threadSafePacketConn{
-			PacketConn: vmessSing.NewXUDPConn(c, M.SocksaddrFromNet(metadata.UDPAddr())),
-		}, v), nil
+		return newPacketConn(N.NewThreadSafePacketConn(
+			vmessSing.NewXUDPConn(c, M.SocksaddrFromNet(metadata.UDPAddr())),
+		), v), nil
 	} else if v.option.PacketAddr {
-		return newPacketConn(&threadSafePacketConn{
-			PacketConn: packetaddr.NewConn(&vlessPacketConn{
+		return newPacketConn(N.NewThreadSafePacketConn(
+			packetaddr.NewConn(&vlessPacketConn{
 				Conn: c, rAddr: metadata.UDPAddr(),
 			}, M.SocksaddrFromNet(metadata.UDPAddr())),
-		}, v), nil
+		), v), nil
 	}
 	return newPacketConn(&vlessPacketConn{Conn: c, rAddr: metadata.UDPAddr()}, v), nil
 }
diff --git a/adapter/outbound/vmess.go b/adapter/outbound/vmess.go
index 8901f3d5..91e509d7 100644
--- a/adapter/outbound/vmess.go
+++ b/adapter/outbound/vmess.go
@@ -379,7 +379,7 @@ func (v *Vmess) ListenPacketOnStreamConn(ctx context.Context, c net.Conn, metada
 	}
 
 	if pc, ok := c.(net.PacketConn); ok {
-		return newPacketConn(&threadSafePacketConn{PacketConn: pc}, v), nil
+		return newPacketConn(N.NewThreadSafePacketConn(pc), v), nil
 	}
 	return newPacketConn(&vmessPacketConn{Conn: c, rAddr: metadata.UDPAddr()}, v), nil
 }
@@ -489,17 +489,6 @@ func NewVmess(option VmessOption) (*Vmess, error) {
 	return v, nil
 }
 
-type threadSafePacketConn struct {
-	net.PacketConn
-	access sync.Mutex
-}
-
-func (c *threadSafePacketConn) WriteTo(b []byte, addr net.Addr) (int, error) {
-	c.access.Lock()
-	defer c.access.Unlock()
-	return c.PacketConn.WriteTo(b, addr)
-}
-
 type vmessPacketConn struct {
 	net.Conn
 	rAddr  net.Addr
diff --git a/common/net/packet.go b/common/net/packet.go
index 30f1104a..43a2bc9c 100644
--- a/common/net/packet.go
+++ b/common/net/packet.go
@@ -2,6 +2,7 @@ package net
 
 import (
 	"net"
+	"sync"
 
 	"github.com/Dreamacro/clash/common/pool"
 )
@@ -66,3 +67,22 @@ func waitReadFrom(pc net.PacketConn) (data []byte, put func(), addr net.Addr, er
 	}
 	return
 }
+
+type threadSafePacketConn struct {
+	net.PacketConn
+	access sync.Mutex
+}
+
+func (c *threadSafePacketConn) WriteTo(b []byte, addr net.Addr) (int, error) {
+	c.access.Lock()
+	defer c.access.Unlock()
+	return c.PacketConn.WriteTo(b, addr)
+}
+
+func (c *threadSafePacketConn) Upstream() any {
+	return c.PacketConn
+}
+
+func NewThreadSafePacketConn(pc net.PacketConn) net.PacketConn {
+	return &threadSafePacketConn{PacketConn: pc}
+}