handler_test.go

package icx_test

import (
	"encoding/binary"
	"log/slog"
	"net"
	"net/netip"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/pkg/tcpip/header"

	"github.com/apoxy-dev/icx"
	"github.com/apoxy-dev/icx/psp"
	"github.com/apoxy-dev/icx/udp"
)

func TestHandler(t *testing.T) {
	if testing.Verbose() {
		slog.SetLogLoggerLevel(slog.LevelDebug)
	}

	localAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 1).To4()),
		Port: 1234,
	}

	peerAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 2).To4()),
		Port: 4321,
	}

	virtMAC := tcpip.GetRandMacAddr()

	var key [16]byte
	copy(key[:], []byte("0123456789abcdef"))

	h, err := icx.NewHandler(icx.WithLocalAddr(localAddr),
		icx.WithVirtMAC(virtMAC),
		icx.WithSourcePortHashing(),
		icx.WithKeepAliveInterval(500*time.Millisecond),
	)
	require.NoError(t, err)

	wildcardPrefix := netip.MustParsePrefix("0.0.0.0/0")

	err = h.AddVirtualNetwork(0x12345, peerAddr, []icx.Route{{Src: wildcardPrefix, Dst: wildcardPrefix}})
	require.NoError(t, err)

	err = h.InstallKeysForTest(0x12345, 1, key, key, time.Now().Add(time.Hour))
	require.NoError(t, err)

	virtFrame := makeIPv4UDPEthernetFrame(virtMAC)

	phyFrame := make([]byte, 1500)
	frameLen, loopback := h.VirtToPhy(virtFrame, phyFrame)
	require.NotZero(t, frameLen)
	require.False(t, loopback)

	receivedFrame := make([]byte, 1500)
	decodedLen := h.PhyToVirt(phyFrame[:frameLen], receivedFrame)
	require.NotZero(t, decodedLen)

	receivedFrame = receivedFrame[:decodedLen]

	require.Equal(t, virtFrame[header.EthernetMinimumSize:], receivedFrame[header.EthernetMinimumSize:])

	eth := header.Ethernet(receivedFrame)

	require.Equal(t, virtMAC, eth.DestinationAddress())

	frameLen = h.ToPhy(phyFrame)
	require.NotZero(t, frameLen)

	require.NoError(t, h.RemoveVirtualNetwork(0x12345))

	frameLen, loopback = h.VirtToPhy(virtFrame, phyFrame)
	require.Zero(t, frameLen)
	require.False(t, loopback)

	frameLen = h.ToPhy(phyFrame)
	require.Zero(t, frameLen)
}

func TestHandler_Layer3(t *testing.T) {
	if testing.Verbose() {
		slog.SetLogLoggerLevel(slog.LevelDebug)
	}

	localAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 1).To4()),
		Port: 1234,
	}

	peerAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 2).To4()),
		Port: 4321,
	}

	var key [16]byte
	copy(key[:], []byte("0123456789abcdef"))

	// Create handler with layer3 mode enabled
	h, err := icx.NewHandler(icx.WithLocalAddr(localAddr),
		icx.WithLayer3VirtFrames(),
		icx.WithSourcePortHashing())
	require.NoError(t, err)

	privatePrefix := netip.MustParsePrefix("192.168.1.0/24")
	err = h.AddVirtualNetwork(0x12345, peerAddr, []icx.Route{{Src: privatePrefix, Dst: privatePrefix}})
	require.NoError(t, err)

	err = h.InstallKeysForTest(0x12345, 1, key, key, time.Now().Add(time.Hour))
	require.NoError(t, err)

	ipPacket := makeIPv4UDPPacket()

	phyFrame := make([]byte, 1500)
	frameLen, loopback := h.VirtToPhy(ipPacket, phyFrame)
	require.NotZero(t, frameLen)
	require.False(t, loopback)

	decoded := make([]byte, 1500)
	decodedLen := h.PhyToVirt(phyFrame[:frameLen], decoded)
	require.NotZero(t, decodedLen)

	require.Equal(t, ipPacket, decoded[:decodedLen])

	require.NoError(t, h.RemoveVirtualNetwork(0x12345))
}

