fix(amneziawg): respect tunnel MTU when writing UDP replies (#6780)

* fix(amneziawg): respect tunnel MTU when writing UDP replies

Send replies through gVisor FindRoute/Route.WritePacket instead of WriteRawPacket so the network endpoints apply the tunnel MTU and fragment oversized IPv4/IPv6 datagrams. Preserve UDP checksums including the zero-complement IPv6 case and validate lengths before narrowing to uint16.

Regression tests inspect and reassemble actual channel-backed TUN output at and above the MTU, verify checksums and exercise concurrent sources.

* docs(review): put the tally before Design and name who can store a trigger

internal/amneziawgnet/udp.go

@@ -42,22 +42,33 @@ func AttachUDPHandler(gstack *stack.Stack, handler UDPHandler) {
 	})
 }
 
-// WriteUDPReply injects a UDP packet into gstack as if it arrived from
-// `from` addressed to `to` -- i.e. a reply travelling back into the tunnel
-// toward the client -- constructed by hand since gVisor exposes no
-// connected-socket-style Write for an address the stack doesn't itself own.
+// WriteUDPReply sends replies through gVisor's network route,
+// which applies the NIC MTU and fragments locally generated IP datagrams.
 func WriteUDPReply(gstack *stack.Stack, from, to netip.AddrPort, payload []byte) error {
+	if !from.Addr().IsValid() || !to.Addr().IsValid() || from.Addr().Zone() != "" || to.Addr().Zone() != "" || from.Addr().Is4() != to.Addr().Is4() {
+		return fmt.Errorf("amneziawgnet: invalid UDP reply addresses: %s -> %s", from, to)
+	}
 	udpLen := header.UDPMinimumSize + len(payload)
+	if udpLen > 65535 {
+		return fmt.Errorf("amneziawgnet: UDP reply too large: %d bytes", udpLen)
+	}
 	srcIP := tcpip.AddrFromSlice(from.Addr().AsSlice())
 	dstIP := tcpip.AddrFromSlice(to.Addr().AsSlice())
 
-	isIPv4 := from.Addr().Is4()
 	ipHdrSize := header.IPv6MinimumSize
 	ipProtocol := header.IPv6ProtocolNumber
-	if isIPv4 {
+	if from.Addr().Is4() {
 		ipHdrSize = header.IPv4MinimumSize
 		ipProtocol = header.IPv4ProtocolNumber
+		if ipHdrSize+udpLen > 65535 {
+			return fmt.Errorf("amneziawgnet: IPv4 reply too large: %d bytes", ipHdrSize+udpLen)
+		}
 	}
+	route, tcpipErr := gstack.FindRoute(1, srcIP, dstIP, ipProtocol, false)
+	if tcpipErr != nil {
+		return fmt.Errorf("amneziawgnet: FindRoute: %s", tcpipErr)
+	}
+	defer route.Release()
 
 	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
 		ReserveHeaderBytes: ipHdrSize + header.UDPMinimumSize,
@@ -72,31 +83,17 @@ func WriteUDPReply(gstack *stack.Stack, from, to netip.AddrPort, payload []byte)
 		Length:  uint16(udpLen),
 	})
 	xsum := header.PseudoHeaderChecksum(header.UDPProtocolNumber, srcIP, dstIP, uint16(udpLen))
