Implement read waiter for UDP

SagerNet · Dec 6, 2023 · 6d14d0a · 6d14d0a
1 parent 13614c0
commit 6d14d0a
Show file tree

Hide file tree

Showing 4 changed files with 78 additions and 5 deletions.
diff --git a/client.go b/client.go
@@ -86,7 +86,7 @@ func (c *Client) DialContext(ctx context.Context, network string, destination M.
 		if err != nil {
 			return nil, err
 		}
-		return bufio.NewUnbindPacketConn(&clientPacketConn{ExtendedConn: bufio.NewExtendedConn(stream), destination: destination}), nil
+		return &clientPacketConn{ExtendedConn: bufio.NewExtendedConn(stream), destination: destination}, nil
 	default:
 		return nil, E.Extend(N.ErrUnknownNetwork, network)
 	}

diff --git a/client_conn.go b/client_conn.go
@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"io"
 	"net"
+	"os"
 	"sync"
 
 	"github.com/sagernet/sing/common"
@@ -93,12 +94,18 @@ func (c *clientConn) Upstream() any {
 	return c.Conn
 }
 
+var (
+	_ N.NetPacketConn    = (*clientPacketConn)(nil)
+	_ N.PacketReadWaiter = (*clientPacketConn)(nil)
+)
+
 type clientPacketConn struct {
 	N.ExtendedConn
 	access         sync.Mutex
 	destination    M.Socksaddr
 	requestWritten bool
 	responseRead   bool
+	newBuffer      func() *buf.Buffer
 }
 
 func (c *clientPacketConn) NeedHandshake() bool {
@@ -214,6 +221,35 @@ func (c *clientPacketConn) WriteBuffer(buffer *buf.Buffer) error {
 	return c.ExtendedConn.WriteBuffer(buffer)
 }
 
+func (c *clientPacketConn) InitializeReadWaiter(newBuffer func() *buf.Buffer) {
+	c.newBuffer = newBuffer
+}
+
+func (c *clientPacketConn) WaitReadPacket() (buffer *buf.Buffer, destination M.Socksaddr, err error) {
+	if c.newBuffer == nil {
+		return nil, M.Socksaddr{}, os.ErrInvalid
+	}
+	if !c.responseRead {
+		err = c.readResponse()
+		if err != nil {
+			return
+		}
+		c.responseRead = true
+	}
+	var length uint16
+	err = binary.Read(c.ExtendedConn, binary.BigEndian, &length)
+	if err != nil {
+		return
+	}
+	buffer = c.newBuffer()
+	_, err = buffer.ReadFullFrom(c.ExtendedConn, int(length))
+	if err != nil {
+		buffer.Release()
+		return nil, M.Socksaddr{}, err
+	}
+	return
+}
+
 func (c *clientPacketConn) FrontHeadroom() int {
 	return 2
 }
@@ -280,14 +316,18 @@ func (c *clientPacketConn) Upstream() any {
 	return c.ExtendedConn
 }
 
-var _ N.NetPacketConn = (*clientPacketAddrConn)(nil)
+var (
+	_ N.NetPacketConn    = (*clientPacketAddrConn)(nil)
+	_ N.PacketReadWaiter = (*clientPacketAddrConn)(nil)
+)
 
 type clientPacketAddrConn struct {
 	N.ExtendedConn
 	access         sync.Mutex
 	destination    M.Socksaddr
 	requestWritten bool
 	responseRead   bool
+	newBuffer      func() *buf.Buffer
 }
 
 func (c *clientPacketAddrConn) NeedHandshake() bool {
@@ -431,6 +471,39 @@ func (c *clientPacketAddrConn) WritePacket(buffer *buf.Buffer, destination M.Soc
 	return c.ExtendedConn.WriteBuffer(buffer)
 }
 
+func (c *clientPacketAddrConn) InitializeReadWaiter(newBuffer func() *buf.Buffer) {
+	c.newBuffer = newBuffer
+}
+
+func (c *clientPacketAddrConn) WaitReadPacket() (buffer *buf.Buffer, destination M.Socksaddr, err error) {
+	if c.newBuffer == nil {
+		return nil, M.Socksaddr{}, os.ErrInvalid
+	}
+	if !c.responseRead {
+		err = c.readResponse()
+		if err != nil {
+			return
+		}
+		c.responseRead = true
+	}
+	destination, err = M.SocksaddrSerializer.ReadAddrPort(c.ExtendedConn)
+	if err != nil {
+		return
+	}
+	var length uint16
+	err = binary.Read(c.ExtendedConn, binary.BigEndian, &length)
+	if err != nil {
+		return
+	}
+	buffer = c.newBuffer()
+	_, err = buffer.ReadFullFrom(c.ExtendedConn, int(length))
+	if err != nil {
+		buffer.Release()
+		return nil, M.Socksaddr{}, err
+	}
+	return
+}
+
 func (c *clientPacketAddrConn) LocalAddr() net.Addr {
 	return c.ExtendedConn.LocalAddr()
 }

diff --git a/go.mod b/go.mod
@@ -4,7 +4,7 @@ go 1.18
 
 require (
 	github.com/hashicorp/yamux v0.1.1
-	github.com/sagernet/sing v0.2.18
+	github.com/sagernet/sing v0.2.19-0.20231206123939-86c131fc0547
 	github.com/sagernet/smux v0.0.0-20230312102458-337ec2a5af37
 	golang.org/x/net v0.19.0
 	golang.org/x/sys v0.15.0

diff --git a/go.sum b/go.sum
@@ -1,8 +1,8 @@
 github.com/hashicorp/yamux v0.1.1 h1:yrQxtgseBDrq9Y652vSRDvsKCJKOUD+GzTS4Y0Y8pvE=
 github.com/hashicorp/yamux v0.1.1/go.mod h1:CtWFDAQgb7dxtzFs4tWbplKIe2jSi3+5vKbgIO0SLnQ=
 github.com/sagernet/sing v0.1.8/go.mod h1:jt1w2u7lJQFFSGLiRrRIs5YWmx4kAPfWuOejuDW9qMk=
-github.com/sagernet/sing v0.2.18 h1:2Ce4dl0pkWft+4914NGXPb8OiQpgA8UHQ9xFOmgvKuY=
-github.com/sagernet/sing v0.2.18/go.mod h1:OL6k2F0vHmEzXz2KW19qQzu172FDgSbUSODylighuVo=
+github.com/sagernet/sing v0.2.19-0.20231206123939-86c131fc0547 h1:Mco2TffcdPMlmzbcGmbKd6Vm6dEBJH2VgjO6H1XfZzY=
+github.com/sagernet/sing v0.2.19-0.20231206123939-86c131fc0547/go.mod h1:Ce5LNojQOgOiWhiD8pPD6E9H7e2KgtOe3Zxx4Ou5u80=
 github.com/sagernet/smux v0.0.0-20230312102458-337ec2a5af37 h1:HuE6xSwco/Xed8ajZ+coeYLmioq0Qp1/Z2zczFaV8as=
 github.com/sagernet/smux v0.0.0-20230312102458-337ec2a5af37/go.mod h1:3skNSftZDJWTGVtVaM2jfbce8qHnmH/AGDRe62iNOg0=
 golang.org/x/net v0.19.0 h1:zTwKpTd2XuCqf8huc7Fo2iSy+4RHPd10s4KzeTnVr1c=