netbird/client/iface/configurer/usp_allowedips_test.go

package configurer

import (
	"net"
	"net/netip"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	wgconn "golang.zx2c4.com/wireguard/conn"
	wgdevice "golang.zx2c4.com/wireguard/device"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	"github.com/netbirdio/netbird/client/iface/bind"
)

// newTestUSPConfigurer builds a configurer over a real wireguard-go device backed by an
// in-memory TUN. The device stays down, so no socket is opened and no privileges are needed.
func newTestUSPConfigurer(t *testing.T) *WGUSPConfigurer {
	t.Helper()

	dev := wgdevice.NewDevice(newIdleTUN(), wgconn.NewDefaultBind(), wgdevice.NewLogger(wgdevice.LogLevelSilent, ""))
	t.Cleanup(dev.Close)

	c := NewUSPConfigurerNoUAPI(dev, "wgtest0", bind.NewActivityRecorder())

	key, err := wgtypes.GeneratePrivateKey()
	require.NoError(t, err, "generate device private key")
	require.NoError(t, c.ConfigureInterface(key.String(), 0), "configure test device")

	return c
}

// seedPeers adds count peers, each with a /32 overlay address, and returns their public keys.
func seedPeers(t *testing.T, c *WGUSPConfigurer, count int) []string {
	t.Helper()

	keys := make([]string, 0, count)
	for i := 0; i < count; i++ {
		priv, err := wgtypes.GeneratePrivateKey()
		require.NoError(t, err, "generate peer private key")
		pub := priv.PublicKey().String()

		addr := netip.PrefixFrom(netip.AddrFrom4([4]byte{100, 64, byte(i >> 8), byte(i)}), 32)
		require.NoError(t, c.UpdatePeer(pub, []netip.Prefix{addr}, 25*time.Second, nil, nil), "add peer")
		keys = append(keys, pub)
	}
	return keys
}

func peerAllowedIPs(t *testing.T, c *WGUSPConfigurer, peerKey string) []string {
	t.Helper()

	stats, err := c.FullStats()
	require.NoError(t, err, "read device stats")

	for _, p := range stats.Peers {
		if p.PublicKey != peerKey {
			continue
		}
		got := make([]string, 0, len(p.AllowedIPs))
		for _, ipNet := range p.AllowedIPs {
			got = append(got, ipNet.String())
		}
		return got
	}
	t.Fatalf("peer %s not found on device", peerKey)
	return nil
}

// TestRemoveEndpointAddressPreservesRoutedAllowedIPs covers the prefixes the route manager
// attaches to a routing peer through AddAllowedIP. Those are not known to the peer.Conn that
// triggers the endpoint removal, so dropping them here would silently blackhole every route
// behind that peer on each relay or ICE disconnect.
func TestRemoveEndpointAddressPreservesRoutedAllowedIPs(t *testing.T) {
	c := newTestUSPConfigurer(t)
	peerKey := seedPeers(t, c, 3)[1]

	routed := []netip.Prefix{
		netip.MustParsePrefix("10.20.0.0/16"),
		netip.MustParsePrefix("192.168.7.0/24"),
	}
	for _, prefix := range routed {
		require.NoError(t, c.AddAllowedIP(peerKey, prefix), "add routed prefix")
	}

	before := peerAllowedIPs(t, c, peerKey)
	require.Len(t, before, 3, "peer should hold its overlay address plus both routed prefixes")

	require.NoError(t, c.RemoveEndpointAddress(peerKey), "remove endpoint address")

	assert.ElementsMatch(t, before, peerAllowedIPs(t, c, peerKey),
		"allowed IPs must survive the endpoint removal unchanged")
}

