Fix UDP async write

transport/trojan/protocol.go

@@ -6,6 +6,7 @@ import (
 	"encoding/hex"
 	"net"
 	"os"
+	"sync"
 
 	"github.com/sagernet/sing/common"
 	"github.com/sagernet/sing/common/buf"
@@ -84,6 +85,7 @@ func (c *ClientConn) Upstream() any {
 
 type ClientPacketConn struct {
 	net.Conn
+	access        sync.Mutex
 	key           [KeyLength]byte
 	headerWritten bool
 }
@@ -105,9 +107,15 @@ func (c *ClientPacketConn) ReadPacket(buffer *buf.Buffer) (M.Socksaddr, error) {
 
 func (c *ClientPacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
 	if !c.headerWritten {
-		err := ClientHandshakePacket(c.Conn, c.key, destination, buffer)
-		c.headerWritten = true
-		return err
+		c.access.Lock()
+		if c.headerWritten {
+			c.access.Unlock()
+		} else {
+			err := ClientHandshakePacket(c.Conn, c.key, destination, buffer)
+			c.headerWritten = true
+			c.access.Unlock()
+			return err
+		}
 	}
 	return WritePacket(c.Conn, buffer, destination)
 }

transport/vless/client.go

@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"io"
 	"net"
+	"sync"
 
 	"github.com/sagernet/sing-vmess"
 	"github.com/sagernet/sing/common"
@@ -190,6 +191,7 @@ func (c *Conn) Upstream() any {
 
 type PacketConn struct {
 	net.Conn
+	access         sync.Mutex
 	key            [16]byte
 	destination    M.Socksaddr
 	flow           string
@@ -218,11 +220,17 @@ func (c *PacketConn) Read(b []byte) (n int, err error) {
 
 func (c *PacketConn) Write(b []byte) (n int, err error) {
 	if !c.requestWritten {
-		err = WritePacketRequest(c.Conn, Request{c.key, vmess.CommandUDP, c.destination, c.flow}, nil)
-		if err == nil {
-			n = len(b)
+		c.access.Lock()
+		if c.requestWritten {
+			c.access.Unlock()
+		} else {
+			err = WritePacketRequest(c.Conn, Request{c.key, vmess.CommandUDP, c.destination, c.flow}, nil)
+			if err == nil {
+				n = len(b)
+			}
+			c.requestWritten = true
+			c.access.Unlock()
 		}
-		c.requestWritten = true
 	}
 	err = binary.Write(c.Conn, binary.BigEndian, uint16(len(b)))
 	if err != nil {
@@ -236,9 +244,15 @@ func (c *PacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) er
 	dataLen := buffer.Len()
 	binary.BigEndian.PutUint16(buffer.ExtendHeader(2), uint16(dataLen))
 	if !c.requestWritten {
-		err := WritePacketRequest(c.Conn, Request{c.key, vmess.CommandUDP, c.destination, c.flow}, buffer.Bytes())
-		c.requestWritten = true
-		return err
+		c.access.Lock()
+		if c.requestWritten {
+			c.access.Unlock()
+		} else {
+			err := WritePacketRequest(c.Conn, Request{c.key, vmess.CommandUDP, c.destination, c.flow}, buffer.Bytes())
+			c.requestWritten = true
+			c.access.Unlock()
+			return err
+		}
 	}
 	return common.Error(c.Conn.Write(buffer.Bytes()))
 }