From 0a97c0adb74c50d267918f20b38a205c12f65c21 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Mon, 15 Apr 2024 21:07:24 +0800
Subject: [PATCH] Add EXP_UDPConn interface

---
 common/bufio/bind.go   | 88 +++++++++++++++++++++++++++++++++++++-----
 common/bufio/conn.go   | 32 +++++++--------
 common/network/conn.go | 10 +++++
 3 files changed, 105 insertions(+), 25 deletions(-)

diff --git a/common/bufio/bind.go b/common/bufio/bind.go
index caa5bbd8..a6fd08a7 100644
--- a/common/bufio/bind.go
+++ b/common/bufio/bind.go
@@ -18,10 +18,25 @@ type bindPacketConn struct {
 	addr net.Addr
 }
 
+type bindUDPConn struct {
+	bindPacketConn
+	N.EXP_UDPConn
+}
+
 func NewBindPacketConn(conn net.PacketConn, addr net.Addr) BindPacketConn {
-	return &bindPacketConn{
-		NewPacketConn(conn),
-		addr,
+	if udpConn, isUDPConn := conn.(N.EXP_UDPConn); isUDPConn {
+		return &bindUDPConn{
+			bindPacketConn{
+				NewPacketConn(conn),
+				addr,
+			},
+			udpConn,
+		}
+	} else {
+		return &bindPacketConn{
+			NewPacketConn(conn),
+			addr,
+		}
 	}
 }
 
@@ -61,16 +76,36 @@ type UnbindPacketConn struct {
 }
 
 func NewUnbindPacketConn(conn net.Conn) N.NetPacketConn {
-	return &UnbindPacketConn{
-		NewExtendedConn(conn),
-		M.SocksaddrFromNet(conn.RemoteAddr()),
+	if udpConn, isUDPConn := conn.(N.EXP_UDPConn); isUDPConn {
+		return &UnbindUDPConn{
+			UnbindPacketConn{
+				NewExtendedConn(conn),
+				M.SocksaddrFromNet(conn.RemoteAddr()),
+			},
+			udpConn,
+		}
+	} else {
+		return &UnbindPacketConn{
+			NewExtendedConn(conn),
+			M.SocksaddrFromNet(conn.RemoteAddr()),
+		}
 	}
 }
 
 func NewUnbindPacketConnWithAddr(conn net.Conn, addr M.Socksaddr) N.NetPacketConn {
-	return &UnbindPacketConn{
-		NewExtendedConn(conn),
-		addr,
+	if udpConn, isUDPConn := conn.(N.EXP_UDPConn); isUDPConn {
+		return &UnbindUDPConn{
+			UnbindPacketConn{
+				NewExtendedConn(conn),
+				addr,
+			},
+			udpConn,
+		}
+	} else {
+		return &UnbindPacketConn{
+			NewExtendedConn(conn),
+			addr,
+		}
 	}
 }
 
@@ -110,3 +145,38 @@ func (c *UnbindPacketConn) CreateReadWaiter() (N.PacketReadWaiter, bool) {
 func (c *UnbindPacketConn) Upstream() any {
 	return c.ExtendedConn
 }
+
+var _ N.EXP_UDPConn = (*UnbindUDPConn)(nil)
+
+type UnbindUDPConn struct {
+	UnbindPacketConn
+	N.EXP_UDPConn
+}
+
+func (c *UnbindUDPConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
+	return c.UnbindPacketConn.ReadFrom(p)
+}
+
+func (c *UnbindUDPConn) WriteTo(p []byte, _ net.Addr) (n int, err error) {
+	return c.ExtendedConn.Write(p)
+}
+
+func (c *UnbindUDPConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
+	return c.UnbindPacketConn.ReadPacket(buffer)
+}
+
+func (c *UnbindUDPConn) WritePacket(buffer *buf.Buffer, _ M.Socksaddr) error {
+	return c.ExtendedConn.WriteBuffer(buffer)
+}
+
+func (c *UnbindUDPConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) {
+	n, oobn, flags, addr, err = c.EXP_UDPConn.ReadMsgUDP(b, oob)
+	if err == nil {
+		addr = c.addr.UDPAddr()
+	}
+	return
+}
+
+func (c *UnbindUDPConn) WriteMsgUDP(b, oob []byte, _ *net.UDPAddr) (n, oobn int, err error) {
+	return c.EXP_UDPConn.WriteMsgUDP(b, oob, c.addr.UDPAddr())
+}
diff --git a/common/bufio/conn.go b/common/bufio/conn.go
index 1b925895..a1855c9b 100644
--- a/common/bufio/conn.go
+++ b/common/bufio/conn.go
@@ -148,40 +148,40 @@ type ExtendedConnWrapper struct {
 	writer N.ExtendedWriter
 }
 
-func (w *ExtendedConnWrapper) ReadBuffer(buffer *buf.Buffer) error {
-	return w.reader.ReadBuffer(buffer)
+func (c *ExtendedConnWrapper) ReadBuffer(buffer *buf.Buffer) error {
+	return c.reader.ReadBuffer(buffer)
 }
 
-func (w *ExtendedConnWrapper) WriteBuffer(buffer *buf.Buffer) error {
-	return w.writer.WriteBuffer(buffer)
+func (c *ExtendedConnWrapper) WriteBuffer(buffer *buf.Buffer) error {
+	return c.writer.WriteBuffer(buffer)
 }
 
-func (w *ExtendedConnWrapper) ReadFrom(r io.Reader) (n int64, err error) {
-	return Copy(w.writer, r)
+func (c *ExtendedConnWrapper) ReadFrom(r io.Reader) (n int64, err error) {
+	return Copy(c.writer, r)
 }
 
-func (r *ExtendedConnWrapper) WriteTo(w io.Writer) (n int64, err error) {
-	return Copy(w, r.reader)
+func (c *ExtendedConnWrapper) WriteTo(w io.Writer) (n int64, err error) {
+	return Copy(w, c.reader)
 }
 
-func (w *ExtendedConnWrapper) UpstreamReader() any {
-	return w.reader
+func (c *ExtendedConnWrapper) UpstreamReader() any {
+	return c.reader
 }
 
-func (w *ExtendedConnWrapper) ReaderReplaceable() bool {
+func (c *ExtendedConnWrapper) ReaderReplaceable() bool {
 	return true
 }
 
-func (w *ExtendedConnWrapper) UpstreamWriter() any {
-	return w.writer
+func (c *ExtendedConnWrapper) UpstreamWriter() any {
+	return c.writer
 }
 
-func (w *ExtendedConnWrapper) WriterReplaceable() bool {
+func (c *ExtendedConnWrapper) WriterReplaceable() bool {
 	return true
 }
 
-func (w *ExtendedConnWrapper) Upstream() any {
-	return w.Conn
+func (c *ExtendedConnWrapper) Upstream() any {
+	return c.Conn
 }
 
 func NewExtendedConn(conn net.Conn) N.ExtendedConn {
diff --git a/common/network/conn.go b/common/network/conn.go
index a920ab6e..403321b3 100644
--- a/common/network/conn.go
+++ b/common/network/conn.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"io"
 	"net"
+	"syscall"
 	"time"
 
 	"github.com/sagernet/sing/common"
@@ -85,6 +86,15 @@ type BindPacketConn interface {
 	net.Conn
 }
 
+// EXP_UDPConn is a interface used x/net/ipv4 and quic-go
+type EXP_UDPConn interface {
+	net.PacketConn
+	syscall.Conn
+	SetReadBuffer(bytes int) error
+	ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error)
+	WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error)
+}
+
 type UDPHandler interface {
 	NewPacket(ctx context.Context, conn PacketConn, buffer *buf.Buffer, metadata M.Metadata) error
 }