// TestRemoveEndpointAddressDoesNotScaleWithPeerCount is the regression guard for the actual
// defect: clearing one peer's endpoint used to dump and parse the whole device, so its cost
// grew with the size of the network map. On a routing peer with thousands of peers that dump
// runs on every relay and ICE transition, under the interface lock.
func TestRemoveEndpointAddressDoesNotScaleWithPeerCount(t *testing.T) {
	measure := func(peerCount int) float64 {
		c := newTestUSPConfigurer(t)
		peerKey := seedPeers(t, c, peerCount)[peerCount/2]

		return testing.AllocsPerRun(5, func() {
			require.NoError(t, c.RemoveEndpointAddress(peerKey), "remove endpoint address")
		})
	}

	small := measure(64)
	large := measure(1024)

	assert.Less(t, large, small*2,
		"clearing one endpoint allocated %.0f objects with 1024 peers against %.0f with 64: the cost still scales with the peer count",
		large, small)
}

func TestRemoveAllowedIPKeepsTheOtherPrefixes(t *testing.T) {
	c := newTestUSPConfigurer(t)
	peerKey := seedPeers(t, c, 3)[0]
	routed := netip.MustParsePrefix("10.20.0.0/16")
	require.NoError(t, c.AddAllowedIP(peerKey, routed), "add routed prefix")
	require.NoError(t, c.AddAllowedIP(peerKey, netip.MustParsePrefix("192.168.7.0/24")), "add routed prefix")

	require.NoError(t, c.RemoveAllowedIP(peerKey, routed), "remove routed prefix")

	assert.ElementsMatch(t, []string{"100.64.0.0/32", "192.168.7.0/24"}, peerAllowedIPs(t, c, peerKey),
		"only the removed prefix should be gone")

	assert.ErrorIs(t, c.RemoveAllowedIP(peerKey, routed), ErrAllowedIPNotFound,
		"removing a prefix that is no longer configured must be reported")
}

// TestAddAllowedIPOnAbsentPeerDoesNotResurrectIt covers the lazy connection window documented
// in #6863: AddAllowedIP is update-only, a silent no-op when the peer is absent, and clearing
// the endpoint of an absent peer must not create one either.
func TestAddAllowedIPOnAbsentPeerDoesNotResurrectIt(t *testing.T) {
	c := newTestUSPConfigurer(t)
	seedPeers(t, c, 2)

	priv, err := wgtypes.GeneratePrivateKey()
	require.NoError(t, err, "generate peer private key")
	absent := priv.PublicKey().String()

	require.NoError(t, c.AddAllowedIP(absent, netip.MustParsePrefix("10.20.0.0/16")),
		"update-only add on an absent peer is a silent no-op")

	stats, err := c.FullStats()
	require.NoError(t, err, "read device stats")
	require.Len(t, stats.Peers, 2, "the absent peer must not have been created by AddAllowedIP")

	assert.ErrorIs(t, c.RemoveEndpointAddress(absent), ErrPeerNotFound,
		"clearing the endpoint of a peer the device does not have must fail")

	stats, err = c.FullStats()
	require.NoError(t, err, "read device stats")
	assert.Len(t, stats.Peers, 2, "no peer may be created while clearing an endpoint")
}

// TestRemoveEndpointAddressDoesNotStealAPrefixFromAnotherPeer covers WireGuard's rule that an
// allowed IP belongs to exactly one peer: configuring a prefix on a peer takes it away from
// whichever peer held it before. UpdatePeer relies on that rule rather than removing the prefix
// from the previous holder itself, so a prefix handed over between peers must not come back.
func TestRemoveEndpointAddressDoesNotStealAPrefixFromAnotherPeer(t *testing.T) {
	c := newTestUSPConfigurer(t)
	keys := seedPeers(t, c, 2)
	peerA, peerB := keys[0], keys[1]
	routed := netip.MustParsePrefix("10.20.0.0/16")

	require.NoError(t, c.AddAllowedIP(peerA, routed), "give the prefix to A")
	require.Contains(t, peerAllowedIPs(t, c, peerA), routed.String(), "A must hold the prefix")

	// The route moves to B. The device takes it away from A on its own.
	require.NoError(t, c.AddAllowedIP(peerB, routed), "hand the prefix over to B")
	require.Contains(t, peerAllowedIPs(t, c, peerB), routed.String(), "B must hold the prefix")
	require.NotContains(t, peerAllowedIPs(t, c, peerA), routed.String(), "the device must have taken it from A")

	require.NoError(t, c.RemoveEndpointAddress(peerA), "clear A's endpoint")

	assert.NotContains(t, peerAllowedIPs(t, c, peerA), routed.String(),
		"clearing A's endpoint must not take the prefix back from B")
	assert.Contains(t, peerAllowedIPs(t, c, peerB), routed.String(),
		"B must still hold the prefix")
}

