From 8b631f11b86994c370ed5e35d14d2882ae21f9e7 Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Mon, 15 May 2023 22:45:08 +0800
Subject: [PATCH] chore: better sing's udp api support

---
 common/net/packet/packet.go      |  3 ++
 common/net/packet/packet_sing.go | 74 ++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+)
 create mode 100644 common/net/packet/packet_sing.go

diff --git a/common/net/packet/packet.go b/common/net/packet/packet.go
index d823334e..6329803b 100644
--- a/common/net/packet/packet.go
+++ b/common/net/packet/packet.go
@@ -18,6 +18,9 @@ func NewEnhancePacketConn(pc net.PacketConn) EnhancePacketConn {
 	if enhancePC, isEnhancePC := pc.(EnhancePacketConn); isEnhancePC {
 		return enhancePC
 	}
+	if singPC, isSingPC := pc.(SingPacketConn); isSingPC {
+		return newEnhanceSingPacketConn(singPC)
+	}
 	return &enhancePacketConn{PacketConn: pc}
 }
 
diff --git a/common/net/packet/packet_sing.go b/common/net/packet/packet_sing.go
new file mode 100644
index 00000000..9be1a4a1
--- /dev/null
+++ b/common/net/packet/packet_sing.go
@@ -0,0 +1,74 @@
+package packet
+
+import (
+	"net"
+
+	"github.com/sagernet/sing/common/buf"
+	"github.com/sagernet/sing/common/bufio"
+	M "github.com/sagernet/sing/common/metadata"
+	N "github.com/sagernet/sing/common/network"
+)
+
+type SingPacketConn = N.NetPacketConn
+
+type enhanceSingPacketConn struct {
+	N.NetPacketConn
+	readWaiter N.PacketReadWaiter
+}
+
+func (c *enhanceSingPacketConn) WaitReadFrom() (data []byte, put func(), addr net.Addr, err error) {
+	var buff *buf.Buffer
+	var dest M.Socksaddr
+	newBuffer := func() *buf.Buffer {
+		buff = buf.NewPacket() // do not use stack buffer
+		return buff
+	}
+	if c.readWaiter != nil {
+		c.readWaiter.InitializeReadWaiter(newBuffer)
+		defer c.readWaiter.InitializeReadWaiter(nil)
+		dest, err = c.readWaiter.WaitReadPacket()
+	} else {
+		dest, err = c.NetPacketConn.ReadPacket(newBuffer())
+	}
+	if dest.IsFqdn() {
+		addr = dest
+	} else {
+		addr = dest.UDPAddr()
+	}
+	if err != nil {
+		if buff != nil {
+			buff.Release()
+		}
+		return
+	}
+	if buff == nil {
+		return
+	}
+	if buff.IsEmpty() {
+		buff.Release()
+		return
+	}
+	data = buff.Bytes()
+	put = buff.Release
+	return
+}
+
+func (c *enhanceSingPacketConn) Upstream() any {
+	return c.NetPacketConn
+}
+
+func (c *enhanceSingPacketConn) WriterReplaceable() bool {
+	return true
+}
+
+func (c *enhanceSingPacketConn) ReaderReplaceable() bool {
+	return true
+}
+
+func newEnhanceSingPacketConn(conn N.NetPacketConn) *enhanceSingPacketConn {
+	epc := &enhanceSingPacketConn{NetPacketConn: conn}
+	if readWaiter, isReadWaiter := bufio.CreatePacketReadWaiter(conn); isReadWaiter {
+		epc.readWaiter = readWaiter
+	}
+	return epc
+}