-	udpHdr.SetChecksum(^udpHdr.CalculateChecksum(checksum.Checksum(payload, xsum)))
-
-	if isIPv4 {
-		ipHdr := header.IPv4(pkt.NetworkHeader().Push(header.IPv4MinimumSize))
-		ipHdr.Encode(&header.IPv4Fields{
-			TotalLength: uint16(header.IPv4MinimumSize + udpLen),
-			TTL:         64,
-			Protocol:    uint8(header.UDPProtocolNumber),
-			SrcAddr:     srcIP,
-			DstAddr:     dstIP,
-		})
-		ipHdr.SetChecksum(^ipHdr.CalculateChecksum())
-	} else {
-		ipHdr := header.IPv6(pkt.NetworkHeader().Push(header.IPv6MinimumSize))
-		ipHdr.Encode(&header.IPv6Fields{
-			PayloadLength:     uint16(udpLen),
-			TransportProtocol: header.UDPProtocolNumber,
-			HopLimit:          64,
-			SrcAddr:           srcIP,
-			DstAddr:           dstIP,
-		})
+	udpChecksum := ^udpHdr.CalculateChecksum(checksum.Checksum(payload, xsum))
+	if udpChecksum == 0 {
+		udpChecksum = 0xffff
 	}
+	udpHdr.SetChecksum(udpChecksum)
 
-	if tcpipErr := gstack.WriteRawPacket(1, ipProtocol, buffer.MakeWithView(pkt.ToView())); tcpipErr != nil {
-		return fmt.Errorf("amneziawgnet: WriteRawPacket: %s", tcpipErr)
+	if tcpipErr := route.WritePacket(stack.NetworkHeaderParams{
+		Protocol: header.UDPProtocolNumber,
+		TTL:      64,
+	}, pkt); tcpipErr != nil {
+		return fmt.Errorf("amneziawgnet: WritePacket: %s", tcpipErr)
 	}
 	return nil
 }

internal/amneziawgnet/udp_mtu_concurrent_test.go

@@ -0,0 +1,62 @@
+package amneziawgnet
+
+import (
+	"encoding/binary"
+	"net/netip"
+	"sync"
+	"testing"
+	"time"
+
+	"gvisor.dev/gvisor/pkg/tcpip/header"
+)
+
+func TestWriteUDPReplyConcurrentSources(t *testing.T) {
+	tun, gstack, err := createNetTUNWithStack([]netip.Addr{netip.MustParseAddr("10.77.0.1")}, 1280)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer tun.Close()
+	enablePromiscuousRouting(gstack)
+	const replies = 32
+	to := netip.MustParseAddrPort("10.77.0.2:49152")
+	var wg sync.WaitGroup
+	errs := make(chan error, replies)
+	for i := range replies {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			from := netip.AddrPortFrom(netip.MustParseAddr("203.0.113.9"), uint16(10000+i))
+			var payload [2]byte
+			binary.BigEndian.PutUint16(payload[:], uint16(i))
+			errs <- WriteUDPReply(gstack, from, to, payload[:])
+		}()
+	}
+	wg.Wait()
+	close(errs)
+	for err := range errs {
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+	seen := make(map[uint16]bool)
+	for range replies {
+		select {
+		case v := <-tun.(*stackTun).incomingPacket:
+			packet := append([]byte(nil), v.AsSlice()...)
+			v.Release()
+			if len(packet) > 1280 {
+				t.Fatalf("packet length %d", len(packet))
+			}
+			ip := header.IPv4(packet)
+			udp := packet[ip.HeaderLength():]
+			port := binary.BigEndian.Uint16(udp[0:2])
+			index := binary.BigEndian.Uint16(udp[8:10])
+			if index >= replies || port != 10000+index || seen[index] || binary.BigEndian.Uint16(udp[2:4]) != to.Port() {
+				t.Fatalf("duplicate/corrupt concurrent reply: port %d index %d", port, index)
+			}
+			seen[index] = true
+		case <-time.After(time.Second):
+			t.Fatal("concurrent replies missing")
+		}
+	}
+}

internal/amneziawgnet/udp_mtu_test.go

@@ -0,0 +1,97 @@
+package amneziawgnet
+
+import (
+	"bytes"
+	"encoding/binary"
+	"net/netip"
+	"testing"
+	"time"
+
+	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/checksum"
+	"gvisor.dev/gvisor/pkg/tcpip/header"
+)
+
+// TestWriteUDPReplyMTU exercises the actual channel-backed TUN output, not a packet builder.
+func TestWriteUDPReplyMTU(t *testing.T) {
+	for _, tc := range []struct {
+		name      string
+		payload   int
+		fragments bool
+	}{
+		{"small", 32, false}, {"boundary", 1252, false}, {"oversize", 1452, true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			tun, gstack, err := createNetTUNWithStack([]netip.Addr{netip.MustParseAddr("10.77.0.1")}, 1280)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer tun.Close()
+			enablePromiscuousRouting(gstack)
+			from := netip.MustParseAddrPort("203.0.113.9:5353")
+			to := netip.MustParseAddrPort("10.77.0.2:49152")
+			payload := make([]byte, tc.payload)
+			for i := range payload {
+				payload[i] = byte(i*37 + 11)
+			}
+			if err := WriteUDPReply(gstack, from, to, payload); err != nil {
+				t.Fatal(err)
+			}
+			st := tun.(*stackTun)
+			var data []byte
+			var received int
+			var id uint16
+			for {
+				select {
+				case v := <-st.incomingPacket:
+					p := v.AsSlice()
+					packet := append([]byte(nil), p...)
+					v.Release()
+					if len(packet) > 1280 {
+						t.Fatalf("TUN packet length %d exceeds MTU 1280", len(packet))
+					}
+					ip := header.IPv4(packet)
+					if ip.SourceAddress().String() != from.Addr().String() || ip.DestinationAddress().String() != to.Addr().String() {
+						t.Fatalf("IP addresses %s -> %s", ip.SourceAddress(), ip.DestinationAddress())
+					}
+					if ip.TotalLength() != uint16(len(packet)) || checksum.Checksum(packet[:ip.HeaderLength()], 0) != 0xffff {
+						t.Fatalf("bad IPv4 header, len %d", len(packet))
+					}
+					if received == 0 {
+						id = ip.ID()
+					} else if id != ip.ID() {
+						t.Fatalf("fragment IDs differ: %d != %d", id, ip.ID())
+					}
+					offset := int(ip.FragmentOffset())
+					part := packet[ip.HeaderLength():]
+					if offset != received {
+						t.Fatalf("fragment offset %d, expected %d", offset, received)
+					}
+					data = append(data, part...)
+					received += len(part)
+					if ip.Flags()&header.IPv4FlagMoreFragments == 0 {
+						if tc.fragments && offset == 0 {
+							t.Fatal("oversize reply was not fragmented")
+						}
+						goto reassembled
+					}
+				case <-time.After(time.Second):
+					t.Fatal("timed out waiting for TUN packets")
+				}
+			}
+		reassembled:
+			if len(data) != header.UDPMinimumSize+len(payload) {
+				t.Fatalf("UDP bytes %d, want %d", len(data), 8+len(payload))
+			}
+			if binary.BigEndian.Uint16(data[:2]) != from.Port() || binary.BigEndian.Uint16(data[2:4]) != to.Port() || binary.BigEndian.Uint16(data[4:6]) != uint16(len(data)) {
+				t.Fatalf("UDP ports/length incorrect: %x", data[:8])
+			}
+			if !bytes.Equal(data[8:], payload) {
+				t.Fatal("UDP payload changed")
+			}
+			if checksum.Checksum(data, header.PseudoHeaderChecksum(header.UDPProtocolNumber, tcpip.AddrFromSlice(from.Addr().AsSlice()), tcpip.AddrFromSlice(to.Addr().AsSlice()), uint16(len(data)))) != 0xffff {
+				t.Fatal("UDP checksum invalid")
+			}
+		})
+	}
+}