func TestHandler_Layer3_IPv6(t *testing.T) {
	localAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 1).To4()),
		Port: 1234,
	}
	peerAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 2).To4()),
		Port: 4321,
	}

	var key [16]byte
	copy(key[:], []byte("0123456789abcdef"))

	h, err := icx.NewHandler(
		icx.WithLocalAddr(localAddr),
		icx.WithLayer3VirtFrames(),
	)
	require.NoError(t, err)

	// Prefix contains src 2001:db8::1
	privatePrefix := netip.MustParsePrefix("2001:db8::/64")
	require.NoError(t, h.AddVirtualNetwork(0x45678, peerAddr, []icx.Route{{Src: privatePrefix, Dst: privatePrefix}}))
	require.NoError(t, h.InstallKeysForTest(0x45678, 1, key, key, time.Now().Add(time.Hour)))

	ip6 := makeIPv6UDPPacket()
	phy := make([]byte, 1500)

	n, loop := h.VirtToPhy(ip6, phy)
	require.NotZero(t, n)
	require.False(t, loop)

	out := make([]byte, 1500)
	m := h.PhyToVirt(phy[:n], out)
	require.Equal(t, len(ip6), m)
	require.Equal(t, ip6, out[:m])
}

func TestUpdateVirtualNetworkRoutes(t *testing.T) {
	localAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 1).To4()),
		Port: 1234,
	}
	peerAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 2).To4()),
		Port: 4321,
	}

	var key [16]byte
	copy(key[:], []byte("0123456789abcdef"))

	h, err := icx.NewHandler(
		icx.WithLocalAddr(localAddr),
		icx.WithVirtMAC(tcpip.GetRandMacAddr()),
	)
	require.NoError(t, err)

	// Start with a prefix that contains 192.168.1.2 (dst of our inner packet)
	privatePrefix := netip.MustParsePrefix("192.168.1.0/24")
	err = h.AddVirtualNetwork(0x23456, peerAddr, []icx.Route{{Src: privatePrefix, Dst: privatePrefix}})
	require.NoError(t, err)
	require.NoError(t, h.InstallKeysForTest(0x23456, 1, key, key, time.Now().Add(time.Hour)))

	virt := makeIPv4UDPEthernetFrame(tcpip.GetRandMacAddr())
	phy := make([]byte, 1500)

	// Should select this virtual network (address mapping exists).
	n, loop := h.VirtToPhy(virt, phy)
	require.NotZero(t, n)
	require.False(t, loop)

	// Now change allowed addresses to *not* include 192.168.1.2.
	privatePrefix = netip.MustParsePrefix("10.0.0.0/8")
	err = h.UpdateVirtualNetworkRoutes(0x23456, []icx.Route{{Src: privatePrefix, Dst: privatePrefix}})
	require.NoError(t, err)

	// Mapping removed -> VirtToPhy should drop as unknown tunnel destination address.
	n, loop = h.VirtToPhy(virt, phy)
	require.Zero(t, n)
	require.False(t, loop)

	// Add back a prefix which matches the destination again.
	privatePrefix = netip.MustParsePrefix("192.168.1.0/24")
	err = h.UpdateVirtualNetworkRoutes(0x23456, []icx.Route{{Src: privatePrefix, Dst: privatePrefix}})
	require.NoError(t, err)

	n, loop = h.VirtToPhy(virt, phy)
	require.NotZero(t, n)
	require.False(t, loop)
}

