From e924a91e998843e753999047e1a9785290596030 Mon Sep 17 00:00:00 2001 From: Brad Fitzpatrick Date: Wed, 12 Nov 2025 13:04:17 -0800 Subject: [PATCH] device: add API for on-demand configuration of peers Updates tailscale/tailscale#17858 Signed-off-by: Brad Fitzpatrick --- device/allowedips.go | 29 +++++++++++--- device/device.go | 92 +++++++++++++++++++++++++++++++++++++++++++- device/peer.go | 30 ++++++++++++++- device/timers.go | 8 ++++ go.mod | 2 +- 5 files changed, 151 insertions(+), 10 deletions(-) diff --git a/device/allowedips.go b/device/allowedips.go index 84fbd90..05081ac 100644 --- a/device/allowedips.go +++ b/device/allowedips.go @@ -257,17 +257,17 @@ func (node *trieEntry) remove() { } func (table *AllowedIPs) Remove(prefix netip.Prefix, peer *Peer) { - table.mutex.Lock() - defer table.mutex.Unlock() + table.mu.Lock() + 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) diff --git a/device/device.go b/device/device.go index 529951a..1a05490 100644 --- a/device/device.go +++ b/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 + } - return device.peers.keyMap[pk] + 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() diff --git a/device/peer.go b/device/peer.go index a2703bd..875d6ba 100644 --- a/device/peer.go +++ b/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. diff --git a/device/timers.go b/device/timers.go index 80fb7d9..97ed451 100644 --- a/device/timers.go +++ b/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) { diff --git a/go.mod b/go.mod index 773bf00..d445678 100644 --- a/go.mod +++ b/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