plumb ECN through via relays

udp/udp_ecn_outer_linux_test.go

@@ -3,8 +3,11 @@
 package udp
 
 import (
+	"encoding/binary"
 	"net/netip"
 	"testing"
+
+	"golang.org/x/sys/unix"
 )
 
 // TestPlanRunBreaksOnECNChange confirms that two same-destination, same-size
@@ -59,3 +62,130 @@ func TestPlanRunBreaksOnECNChange(t *testing.T) {
 		}
 	})
 }
+
+// ecnReceiver is a raw UDP socket with IP_RECVTOS / IPV6_RECVTCLASS enabled,
+// used to observe the outer ECN codepoint WriteTo stamps on the wire.
+type ecnReceiver struct {
+	fd   int
+	addr netip.AddrPort
+}
+
+func newEcnReceiver(t *testing.T, v6 bool) *ecnReceiver {
+	t.Helper()
+	family := unix.AF_INET
+	if v6 {
+		family = unix.AF_INET6
+	}
+	fd, err := unix.Socket(family, unix.SOCK_DGRAM, 0)
+	if err != nil {
+		t.Fatalf("socket: %v", err)
+	}
+	t.Cleanup(func() { unix.Close(fd) })
+
+	var bindAddr netip.Addr
+	if v6 {
+		if err = unix.SetsockoptInt(fd, unix.IPPROTO_IPV6, unix.IPV6_RECVTCLASS, 1); err != nil {
+			t.Fatalf("IPV6_RECVTCLASS: %v", err)
+		}
+		if err = unix.Bind(fd, &unix.SockaddrInet6{Addr: [16]byte{15: 1}}); err != nil {
+			t.Fatalf("bind ::1: %v", err)
+		}
+		bindAddr = netip.MustParseAddr("::1")
+	} else {
+		if err = unix.SetsockoptInt(fd, unix.IPPROTO_IP, unix.IP_RECVTOS, 1); err != nil {
+			t.Fatalf("IP_RECVTOS: %v", err)
+		}
+		if err = unix.Bind(fd, &unix.SockaddrInet4{Addr: [4]byte{127, 0, 0, 1}}); err != nil {
+			t.Fatalf("bind 127.0.0.1: %v", err)
+		}
+		bindAddr = netip.MustParseAddr("127.0.0.1")
+	}
+	tv := unix.Timeval{Sec: 5}
+	if err = unix.SetsockoptTimeval(fd, unix.SOL_SOCKET, unix.SO_RCVTIMEO, &tv); err != nil {
+		t.Fatalf("SO_RCVTIMEO: %v", err)
+	}
+
+	sa, err := unix.Getsockname(fd)
+	if err != nil {
+		t.Fatalf("getsockname: %v", err)
+	}
+	var port int
+	switch v := sa.(type) {
+	case *unix.SockaddrInet4:
+		port = v.Port
+	case *unix.SockaddrInet6:
+		port = v.Port
+	default:
+		t.Fatalf("unexpected sockaddr %T", sa)
+	}
+	return &ecnReceiver{fd: fd, addr: netip.AddrPortFrom(bindAddr, uint16(port))}
+}
+
+// recvECN receives one datagram and returns the 2-bit ECN codepoint from its
+// TOS / TCLASS cmsg.
+func (r *ecnReceiver) recvECN(t *testing.T) byte {
+	t.Helper()
+	buf := make([]byte, 128)
+	oob := make([]byte, 128)
+	_, oobn, _, _, err := unix.Recvmsg(r.fd, buf, oob, 0)
+	if err != nil {
+		t.Fatalf("recvmsg: %v", err)
+	}
+	cmsgs, err := unix.ParseSocketControlMessage(oob[:oobn])
+	if err != nil {
+		t.Fatalf("parse cmsg: %v", err)
+	}
+	for _, m := range cmsgs {
+		switch {
+		case m.Header.Level == unix.IPPROTO_IP && m.Header.Type == unix.IP_TOS:
+			return m.Data[0] & 0x03
+		case m.Header.Level == unix.IPPROTO_IPV6 && m.Header.Type == unix.IPV6_TCLASS:
+			return byte(binary.NativeEndian.Uint32(m.Data)) & 0x03
+		}
+	}
+	t.Fatal("no TOS/TCLASS cmsg received")
+	return 0
+}
+
+// TestWriteToStampsOuterECN sends single packets through StdConn.WriteTo and
+// asserts the requested ECN codepoint lands on the outer IP header, for a
+// v4 socket, a v6 socket, and the dual-stack case where a v4-mapped
+// destination must be stamped via IP_TOS rather than IPV6_TCLASS.
+func TestWriteToStampsOuterECN(t *testing.T) {
+	cases := []struct {
+		name    string
+		bind    string
+		recvV6  bool
+		sendECN byte
+	}{
+		{"v4_socket_to_v4", "127.0.0.1", false, 0x03},
+		{"v6_socket_to_v6", "::1", true, 0x01},
+		{"dualstack_v6_socket_to_v4_mapped", "::", false, 0x02},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			c, err := NewListener(testLogger(), netip.MustParseAddr(tc.bind), 0, false, 8)
+			if err != nil {
+				t.Fatalf("NewListener: %v", err)
+			}
+			defer c.Close()
+			rx := newEcnReceiver(t, tc.recvV6)
+
+			if err = c.WriteTo([]byte("ecn"), rx.addr, tc.sendECN); err != nil {
+				t.Fatalf("WriteTo(ecn=%#02x): %v", tc.sendECN, err)
+			}
+			if got := rx.recvECN(t); got != tc.sendECN {
+				t.Errorf("outer ECN = %#02x, want %#02x", got, tc.sendECN)
+			}
+
+			// The zero codepoint sends no TOS cmsg and must arrive Not-ECT
+			// (the socket-default TOS byte).
+			if err = c.WriteTo([]byte("ecn"), rx.addr, 0); err != nil {
+				t.Fatalf("WriteTo(ecn=0): %v", err)
+			}
+			if got := rx.recvECN(t); got != 0 {
+				t.Errorf("outer ECN = %#02x, want 0 (Not-ECT)", got)
+			}
+		})
+	}
+}