device: add API for on-demand configuration of peers

Updates tailscale/tailscale#17858

Signed-off-by: Brad Fitzpatrick <brad@danga.com>

device/allowedips.go

@@ -257,17 +257,17 @@ func (node *trieEntry) remove() {
 }
 
 func (table *AllowedIPs) Remove(prefix netip.Prefix, peer *Peer) {
-	table.mutex.Lock()
+	table.mu.Lock()
-	defer table.mutex.Unlock()
+	defer table.mu.Unlock()
 	var node *trieEntry
 	var exact bool
 
 	if prefix.Addr().Is6() {
 		ip := prefix.Addr().As16()
-		node, exact = table.IPv6.nodePlacement(ip[:], uint8(prefix.Bits()))
+		node, exact = table.ipv6.nodePlacement(ip[:], uint8(prefix.Bits()))
 	} else if prefix.Addr().Is4() {
 		ip := prefix.Addr().As4()
-		node, exact = table.IPv4.nodePlacement(ip[:], uint8(prefix.Bits()))
+		node, exact = table.ipv4.nodePlacement(ip[:], uint8(prefix.Bits()))
 	} else {
 		panic(errors.New("removing unknown address type"))
 	}
@@ -277,10 +277,26 @@ func (table *AllowedIPs) Remove(prefix netip.Prefix, peer *Peer) {
 	node.remove()
 }
 
-func (table *AllowedIPs) RemoveByPeer(peer *Peer) {
+
+// setPeerPrefixes atomically removes all of peer's existing prefixes and adds
+// the provided ones.
+func (table *AllowedIPs) setPeerPrefixes(peer *Peer, prefixes []netip.Prefix) {
 	table.mu.Lock()
 	defer table.mu.Unlock()
 
+	table.removeByPeerLocked(peer)
+	for _, prefix := range prefixes {
+		table.insertLocked(prefix, peer)
+	}
+}
+
+func (table *AllowedIPs) RemoveByPeer(peer *Peer) {
+	table.mu.Lock()
+	defer table.mu.Unlock()
+	table.removeByPeerLocked(peer)
+}
+
+func (table *AllowedIPs) removeByPeerLocked(peer *Peer) {
 	var next *list.Element
 	for elem := peer.trieEntries.Front(); elem != nil; elem = next {
 		next = elem.Next()
@@ -291,7 +307,10 @@ func (table *AllowedIPs) RemoveByPeer(peer *Peer) {
 func (table *AllowedIPs) Insert(prefix netip.Prefix, peer *Peer) {
 	table.mu.Lock()
 	defer table.mu.Unlock()
+	table.insertLocked(prefix, peer)
+}
 
+func (table *AllowedIPs) insertLocked(prefix netip.Prefix, peer *Peer) {
 	if prefix.Addr().Is6() {
 		ip := prefix.Addr().As16()
 		parentIndirection{&table.ipv6, 2}.insert(ip[:], uint8(prefix.Bits()), peer)

device/device.go

@@ -7,6 +7,8 @@ package device
 
 import (
 	"context"
+	"errors"
+	"net/netip"
 	"runtime"
 	"sync"
 	"sync/atomic"
@@ -59,6 +61,7 @@ type Device struct {
 	peers struct {
 		sync.RWMutex // protects keyMap
 		keyMap       map[NoisePublicKey]*Peer
+		lookupFunc   PeerLookupFunc // or nil if unused
 	}
 
 	rate struct {
@@ -345,13 +348,63 @@ func (device *Device) BatchSize() int {
 	return size
 }
 
+// LookupPeer looks up a peer by its public key.
+//
+// If the peer does not exist and a [PeerLookupFunc] is set (via
+// [Device.SetPeerLookupFunc]), then that function is used to create the peer
+// before returning it. Peers created via this mechanism exist only until their
+// state machine reaches idle, and then the peers are removed.
+//
+// If the peer does not exist and no [PeerLookupFunc] is set, nil is returned.
+//
+// Use [Device.LookupActivePeer] to only return already-existing peers, without
+// using a [PeerLookupFunc].
 func (device *Device) LookupPeer(pk NoisePublicKey) *Peer {
 	device.peers.RLock()
-	defer device.peers.RUnlock()
+	p, ok := device.peers.keyMap[pk]
+	lookupFunc := device.peers.lookupFunc
+	device.peers.RUnlock()
+	if ok || lookupFunc == nil {
+		return p
+	}
 
+	allowedIPs := lookupFunc(pk)
+	if allowedIPs == nil {
+		return nil
+	}
+
+	p, err := device.NewPeer(pk)
+	if err != nil {
+		if errors.Is(err, errAddExistingPeer) {
+			device.peers.RLock()
+			defer device.peers.RUnlock()
 			return device.peers.keyMap[pk]
+		}
+		device.log.Errorf("Failed to create peer: %v", err)
+		return nil
+	}
+	p.SetAllowedIPs(allowedIPs)
+	p.deleteOnIdle = true
+	p.Start()
+	return p
 }
 
+// LookupActivePeer looks up a peer by its public key.
+//
+// Unlike [Device.LookupPeer], this function does not use a [PeerLookupFunc] to
+// create the peer if it does not already exist.
+//
+// If the peer does not exist or was created lazily via [PeerLookupFunc]
+// and has subsequently idled away, it returns (nil, false).
+func (device *Device) LookupActivePeer(pk NoisePublicKey) (_ *Peer, ok bool) {
+	device.peers.RLock()
+	defer device.peers.RUnlock()
+	p, ok := device.peers.keyMap[pk]
+	return p, ok
+}
+
+var errAddExistingPeer = errors.New("adding existing peer")
+
 func (device *Device) RemovePeer(key NoisePublicKey) {
 	device.peers.Lock()
 	defer device.peers.Unlock()
@@ -374,6 +427,41 @@ func (device *Device) RemoveAllPeers() {
 	device.peers.keyMap = make(map[NoisePublicKey]*Peer)
 }
 
+// RemoveMatchingPeers removes all peers for which shouldRemove returns true.
+//
+// It returns the number of peers removed.
+func (device *Device) RemoveMatchingPeers(shouldRemove func(NoisePublicKey) bool) (numRemoved int) {
+	device.peers.Lock()
+	defer device.peers.Unlock()
+
+	for key, peer := range device.peers.keyMap {
+		if shouldRemove(key) {
+			removePeerLocked(device, peer, key)
+			numRemoved++
+		}
+	}
+	return numRemoved
+}
+
+// PeerLookupFunc is the type of function used to look up peers by public key
+// when receiving packets for unknown peers.
+//
+// If it returns nil, the peer is not known.
+//
+// Otherwise, returning non-nil signals that wireguard-go should create the peer
+// with the provided allowed IPs.
+//
+// See [Device.SetPeerLookupFunc] and [Device.LookupPeer].
+type PeerLookupFunc func(NoisePublicKey) (allowedIPs []netip.Prefix)
+
+// SetPeerLookupFunc sets the function used to look up peers by public key
+// when receiving packets for unknown peers.
+func (device *Device) SetPeerLookupFunc(f PeerLookupFunc) {
+	device.peers.Lock()
+	defer device.peers.Unlock()
+	device.peers.lookupFunc = f
+}
+
 func (device *Device) Close() {
 	device.state.Lock()
 	defer device.state.Unlock()

device/peer.go

@@ -8,6 +8,8 @@ package device
 import (
 	"container/list"
 	"errors"
+	"net/netip"
+	"slices"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -27,6 +29,12 @@ type Peer struct {
 
 	queuedOutboundPackets atomic.Int32 // packets in staged+outbound queues, for input backpressure
 
+	// deleteOnIdle indicates whether the peer should be deleted when idle
+	// because it was auto-created via a Device.PeerLookupFunc.
+	//
+	// This field should only be set once, before the peer is started.
+	deleteOnIdle bool
+
 	endpoint struct {
 		sync.Mutex
 		val            conn.Endpoint
@@ -46,7 +54,9 @@ type Peer struct {
 	}
 
 	state struct {
-		sync.Mutex // protects against concurrent Start/Stop
+		sync.Mutex // protects against concurrent Start/Stop, and fields below
+
+		allowedIPs []netip.Prefix
 	}
 
 	queue struct {
@@ -89,7 +99,7 @@ func (device *Device) NewPeer(pk NoisePublicKey) (*Peer, error) {
 	// map public key
 	_, ok := device.peers.keyMap[pk]
 	if ok {
-		return nil, errors.New("adding existing peer")
+		return nil, errAddExistingPeer
 	}
 
 	// pre-compute DH
@@ -115,6 +125,22 @@ func (device *Device) NewPeer(pk NoisePublicKey) (*Peer, error) {
 	return peer, nil
 }
 
+// SetAllowedIPs sets the allowed IP prefixes for this peer.
+//
+// If the allowedIPs are unchanged since the last call, this method is a no-op.
+// It's the caller's responsibility to ensure that no two peers have duplicate
+// allowed IPs. If so, the last writer wins.
+func (p *Peer) SetAllowedIPs(allowedIPs []netip.Prefix) {
+	p.state.Lock()
+	defer p.state.Unlock()
+
+	if slices.Equal(p.state.allowedIPs, allowedIPs) {
+		return
+	}
+	p.device.allowedips.setPeerPrefixes(p, allowedIPs)
+	p.state.allowedIPs = slices.Clone(allowedIPs) // avoid retaining caller's slice
+}
+
 // SendBuffers sends buffers to peer. WireGuard packet data in each element of
 // buffers must be preceded by MessageEncapsulatingTransportSize number of
 // bytes.

device/timers.go

@@ -129,6 +129,14 @@ func expiredNewHandshake(peer *Peer) {
 func expiredZeroKeyMaterial(peer *Peer) {
 	peer.device.log.Verbosef("%s - Removing all keys, since we haven't received a new one in %d seconds", peer, int((RejectAfterTime * 3).Seconds()))
 	peer.ZeroAndFlushAll()
+	if peer.deleteOnIdle {
+		peer.device.log.Verbosef("%s - Removing idle lazy peer", peer)
+		// Remove the peer from the device in a new goroutine as we're currently
+		// holding timer locks which RemovePeer also needs. This is TOCTOU, but
+		// acceptable since the worst case is we remove the peer and the lazy
+		// peerfunc created it again after. We might lose some packets.
+		go peer.device.RemovePeer(peer.handshake.remoteStatic)
+	}
 }
 
 func expiredPersistentKeepalive(peer *Peer) {

go.mod

@@ -1,6 +1,6 @@
 module github.com/sagernet/wireguard-go
 
-go 1.20
+go 1.25
 
 require (
 	github.com/sagernet/sing v0.7.10