func TestKeyRotation(t *testing.T) {
	clk := &fakeClock{now: time.Unix(1_700_000_000, 0)}

	local := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 1).To4()),
		Port: 1234,
	}
	peer := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 2).To4()),
		Port: 4321,
	}

	var k1, k2, k3, k4 [16]byte
	copy(k1[:], []byte("aaaaaaaaaaaaaaaa"))
	copy(k2[:], []byte("bbbbbbbbbbbbbbbb"))
	copy(k3[:], []byte("cccccccccccccccc"))
	copy(k4[:], []byte("dddddddddddddddd"))

	h, err := icx.NewHandler(
		icx.WithLocalAddr(local),
		icx.WithLayer3VirtFrames(),
		icx.WithClock(clk),
	)
	require.NoError(t, err)

	const vni = 0x4242
	privatePrefix := netip.MustParsePrefix("192.168.1.0/24")
	require.NoError(t, h.AddVirtualNetwork(vni, peer, []icx.Route{{Src: privatePrefix, Dst: privatePrefix}}))

	// Epoch 1
	require.NoError(t, h.InstallKeysForTest(vni, 1, k1, k1, clk.Now().Add(time.Hour)))

	ip := makeIPv4UDPPacket()
	phy := make([]byte, 2000)
	out := make([]byte, 2000)

	// Create TWO distinct epoch-1 ciphertexts (different TX counters), but do not decode them yet.
	n, loop := h.VirtToPhy(ip, phy)
	require.NotZero(t, n)
	require.False(t, loop)
	epoch1A := append([]byte(nil), phy[:n]...)

	n, loop = h.VirtToPhy(ip, phy)
	require.NotZero(t, n)
	require.False(t, loop)
	epoch1B := append([]byte(nil), phy[:n]...)

	// Rotate to epoch 2; epoch 1 gets grace
	require.NoError(t, h.InstallKeysForTest(vni, 2, k2, k2, clk.Now().Add(time.Hour)))

	// Within grace: one of the saved epoch-1 frames must decrypt.
	m := h.PhyToVirt(epoch1A, out)
	require.Equal(t, len(ip), m)
	require.Equal(t, ip, out[:m])

	// After grace: the *other* (previously unseen) epoch-1 frame must now be rejected.
	clk.Advance(31 * time.Second)
	m = h.PhyToVirt(epoch1B, out[:cap(out)])
	require.Zero(t, m)

	// Produce a single epoch-2 ciphertext and save it for later.
	n, loop = h.VirtToPhy(ip, phy)
	require.NotZero(t, n)
	require.False(t, loop)
	epoch2A := append([]byte(nil), phy[:n]...)

	// Rotate to epoch 3 (starts grace for epoch 2)
	require.NoError(t, h.InstallKeysForTest(vni, 3, k3, k3, clk.Now().Add(time.Hour)))

	// Let epoch-2 grace expire.
	clk.Advance(31 * time.Second)

	// Rotate to epoch 4; expired RX keys should be swept here
	require.NoError(t, h.InstallKeysForTest(vni, 4, k4, k4, clk.Now().Add(time.Hour)))

	// The saved epoch-2 frame should now be rejected (no matching key after cleanup).
	m = h.PhyToVirt(epoch2A, out[:cap(out)])
	require.Zero(t, m)

	// Sanity: current epoch (4) packets still round-trip.
	n, loop = h.VirtToPhy(ip, phy)
	require.NotZero(t, n)
	require.False(t, loop)
	m = h.PhyToVirt(phy[:n], out[:cap(out)])
	require.Equal(t, len(ip), m)
	require.Equal(t, ip, out[:m])
}

