socks: report payload bytes from associate packet writes

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) {

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
+}