Improve bufio

common/bufio/conn.go

@@ -299,8 +299,10 @@ func CopyPacketConnTimeout(ctx context.Context, conn N.PacketConn, dest N.Packet
 	})
 }
 
-func NewPacketConn(conn net.PacketConn) N.PacketConn {
-	if udpConn, ok := conn.(*net.UDPConn); ok {
+func NewPacketConn(conn net.PacketConn) N.NetPacketConn {
+	if packetConn, ok := conn.(N.NetPacketConn); ok {
+		return packetConn
+	} else if udpConn, ok := conn.(*net.UDPConn); ok {
 		return &ExtendedUDPConn{udpConn}
 	} else {
 		return &ExtendedPacketConn{udpConn}
@@ -380,37 +382,44 @@ func (c *BindPacketConn) Upstream() any {
 }
 
 type UnbindPacketConn struct {
-	net.Conn
+	N.ExtendedConn
+	Addr M.Socksaddr
 }
 
 func (c *UnbindPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
-	n, err = c.Conn.Read(p)
+	n, err = c.ExtendedConn.Read(p)
 	if err == nil {
-		addr = c.RemoteAddr()
+		addr = c.Addr.UDPAddr()
 	}
 	return
 }
 
 func (c *UnbindPacketConn) WriteTo(p []byte, _ net.Addr) (n int, err error) {
-	return c.Write(p)
+	return c.ExtendedConn.Write(p)
 }
 
 func (c *UnbindPacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
-	_, err = buffer.ReadFrom(c.Conn)
+	err = c.ExtendedConn.ReadBuffer(buffer)
 	if err != nil {
 		return
 	}
-	destination = M.SocksaddrFromNet(c.RemoteAddr())
+	destination = c.Addr
 	return
 }
 
-func (c *UnbindPacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
-	defer buffer.Release()
-	return common.Error(c.Conn.Write(buffer.Bytes()))
+func (c *UnbindPacketConn) WritePacket(buffer *buf.Buffer, _ M.Socksaddr) error {
+	return c.ExtendedConn.WriteBuffer(buffer)
 }
 
 func (c *UnbindPacketConn) Upstream() any {
-	return c.Conn
+	return c.ExtendedConn
+}
+
+func NewUnbindPacketConn(conn net.Conn) N.NetPacketConn {
+	return &UnbindPacketConn{
+		NewExtendedConn(conn),
+		M.SocksaddrFromNet(conn.RemoteAddr()),
+	}
 }
 
 type ExtendedReaderWrapper struct {

common/bufio/limit.go

@@ -0,0 +1,50 @@
+package bufio
+
+import (
+	"io"
+
+	"github.com/sagernet/sing/common"
+	"github.com/sagernet/sing/common/buf"
+	N "github.com/sagernet/sing/common/network"
+)
+
+type LimitedWriter struct {
+	upstream       N.ExtendedWriter
+	maxChunkLength int
+}
+
+func NewLimitedWriter(writer io.Writer, maxChunkLength int) *LimitedWriter {
+	return &LimitedWriter{
+		upstream:       NewExtendedWriter(writer),
+		maxChunkLength: maxChunkLength,
+	}
+}
+
+func (w *LimitedWriter) Write(p []byte) (n int, err error) {
+	for pLen := len(p); pLen > 0; {
+		var data []byte
+		if pLen > w.maxChunkLength {
+			data = p[:w.maxChunkLength]
+			p = p[w.maxChunkLength:]
+			pLen -= w.maxChunkLength
+		} else {
+			data = p
+			pLen = 0
+		}
+		var writeN int
+		writeN, err = w.upstream.Write(data)
+		if err != nil {
+			return
+		}
+		n += writeN
+	}
+	return
+}
+
+func (w *LimitedWriter) WriteBuffer(buffer *buf.Buffer) error {
+	if buffer.Len() <= w.maxChunkLength {
+		return w.upstream.WriteBuffer(buffer)
+	}
+	defer buffer.Release()
+	return common.Error(w.Write(buffer.Bytes()))
+}

common/metadata/addr.go

@@ -116,6 +116,9 @@ func SocksaddrFromNetIP(ap netip.AddrPort) Socksaddr {
 }
 
 func SocksaddrFromNet(ap net.Addr) Socksaddr {
+	if ap == nil {
+		return Socksaddr{}
+	}
 	if socksAddr, ok := ap.(Socksaddr); ok {
 		return socksAddr
 	}

common/network/conn.go

@@ -60,6 +60,11 @@ type NetPacketConn interface {
 	PacketConn
 }
 
+type BindPacketConn interface {
+	NetPacketConn
+	net.Conn
+}
+
 type UDPHandler interface {
 	NewPacket(ctx context.Context, conn PacketConn, buffer *buf.Buffer, metadata M.Metadata) error
 }

go.mod

@@ -2,4 +2,4 @@ module github.com/sagernet/sing
 
 go 1.18
 
-require golang.org/x/sys v0.0.0-20220608164250-635b8c9b7f68
+require golang.org/x/sys v0.0.0-20220610221304-9f5ed59c137d

go.sum

@@ -1,2 +1,2 @@
-golang.org/x/sys v0.0.0-20220608164250-635b8c9b7f68 h1:z8Hj/bl9cOV2grsOpEaQFUaly0JWN3i97mo3jXKJNp0=
-golang.org/x/sys v0.0.0-20220608164250-635b8c9b7f68/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220610221304-9f5ed59c137d h1:Zu/JngovGLVi6t2J3nmAf3AoTDwuzw85YZ3b9o4yU7s=
+golang.org/x/sys v0.0.0-20220610221304-9f5ed59c137d/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

protocol/socks/packet.go

@@ -33,7 +33,7 @@ func NewAssociatePacketConn(conn net.PacketConn, remoteAddr M.Socksaddr, underly
 
 func NewAssociateConn(conn net.Conn, remoteAddr M.Socksaddr, underlying net.Conn) *AssociatePacketConn {
 	return &AssociatePacketConn{
-		PacketConn: &bufio.UnbindPacketConn{Conn: conn},
+		PacketConn: bufio.NewUnbindPacketConn(conn),
 		remoteAddr: remoteAddr,
 		underlying: underlying,
 	}