// TestUpdateVirtualNetworkSecretGuards exercises the fail-closed guards on the derived-key
// install seam. Reserved SPIs (low 31 bits zero) and rxSPI == txSPI (which would derive
// equal keys for both directions) are always rejected. The TX nonce-safety guard is
// scoped by master secret: under an UNCHANGED master the txSPI must strictly increase
// (re-installing at or below the live SPI would reset the from-zero counter into an
// already-used (key, nonce) space), while a DIFFERENT master accepts any SPI — a fresh
// master is a fresh derivation space, which is exactly the transient-reconnect case
// (the peer's allocator reset to low SPI values under a fresh session).
func TestUpdateVirtualNetworkSecretGuards(t *testing.T) {
	localAddr := &tcpip.FullAddress{Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 1).To4()), Port: 1234}
	peerAddr := &tcpip.FullAddress{Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 2).To4()), Port: 4321}

	h, err := icx.NewHandler(icx.WithLocalAddr(localAddr), icx.WithLayer3VirtFrames())
	require.NoError(t, err)

	const vni = 0x9999
	prefix := netip.MustParsePrefix("192.168.1.0/24")
	require.NoError(t, h.AddVirtualNetwork(vni, peerAddr, []icx.Route{{Src: prefix, Dst: prefix}}))

	var mA, mB [32]byte
	copy(mA[:], []byte("master-secret-A-aaaaaaaaaaaaaaaa"))
	copy(mB[:], []byte("master-secret-B-bbbbbbbbbbbbbbbb"))
	exp := time.Now().Add(time.Hour)

	// An all-zero master is rejected: its derived keys are publicly computable, so
	// it must never key a tunnel regardless of the SPIs.
	require.Error(t, h.UpdateVirtualNetworkSecret(vni, [32]byte{}, 10, 20, exp),
		"an all-zero master secret must be rejected")

	// Reserved SPIs (low 31 bits zero) are rejected in either direction — including
	// the master-key-selector MSB alone.
	require.Error(t, h.UpdateVirtualNetworkSecret(vni, mA, 0, 20, exp), "rx SPI 0 must be rejected")
	require.Error(t, h.UpdateVirtualNetworkSecret(vni, mA, 10, 0, exp), "tx SPI 0 must be rejected")
	require.Error(t, h.UpdateVirtualNetworkSecret(vni, mA, 0x80000000, 20, exp),
		"an SPI with only the MSB set is still reserved (low 31 bits zero)")

	// Equal SPIs are rejected: both directions would derive the same key.
	require.Error(t, h.UpdateVirtualNetworkSecret(vni, mA, 20, 20, exp),
		"rxSPI == txSPI must be rejected")

	// First install with distinct per-direction SPIs succeeds. Live tx SPI is now 20.
	require.NoError(t, h.UpdateVirtualNetworkSecret(vni, mA, 10, 20, exp))

	// Under the UNCHANGED master, a txSPI at or below the live one is rejected —
	// including the identical re-install, regardless of the rx side.
	require.Error(t, h.UpdateVirtualNetworkSecret(vni, mA, 10, 20, exp),
		"re-installing the identical live tx SA must be rejected")
	require.Error(t, h.UpdateVirtualNetworkSecret(vni, mA, 99, 20, exp),
		"same master + same tx SPI is unsafe no matter the rx SPI")
	require.Error(t, h.UpdateVirtualNetworkSecret(vni, mA, 3, 5, exp),
		"same master + lower tx SPI must be rejected (monotonicity)")

	// Under the same master a strictly higher txSPI succeeds; the rx SPI may repeat
	// (the receive side never emits a nonce, so there is no rx monotonicity guard).
	// Live tx SPI is now 21.
	require.NoError(t, h.UpdateVirtualNetworkSecret(vni, mA, 10, 21, exp),
		"same master + higher tx SPI with a reused rx SPI is accepted")

	// A DIFFERENT master accepts regressed SPIs — reconnect recovery: the fresh
	// master makes the low SPIs a fresh derivation space.
	require.NoError(t, h.UpdateVirtualNetworkSecret(vni, mB, 3, 5, exp),
		"a fresh master accepts regressed SPIs (reconnect recovery)")

	// And switching BACK to a previously-used master is treated as changed again —
	// the guard compares fingerprints, not history. This mirrors the control plane,
	// where a master is never resurrected; the SPI just has to differ from the live
	// generation's constraint set.
	require.NoError(t, h.UpdateVirtualNetworkSecret(vni, mA, 10, 21, exp),
		"a master change resets the monotonicity scope")
}

// TestDerivedKeysInterop proves two handlers sharing a master secret with mirrored
// psp.EpochSPIs roles derive matching per-direction keys — Geneve frames round-trip in
// both directions — and pins the handler's internal derivation to the PSP spec KDF by
// decrypting its output with a raw key derived via psp.DeriveSAKey in the test.
func TestDerivedKeysInterop(t *testing.T) {
	const vni = 0x7777
	addrA := tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 1).To4())
	addrB := tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 2).To4())
	hA := newPeerHandler(t, vni, addrA, addrB)
	hB := newPeerHandler(t, vni, addrB, addrA)

	var master [32]byte
	copy(master[:], []byte("interop-master-secret-0123456789"))
	exp := time.Now().Add(time.Hour)

	rxA, txA, err := psp.EpochSPIs(psp.Initiator, 42)
	require.NoError(t, err)
	rxB, txB, err := psp.EpochSPIs(psp.Responder, 42)
	require.NoError(t, err)
	require.Equal(t, txA, rxB, "mirrored roles: A transmits to B's receive SPI")
	require.Equal(t, rxA, txB, "mirrored roles: B transmits to A's receive SPI")

	require.NoError(t, hA.UpdateVirtualNetworkSecret(vni, master, rxA, txA, exp))
	require.NoError(t, hB.UpdateVirtualNetworkSecret(vni, master, rxB, txB, exp))

	ip := makeIPv4UDPPacket()
	phy := make([]byte, 1500)
	out := make([]byte, 1500)

	// A -> B and B -> A round-trip through two independent handler derivations.
	n, loop := hA.VirtToPhy(ip, phy)
	require.NotZero(t, n)
	require.False(t, loop)
	m := hB.PhyToVirt(phy[:n], out)
	require.NotZero(t, m, "B must decrypt A's traffic from the shared master")
	require.Equal(t, ip, out[:m])

	n, loop = hB.VirtToPhy(ip, phy)
	require.NotZero(t, n)
	require.False(t, loop)
	m = hA.PhyToVirt(phy[:n], out)
	require.NotZero(t, m, "A must decrypt B's traffic from the shared master")
	require.Equal(t, ip, out[:m])

	// Spec conformance: a THIRD handler installed with raw keys derived in the test
	// via psp.DeriveSAKey must decrypt hA's output byte-for-byte — the handler's
	// internal derivation IS the PSP KDF, not merely self-consistent.
	hC := newPeerHandler(t, vni, addrB, addrA)
	keyAtoB, err := psp.DeriveSAKey(master[:], txA, psp.AESGCM128)
	require.NoError(t, err)
	var raw [16]byte
	copy(raw[:], keyAtoB)
	require.NoError(t, hC.InstallKeysForTest(vni, txA, raw, raw, exp))

	n, loop = hA.VirtToPhy(ip, phy)
	require.NotZero(t, n)
	require.False(t, loop)
	m = hC.PhyToVirt(phy[:n], out)
	require.NotZero(t, m, "a raw psp.DeriveSAKey key must decrypt the handler's derived-key output")
	require.Equal(t, ip, out[:m])
}