internal/amneziawgnet/udp_mtu_v6_test.go

@@ -0,0 +1,175 @@
+package amneziawgnet
+
+import (
+	"bytes"
+	"encoding/binary"
+	"net/netip"
+	"strings"
+	"testing"
+	"time"
+
+	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/checksum"
+	"gvisor.dev/gvisor/pkg/tcpip/header"
+)
+
+func TestWriteUDPReplyIPv6MTU(t *testing.T) {
+	for _, tc := range []struct {
+		name       string
+		size       int
+		fragmented bool
+	}{
+		{"small", 80, false}, {"boundary", 1232, false}, {"fragmented", 1452, true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			tun, gstack, err := createNetTUNWithStack([]netip.Addr{netip.MustParseAddr("fd00::1")}, 1280)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer tun.Close()
+			enablePromiscuousRouting(gstack)
+			from := netip.MustParseAddrPort("[2001:db8::abcd]:5353")
+			to := netip.MustParseAddrPort("[fd00::2]:49152")
+			payload := make([]byte, tc.size)
+			for i := range payload {
+				payload[i] = byte(i*51 + 7)
+			}
+			if err := WriteUDPReply(gstack, from, to, payload); err != nil {
+				t.Fatal(err)
+			}
+			var data []byte
+			var id uint32
+			count := 0
+			for {
+				select {
+				case v := <-tun.(*stackTun).incomingPacket:
+					packet := append([]byte(nil), v.AsSlice()...)
+					v.Release()
+					if len(packet) > 1280 {
+						t.Fatalf("IPv6 TUN packet length %d exceeds MTU", len(packet))
+					}
+					ip := header.IPv6(packet)
+					if ip.SourceAddress().String() != from.Addr().String() || ip.DestinationAddress().String() != to.Addr().String() || int(ip.PayloadLength())+40 != len(packet) {
+						t.Fatalf("IPv6 addresses/length invalid: %s -> %s, %d", ip.SourceAddress(), ip.DestinationAddress(), ip.PayloadLength())
+					}
+					var offset int
+					more := false
+					part := packet[40:]
+					if tc.fragmented {
+						if ip.TransportProtocol() != header.IPv6FragmentHeader {
+							t.Fatalf("IPv6 NextHeader %d, want fragment", ip.TransportProtocol())
+						}
+						frag := header.IPv6Fragment(part)
+						if frag.NextHeader() != uint8(header.UDPProtocolNumber) {
+							t.Fatalf("fragment next header %d", frag.NextHeader())
+						}
+						offset = int(frag.FragmentOffset()) * 8
+						more = frag.More()
+						if count == 0 {
+							id = frag.ID()
+						} else if frag.ID() != id {
+							t.Fatalf("fragment ID %d != %d", frag.ID(), id)
+						}
+						part = part[8:]
+					} else if ip.TransportProtocol() != header.UDPProtocolNumber {
+						t.Fatalf("IPv6 NextHeader %d", ip.TransportProtocol())
+					}
+					if offset != len(data) {
+						t.Fatalf("fragment offset %d, want %d", offset, len(data))
+					}
+					data = append(data, part...)
+					count++
+					if !more {
+						goto assembled
+					}
+				case <-time.After(time.Second):
+					t.Fatal("IPv6 reply timed out")
+				}
+			}
+		assembled:
+			if tc.fragmented && count < 2 {
+				t.Fatal("no IPv6 fragments")
+			}
+			if len(data) != len(payload)+8 || binary.BigEndian.Uint16(data[:2]) != from.Port() || binary.BigEndian.Uint16(data[2:4]) != to.Port() || binary.BigEndian.Uint16(data[4:6]) != uint16(len(data)) || !bytes.Equal(data[8:], payload) {
+				t.Fatal("reassembled IPv6 UDP bytes/ports/length differ")
+			}
+			if checksum.Checksum(data, header.PseudoHeaderChecksum(header.UDPProtocolNumber, tcpip.AddrFromSlice(from.Addr().AsSlice()), tcpip.AddrFromSlice(to.Addr().AsSlice()), uint16(len(data)))) != 0xffff {
+				t.Fatal("IPv6 UDP checksum invalid")
+			}
+		})
+	}
+}
+
+func TestWriteUDPReplyRejectsInvalidDatagrams(t *testing.T) {
+	tun, gstack, err := createNetTUNWithStack([]netip.Addr{netip.MustParseAddr("10.77.0.1"), netip.MustParseAddr("fd00::1")}, 1280)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer tun.Close()
+	enablePromiscuousRouting(gstack)
+	for _, tc := range []struct {
+		name, from, to string
+		size           int
+		want           string
+	}{
+		{"mixed", "203.0.113.9:53", "[fd00::2]:4000", 1, "invalid UDP reply addresses"},
+		{"zone", "[fe80::1%eth0]:53", "[fd00::2]:4000", 1, "invalid UDP reply addresses"},
+		{"v4 length", "203.0.113.9:53", "10.77.0.2:4000", 65508, "IPv4 reply too large"},
+		{"v6 length", "[2001:db8::1]:53", "[fd00::2]:4000", 65528, "UDP reply too large"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			err := WriteUDPReply(gstack, netip.MustParseAddrPort(tc.from), netip.MustParseAddrPort(tc.to), make([]byte, tc.size))
+			if err == nil || !strings.Contains(err.Error(), tc.want) {
+				t.Fatalf("error = %v, want %q", err, tc.want)
+			}
+			select {
+			case v := <-tun.(*stackTun).incomingPacket:
+				v.Release()
+				t.Fatal("invalid reply emitted packet")
+			default:
+			}
+		})
+	}
+}
+
+func TestWriteUDPReplyIPv6ZeroChecksumIsEncodedAsOnes(t *testing.T) {
+	tun, gstack, err := createNetTUNWithStack([]netip.Addr{netip.MustParseAddr("fd00::1")}, 1280)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer tun.Close()
+	enablePromiscuousRouting(gstack)
+	from := netip.MustParseAddrPort("[2001:db8::abcd]:5353")
+	to := netip.MustParseAddrPort("[fd00::2]:49152")
+	src := tcpip.AddrFromSlice(from.Addr().AsSlice())
+	dst := tcpip.AddrFromSlice(to.Addr().AsSlice())
+	udp := make([]byte, 10)
+	binary.BigEndian.PutUint16(udp[0:2], from.Port())
+	binary.BigEndian.PutUint16(udp[2:4], to.Port())
+	binary.BigEndian.PutUint16(udp[4:6], uint16(len(udp)))
+	pseudo := header.PseudoHeaderChecksum(header.UDPProtocolNumber, src, dst, uint16(len(udp)))
+	var found bool
+	for n := 0; n <= 65535; n++ {
+		binary.BigEndian.PutUint16(udp[8:], uint16(n))
+		if checksum.Checksum(udp, pseudo) == 0xffff {
+			found = true
+			break
+		}
+	}
+	if !found {
+		t.Fatal("no zero-complement checksum payload found")
+	}
+	if err := WriteUDPReply(gstack, from, to, udp[8:]); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case v := <-tun.(*stackTun).incomingPacket:
+		packet := append([]byte(nil), v.AsSlice()...)
+		v.Release()
+		if got := binary.BigEndian.Uint16(packet[46:48]); got != 0xffff {
+			t.Fatalf("IPv6 UDP checksum = %#04x, want 0xffff (zero is invalid)", got)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("reply timed out")
+	}
+}