diff --git a/protocol/socks/packet.go b/protocol/socks/packet.go index 0560ac72..dd688a68 100644 --- a/protocol/socks/packet.go +++ b/protocol/socks/packet.go @@ -69,7 +69,8 @@ func (c *AssociatePacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error if err != nil { return } - return c.conn.Write(buffer.Bytes()) + n, err = c.conn.Write(buffer.Bytes()) + return max(0, n-(buffer.Len()-len(p))), err } func (c *AssociatePacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) { diff --git a/protocol/socks/packet_test.go b/protocol/socks/packet_test.go new file mode 100644 index 00000000..ce18eaac --- /dev/null +++ b/protocol/socks/packet_test.go @@ -0,0 +1,97 @@ +package socks + +import ( + "bytes" + "errors" + "io" + "net" + "testing" + "time" + + M "github.com/sagernet/sing/common/metadata" +) + +func TestAssociatePacketConnWriteCount(t *testing.T) { + t.Parallel() + for _, destination := range []string{"192.0.2.1:53", "[2001:db8::1]:53", "example.org:53"} { + t.Run(destination, func(t *testing.T) { + t.Parallel() + server, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)}) + if err != nil { + t.Fatal(err) + } + defer server.Close() + client, err := net.DialUDP("udp4", nil, server.LocalAddr().(*net.UDPAddr)) + if err != nil { + t.Fatal(err) + } + defer client.Close() + address := M.ParseSocksaddr(destination) + conn := NewAssociatePacketConn(client, address, nil) + for _, payload := range [][]byte{[]byte("payload"), {}} { + for _, write := range []func([]byte) (int, error){conn.Write, func(p []byte) (int, error) { return conn.WriteTo(p, address) }} { + n, writeErr := write(payload) + if writeErr != nil || n != len(payload) { + t.Errorf("write = %d, %v; want %d, nil", n, writeErr, len(payload)) + } + if err = server.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatal(err) + } + packet := make([]byte, 512) + n, _, err = server.ReadFromUDP(packet) + if err != nil { + t.Fatal(err) + } + if n < 3 || !bytes.Equal(packet[:3], []byte{0, 0, 0}) { + t.Fatalf("invalid SOCKS UDP header: %x", packet[:n]) + } + reader := bytes.NewReader(packet[3:n]) + gotAddress, parseErr := M.SocksaddrSerializer.ReadAddrPort(reader) + if parseErr != nil || gotAddress != address { + t.Fatalf("destination = %v, %v; want %v", gotAddress, parseErr, address) + } + gotPayload, readErr := io.ReadAll(reader) + if readErr != nil || !bytes.Equal(gotPayload, payload) { + t.Fatalf("payload = %q, %v; want %q", gotPayload, readErr, payload) + } + } + } + }) + } +} + +func TestAssociatePacketConnWriteErrorCount(t *testing.T) { + t.Parallel() + writeErr := errors.New("write failed") + for _, test := range []struct { + name string + n int + err error + want int + }{ + {"no bytes", 0, writeErr, 0}, + {"partial header", 5, writeErr, 0}, + {"header only", 10, writeErr, 0}, + {"partial payload", 12, writeErr, 2}, + {"complete with error", 17, writeErr, 7}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + conn := NewAssociatePacketConn(&packetWriteResultConn{n: test.n, err: test.err}, M.ParseSocksaddr("192.0.2.1:53"), nil) + n, err := conn.Write([]byte("payload")) + if n != test.want || !errors.Is(err, test.err) { + t.Fatalf("write = %d, %v; want %d, %v", n, err, test.want, test.err) + } + }) + } +} + +type packetWriteResultConn struct { + net.Conn + n int + err error +} + +func (c *packetWriteResultConn) Write([]byte) (int, error) { + return c.n, c.err +}