// TestRXRejectsSPINonceMismatch proves the RX side rejects a frame whose nonce
// SPI (nonce[:4]) does not match the key epoch it selected, on BOTH the
// cross-buffer and in-place decap paths, and that TX binds the SPI into the
// nonce in the first place (the offset sanity checks below).
func TestRXRejectsSPINonceMismatch(t *testing.T) {
	localAddr := &tcpip.FullAddress{Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 1).To4()), Port: 1234}
	peerAddr := &tcpip.FullAddress{Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 2).To4()), Port: 4321}

	h, err := icx.NewHandler(icx.WithLocalAddr(localAddr), icx.WithLayer3VirtFrames())
	require.NoError(t, err)

	const vni = 0x1234
	prefix := netip.MustParsePrefix("192.168.1.0/24")
	require.NoError(t, h.AddVirtualNetwork(vni, peerAddr, []icx.Route{{Src: prefix, Dst: prefix}}))

	var key [16]byte
	copy(key[:], []byte("0123456789abcdef"))
	// Loopback (single shared key) so this one handler both encaps and decaps;
	// the production seam would reject equal rx/tx keys.
	require.NoError(t, h.InstallKeysForTest(vni, 1, key, key, time.Now().Add(time.Hour)))

	// Mint a real encrypted frame; TX binds epoch 1 into nonce[:4].
	ip := makeIPv4UDPPacket()
	phy := make([]byte, 1500)
	n, loop := h.VirtToPhy(ip, phy)
	require.NotZero(t, n)
	require.False(t, loop)

	// The Geneve header sits at the UDP payload offset; its layout is
	// base(8) + KeyEpoch option(hdr 4 + value 4) + TxCounter option(hdr 4 + value 12).
	// So the key-epoch value and the nonce (= TxCounter value, nonce[:4] = SPI)
	// live at these absolute offsets within the IPv4-underlay physical frame.
	const geneveBase, optHdr, epochValLen = 8, 4, 4
	keyEpochOff := udp.PayloadOffsetIPv4 + geneveBase + optHdr
	nonceOff := keyEpochOff + epochValLen + optHdr
	require.Equal(t, uint32(1), binary.BigEndian.Uint32(phy[keyEpochOff:keyEpochOff+4]),
		"sanity: key-epoch option should carry epoch 1")
	require.Equal(t, uint32(1), binary.BigEndian.Uint32(phy[nonceOff:nonceOff+4]),
		"TX must bind the SPI (epoch 1) into nonce[:4]")

	// Tamper nonce[:4] so it no longer matches the key epoch (which stays 1, so
	// the RX side still selects the installed cipher and reaches the SPI check).
	tampered := append([]byte(nil), phy[:n]...)
	tampered[nonceOff] = 0xFF
	require.NotEqual(t, uint32(1), binary.BigEndian.Uint32(tampered[nonceOff:nonceOff+4]))
	require.Equal(t, uint32(1), binary.BigEndian.Uint32(tampered[keyEpochOff:keyEpochOff+4]),
		"key epoch must remain 1 so the SPI check, not the key lookup, rejects the frame")

	vnet, ok := h.GetVirtualNetwork(vni)
	require.True(t, ok)

	// Cross-buffer decap drops it and counts an SPI mismatch.
	out := make([]byte, 1500)
	require.Zero(t, h.PhyToVirt(append([]byte(nil), tampered...), out))
	require.Equal(t, uint64(1), vnet.Stats.RXDropsSPIMismatch.Load())

	// In-place decap drops it identically (placed at a non-zero offset).
	buf := make([]byte, len(tampered)+256)
	const off = 64
	copy(buf[off:off+len(tampered)], tampered)
	gotOff, gotLen := h.PhyToVirtInPlace(buf, off, len(tampered))
	require.Zero(t, gotLen)
	require.Zero(t, gotOff)
	require.Equal(t, uint64(2), vnet.Stats.RXDropsSPIMismatch.Load())
}