// TestPresharedKeyCreatedPeerTakesPartInPrefixHandover covers a peer created by a preshared
// key write rather than by a peer update. Rosenpass applies a peer's first key without
// updateOnly, which creates the peer on the device, so a store that ignored that operation
// would treat the peer as unknown and would not account for a prefix later handed over to it.
func TestPresharedKeyCreatedPeerTakesPartInPrefixHandover(t *testing.T) {
	c := newTestUSPConfigurer(t)
	peerA := seedPeers(t, c, 1)[0]
	routed := netip.MustParsePrefix("10.20.0.0/16")
	require.NoError(t, c.AddAllowedIP(peerA, routed), "give the prefix to A")

	priv, err := wgtypes.GeneratePrivateKey()
	require.NoError(t, err, "generate peer private key")
	peerB := priv.PublicKey().String()

	psk, err := wgtypes.GenerateKey()
	require.NoError(t, err, "generate preshared key")
	require.NoError(t, c.SetPresharedKey(peerB, psk, false), "a first key creates the peer")

	require.NoError(t, c.AddAllowedIP(peerB, routed), "hand the prefix over to B")
	require.Contains(t, peerAllowedIPs(t, c, peerB), routed.String(), "B must hold the prefix")

	require.NoError(t, c.RemoveEndpointAddress(peerA), "clear A's endpoint")

	assert.NotContains(t, peerAllowedIPs(t, c, peerA), routed.String(),
		"clearing A's endpoint must not take the prefix back from B")
	assert.Contains(t, peerAllowedIPs(t, c, peerB), routed.String(), "B must still hold the prefix")
}

// TestUpdatePeerDoesNotWidenAMappedPrefixOnTheDevice is the end to end form of the
// conversion: a v4-mapped prefix must not reach the device as a zero length allowed IP,
// which would route every v4 address to that peer.
func TestUpdatePeerDoesNotWidenAMappedPrefixOnTheDevice(t *testing.T) {
	c := newTestUSPConfigurer(t)

	priv, err := wgtypes.GeneratePrivateKey()
	require.NoError(t, err, "generate peer private key")
	peerKey := priv.PublicKey().String()

	mapped := netip.MustParsePrefix("::ffff:10.1.2.3/112")
	require.NoError(t, c.UpdatePeer(peerKey, []netip.Prefix{mapped}, 25*time.Second, nil, nil), "add peer")

	onDevice := peerAllowedIPs(t, c, peerKey)
	assert.NotContains(t, onDevice, "0.0.0.0/0", "the device must not be given a catch-all allowed IP")
	assert.Equal(t, []string{"10.1.0.0/16"}, onDevice, "the device holds the normalized prefix")
}

// TestUpdatePeerWithAnUnusableEndpointTouchesNothing pins the ordering: the endpoint is
// parsed before the device is configured, so a failure cannot leave a half-configured peer
// on the device.
func TestUpdatePeerWithAnUnusableEndpointTouchesNothing(t *testing.T) {
	c := newTestUSPConfigurer(t)
	seedPeers(t, c, 2)

	priv, err := wgtypes.GeneratePrivateKey()
	require.NoError(t, err, "generate peer private key")
	peerKey := priv.PublicKey().String()

	// A three byte address has no textual form netip can parse back.
	endpoint := &net.UDPAddr{IP: net.IP{1, 2, 3}, Port: 51820}
	require.Error(t, c.UpdatePeer(peerKey, []netip.Prefix{netip.MustParsePrefix("10.30.0.0/16")},
		25*time.Second, endpoint, nil), "an unusable endpoint must fail the update")

	stats, err := c.FullStats()
	require.NoError(t, err, "read device stats")
	assert.Len(t, stats.Peers, 2, "the peer must not have reached the device")
}