func TestARPRequest_Loopback(t *testing.T) {
	if testing.Verbose() {
		slog.SetLogLoggerLevel(slog.LevelDebug)
	}

	localAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 1).To4()),
		Port: 1234,
	}
	virtMAC := tcpip.GetRandMacAddr()

	h, err := icx.NewHandler(
		icx.WithLocalAddr(localAddr),
		icx.WithVirtMAC(virtMAC), // L2 mode -> virtMAC required
	)
	require.NoError(t, err)

	// Build a minimal ARP request (who-has 192.168.1.2 tell 192.168.1.1).
	arpReq := makeARPEthernetRequestFrame(
		tcpip.GetRandMacAddr(),
		net.IPv4(192, 168, 1, 1),
		net.IPv4(192, 168, 1, 2),
	)

	phy := make([]byte, 2000)

	// ARP should be answered locally with an immediate ARP reply (loopback).
	n, loop := h.VirtToPhy(arpReq, phy)
	require.NotZero(t, n)
	require.True(t, loop)

	// Validate ARP reply.
	eth := header.Ethernet(phy[:n])
	require.Equal(t, header.ARPProtocolNumber, eth.Type())

	arp := phy[header.EthernetMinimumSize:n]
	require.GreaterOrEqual(t, len(arp), 28)

	// Opcode (bytes 6:8) == 2 for reply.
	op := binary.BigEndian.Uint16(arp[6:8])
	require.Equal(t, uint16(2), op, "expected ARP reply opcode")

	// Sender/target protocol addresses are swapped.
	spa := net.IP(arp[14:18])
	tpa := net.IP(arp[24:28])
	require.Equal(t, net.IPv4(192, 168, 1, 2).To4(), spa.To4())
	require.Equal(t, net.IPv4(192, 168, 1, 1).To4(), tpa.To4())
}

func TestNeighborSolicitation_Loopback(t *testing.T) {
	localAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 1).To4()),
		Port: 1234,
	}
	peerAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 2).To4()),
		Port: 4321,
	}
	virtMAC := tcpip.GetRandMacAddr()

	var key [16]byte
	copy(key[:], []byte("0123456789abcdef"))

	h, err := icx.NewHandler(
		icx.WithLocalAddr(localAddr),
		icx.WithVirtMAC(virtMAC),
	)
	require.NoError(t, err)

	privatePrefix := netip.MustParsePrefix("2001:db8::/64")
	require.NoError(t, h.AddVirtualNetwork(0x56789, peerAddr, []icx.Route{{Src: privatePrefix, Dst: privatePrefix}}))
	require.NoError(t, h.InstallKeysForTest(0x56789, 1, key, key, time.Now().Add(time.Hour)))

	nsFrame := makeIPv6NeighborSolicitationEthernetFrame()
	phy := make([]byte, 2000)

	// Should be handled locally with an immediate NA reply (loopback).
	n, loop := h.VirtToPhy(nsFrame, phy)
	require.NotZero(t, n)
	require.True(t, loop)
}

func TestGetAndListVirtualNetworks(t *testing.T) {
	localAddr := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 1).To4()),
		Port: 1234,
	}
	peerA := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 2).To4()),
		Port: 4321,
	}
	peerB := &tcpip.FullAddress{
		Addr: tcpip.AddrFrom4Slice(net.IPv4(10, 0, 0, 3).To4()),
		Port: 4322,
	}

	h, err := icx.NewHandler(
		icx.WithLocalAddr(localAddr),
		icx.WithVirtMAC(tcpip.GetRandMacAddr()),
	)
	require.NoError(t, err)

	prefixA := netip.MustParsePrefix("192.0.2.0/24")
	prefixB := netip.MustParsePrefix("203.0.113.0/24")
	require.NoError(t, h.AddVirtualNetwork(10, peerA, []icx.Route{{Src: prefixA, Dst: prefixA}}))
	require.NoError(t, h.AddVirtualNetwork(5, peerB, []icx.Route{{Src: prefixB, Dst: prefixB}}))

	// GetVirtualNetwork
	v, ok := h.GetVirtualNetwork(10)
	require.True(t, ok)
	require.Equal(t, uint(10), v.ID)

	_, ok = h.GetVirtualNetwork(999)
	require.False(t, ok)

	// ListVirtualNetworks should be sorted by VNI.
	list := h.ListVirtualNetworks()
	require.Len(t, list, 2)
	require.Equal(t, uint(5), list[0].ID)
	require.Equal(t, uint(10), list[1].ID)
}

func BenchmarkHandler(b *testing.B) {
	if testing.Verbose() {
		slog.SetLogLoggerLevel(slog.LevelDebug)
	}

	localAddr := mustNewFullAddress("10.0.0.1:6081")
	remoteAddr := mustNewFullAddress("10.0.0.2:6081")

	h, err := icx.NewHandler(icx.WithLocalAddr(localAddr),
		icx.WithVirtMAC(tcpip.GetRandMacAddr()))
	require.NoError(b, err)

	var key [16]byte
	copy(key[:], []byte("0123456789abcdef"))

	const vni = 0x12345

	privatePrefix := netip.MustParsePrefix("192.168.1.0/24")
	err = h.AddVirtualNetwork(vni, remoteAddr, []icx.Route{{Src: privatePrefix, Dst: privatePrefix}})
	require.NoError(b, err)

	err = h.InstallKeysForTest(0x12345, 1, key, key, time.Now().Add(time.Hour))
	require.NoError(b, err)

	virtMAC := tcpip.GetRandMacAddr()
	virtFrame := makeIPv4UDPEthernetFrame(virtMAC)

	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		phyFrame := make([]byte, 1500)
		decodedFrame := make([]byte, 1400)

		for pb.Next() {
			n, _ := h.VirtToPhy(virtFrame, phyFrame)
			require.NotZero(b, n, "Failed to encode frame")

			n = h.PhyToVirt(phyFrame, decodedFrame)
			require.NotZero(b, n, "Failed to decode frame")
		}
	})
}

func makeIPv4UDPEthernetFrame(virtMAC tcpip.LinkAddress) []byte {
	ipPacket := makeIPv4UDPPacket()

	frame := make([]byte, header.EthernetMinimumSize+len(ipPacket))
	copy(frame[header.EthernetMinimumSize:], ipPacket)

	eth := header.Ethernet(frame)
	eth.Encode(&header.EthernetFields{
		SrcAddr: tcpip.GetRandMacAddr(),
		DstAddr: tcpip.GetRandMacAddr(),
		Type:    header.IPv4ProtocolNumber,
	})

	return frame
}

func makeIPv4UDPPacket() []byte {
	ipPacket := make([]byte, header.IPv4MinimumSize+header.UDPMinimumSize)

	ip := header.IPv4(ipPacket)
	ip.Encode(&header.IPv4Fields{
		TotalLength: uint16(len(ipPacket)),
		TTL:         64,
		Protocol:    uint8(header.UDPProtocolNumber),
		SrcAddr:     tcpip.AddrFrom4Slice(net.IPv4(192, 168, 1, 1).To4()),
		DstAddr:     tcpip.AddrFrom4Slice(net.IPv4(192, 168, 1, 2).To4()),
	})
	ip.SetChecksum(^ip.CalculateChecksum())

	udp := header.UDP(ipPacket[header.IPv4MinimumSize:])
	udp.Encode(&header.UDPFields{
		SrcPort: 1234,
		DstPort: 5678,
		Length:  header.UDPMinimumSize,
	})

	return ipPacket
}

func makeIPv6UDPPacket() []byte {
	packet := make([]byte, header.IPv6MinimumSize+header.UDPMinimumSize)
	ip6 := header.IPv6(packet)
	ip6.Encode(&header.IPv6Fields{
		PayloadLength:     header.UDPMinimumSize,
		TransportProtocol: header.UDPProtocolNumber,
		HopLimit:          64,
		SrcAddr:           tcpip.AddrFrom16Slice(net.ParseIP("2001:db8::1").To16()),
		DstAddr:           tcpip.AddrFrom16Slice(net.ParseIP("2001:db8::2").To16()),
	})
	udp := header.UDP(packet[header.IPv6MinimumSize:])
	udp.Encode(&header.UDPFields{
		SrcPort: 1234,
		DstPort: 5678,
		Length:  header.UDPMinimumSize,
	})
	return packet
}

func makeARPEthernetRequestFrame(srcMAC tcpip.LinkAddress, srcIP, dstIP net.IP) []byte {
	// ARP payload layout (28 bytes total for Ethernet/IPv4):
	//  0-1:  HW type (1 = Ethernet)
	//  2-3:  Proto type (0x0800 = IPv4)
	//    4:  HW addr len (6)
	//    5:  Proto addr len (4)
	//  6-7:  Opcode (1 = request)
	//  8-13: Sender HW addr (6)
	// 14-17: Sender Proto addr (4)
	// 18-23: Target HW addr (6) -> zero for request
	// 24-27: Target Proto addr (4)

	arpPayload := make([]byte, 28)
	binary.BigEndian.PutUint16(arpPayload[0:2], 1)              // HW type: Ethernet
	binary.BigEndian.PutUint16(arpPayload[2:4], uint16(0x0800)) // Proto: IPv4
	arpPayload[4] = 6                                           // HW len
	arpPayload[5] = 4                                           // Proto len
	binary.BigEndian.PutUint16(arpPayload[6:8], 1)              // Opcode: request
	copy(arpPayload[8:14], []byte(srcMAC))                      // Sender MAC
	copy(arpPayload[14:18], srcIP.To4())                        // Sender IP
	// Target MAC (18:24) left zero for request
	copy(arpPayload[24:28], dstIP.To4()) // Target IP

	frame := make([]byte, header.EthernetMinimumSize+len(arpPayload))
	copy(frame[header.EthernetMinimumSize:], arpPayload)

	eth := header.Ethernet(frame)
	eth.Encode(&header.EthernetFields{
		SrcAddr: srcMAC,
		DstAddr: tcpip.GetRandMacAddr(), // arbitrary; not relevant for loopback decision
		Type:    header.ARPProtocolNumber,
	})
	return frame
}

func makeIPv6NeighborSolicitationEthernetFrame() []byte {
	// Minimal NS: IPv6 + ICMPv6 (type 135) + 16-byte target address.
	const nsPayloadLen = header.ICMPv6NeighborSolicitMinimumSize
	total := header.EthernetMinimumSize + header.IPv6MinimumSize + nsPayloadLen
	frame := make([]byte, total)

	ipPayload := frame[header.EthernetMinimumSize:]
	ip6 := header.IPv6(ipPayload)
	ip6.Encode(&header.IPv6Fields{
		PayloadLength:     nsPayloadLen,
		TransportProtocol: header.ICMPv6ProtocolNumber,
		HopLimit:          255,
		SrcAddr:           tcpip.AddrFrom16Slice(net.ParseIP("2001:db8::1").To16()),
		DstAddr:           tcpip.AddrFrom16Slice(net.ParseIP("ff02::1:ff00:2").To16()), // solicited-node multicast (example)
	})
	icmp := ipPayload[header.IPv6MinimumSize:]
	icmp[0] = byte(header.ICMPv6NeighborSolicit) // type 135
	icmp[1] = 0                                  // code
	// checksum left 0; handler doesn't validate it for loopback decision
	// Target Address (16 bytes) starts at offset 8
	copy(icmp[8:24], net.ParseIP("2001:db8::2").To16())
	// zero out reserved (4 bytes) and leave options empty
	binary.BigEndian.PutUint32(icmp[4:8], 0)

	eth := header.Ethernet(frame)
	eth.Encode(&header.EthernetFields{
		SrcAddr: tcpip.GetRandMacAddr(),
		DstAddr: tcpip.GetRandMacAddr(),
		Type:    header.IPv6ProtocolNumber,
	})
	return frame
}

func mustNewFullAddress(addrPortStr string) *tcpip.FullAddress {
	addrPort := netip.MustParseAddrPort(addrPortStr)

	switch addrPort.Addr().BitLen() {
	case 32:
		return &tcpip.FullAddress{
			Addr: tcpip.AddrFrom4Slice(addrPort.Addr().AsSlice()),
			Port: addrPort.Port(),
		}
	case 128:
		return &tcpip.FullAddress{
			Addr: tcpip.AddrFrom16Slice(addrPort.Addr().AsSlice()),
			Port: addrPort.Port(),
		}
	default:
		panic("Unsupported IP address length")
	}
}

// fakeClock lets us control time for grace/expiry testing.
type fakeClock struct {
	now time.Time
}

func (c *fakeClock) Now() time.Time          { return c.now }
func (c *fakeClock) Advance(d time.Duration) { c.now = c.now.Add(d) }