Update

vendor/github.com/tailscale/wireguard-go/device/peer.go

@@ -7,7 +7,11 @@ package device
 
 import (
 	"container/list"
+	"encoding/binary"
 	"errors"
+	"net/netip"
+	"runtime"
+	"slices"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -16,26 +20,53 @@ import (
 )
 
 type Peer struct {
-	isRunning         atomic.Bool
-	keypairs          Keypairs
-	handshake         Handshake
-	device            *Device
-	stopping          sync.WaitGroup // routines pending stop
-	txBytes           atomic.Uint64  // bytes send to peer (endpoint)
-	rxBytes           atomic.Uint64  // bytes received from peer
-	lastHandshakeNano atomic.Int64   // nano seconds since epoch
+	runningState struct {
+		// isRunning is swapped across [Peer.Start] and [Peer.Stop]. Direct
+		// access loads are best-effort and inherently racey.
+		isRunning atomic.Bool
+		// queueWriters is the count of outstanding goroutines performing
+		// potential write operations against [Peer.queue].* channels.
+		// It is semantically similar to a [sync.WaitGroup], but free
+		// from the event order constraints of [sync.WaitGroup.Add]. Increments
+		// from zero may run concurrent to the "waiting" side.
+		queueWriters atomic.Int32
+		// queueReaders is the count of outstanding goroutines performing
+		// potential read operations against [Peer.queue.inbound] and
+		// [Peer.queue.outbound].
+		queueReaders sync.WaitGroup
+	}
+
+	keypairs            Keypairs
+	handshake           Handshake
+	device              *Device
+	txBytes             atomic.Uint64 // bytes send to peer (endpoint)
+	rxBytes             atomic.Uint64 // bytes received from peer
+	lastHandshakeNano   atomic.Int64  // nanoseconds since epoch
+	handshakeOnUserSend atomic.Bool   // attempt handshake following next outbound transport
+
+	sessionState struct {
+		sync.Mutex
+		current        PeerSessionState
+		sessionExpires time.Time
+	}
+
+	// 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
-		clearSrcOnTx   bool // signal to val.ClearSrc() prior to next packet transmission
-		disableRoaming bool
+		val          conn.Endpoint
+		clearSrcOnTx bool // signal to val.ClearSrc() prior to next packet transmission
 	}
 
 	timers struct {
 		retransmitHandshake     *Timer
 		sendKeepalive           *Timer
 		newHandshake            *Timer
+		sessionExpired          *Timer
 		zeroKeyMaterial         *Timer
 		persistentKeepalive     *Timer
 		handshakeAttempts       atomic.Uint32
@@ -44,18 +75,43 @@ 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
+
+		// testAllowedIP, if non-nil, is used to test whether the peer is
+		// allowed to send a packet from the given IP address. It can be read
+		// without locking, but must be set with the state mutex locked.
+		testAllowedIP atomic.Pointer[func(netip.Addr) bool]
 	}
 
+	// queue channel write operations must be wrapped by [Peer.doIfRunning]
 	queue struct {
 		staged   chan *QueueOutboundElementsContainer // staged packets before a handshake is available
-		outbound *autodrainingOutboundQueue           // sequential ordering of udp transmission
-		inbound  *autodrainingInboundQueue            // sequential ordering of tun writing
+		outbound chan *QueueOutboundElementsContainer // sequential ordering of udp transmission
+		inbound  chan *QueueInboundElementsContainer  // sequential ordering of tun writing
 	}
 
 	cookieGenerator             CookieGenerator
 	trieEntries                 list.List
 	persistentKeepaliveInterval atomic.Uint32
+
+	// flowID pins this peer's decrypted packets to a single TUN write queue
+	// when the device is a [tun.MultiQueueDevice].
+	// See [Peer.writeTUN].
+	//
+	// The kernel picks a TUN queue by 5-tuple for new (or quiescent,
+	// TUN_FLOW_EXPIRE=3s) flows. A pin learned from our writes overrides
+	// the hash. Kernel delivers a packet to the stack before recording the pin,
+	// so the first reply to a flow we write leaves on the hashed queue and the
+	// rest on ours. Both queues then carry that flow briefly, which may
+	// introduce overlay reordering.
+	// See https://elixir.bootlin.com/linux/v7.0/source/drivers/net/tun.c#L461
+	//
+	// TODO match or influence kernel hashing when writing to TUN to avoid
+	// reorders. Note that the kernel hash is siphash keyed by an unexported
+	// random value. This can be overridden with TUNSETSTEERINGEBPF ioctl.
+	flowID uint32
 }
 
 func (device *Device) NewPeer(pk NoisePublicKey) (*Peer, error) {
@@ -80,14 +136,17 @@ func (device *Device) NewPeer(pk NoisePublicKey) (*Peer, error) {
 
 	peer.cookieGenerator.Init(pk)
 	peer.device = device
-	peer.queue.outbound = newAutodrainingOutboundQueue(device)
-	peer.queue.inbound = newAutodrainingInboundQueue(device)
-	peer.queue.staged = make(chan *QueueOutboundElementsContainer, QueueStagedSize)
+	// Public keys are uniformly distributed, so the low bytes spread fine.
+	peer.flowID = binary.LittleEndian.Uint32(pk[:4])
+
+	// staged is never closed, and it can be accessed concurrent to [Peer.Start],
+	// so we init here instead of [Peer.Start].
+	peer.queue.staged = make(chan *QueueOutboundElementsContainer, device.config.queueStagedSize)
 
 	// map public key
 	_, ok := device.peers.keyMap[pk]
 	if ok {
-		return nil, errors.New("adding existing peer")
+		return nil, errAddExistingPeer
 	}
 
 	// pre-compute DH
@@ -100,7 +159,6 @@ func (device *Device) NewPeer(pk NoisePublicKey) (*Peer, error) {
 	// reset endpoint
 	peer.endpoint.Lock()
 	peer.endpoint.val = nil
-	peer.endpoint.disableRoaming = false
 	peer.endpoint.clearSrcOnTx = false
 	peer.endpoint.Unlock()
 
@@ -113,6 +171,35 @@ 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)
+
+	allowedIPs = slices.Clone(allowedIPs) // avoid retaining caller's slice
+	p.state.allowedIPs = allowedIPs
+
+	f := mkIPInCIDRsTestFunc(allowedIPs)
+	p.state.testAllowedIP.Store(&f)
+}
+
+// SetPresharedKey sets the optional WireGuard pre-shared key for this peer.
+// The zero value disables the pre-shared-key layer.
+func (p *Peer) SetPresharedKey(psk NoisePresharedKey) {
+	p.handshake.mutex.Lock()
+	defer p.handshake.mutex.Unlock()
+	p.handshake.presharedKey = psk
+}
+
 // SendBuffers sends buffers to peer. WireGuard packet data in each element of
 // buffers must be preceded by MessageEncapsulatingTransportSize number of
 // bytes.
@@ -147,6 +234,12 @@ func (peer *Peer) SendBuffers(buffers [][]byte) error {
 	return err
 }
 
+// writeTUN writes bufs to the TUN queue this peer is pinned to, keeping the
+// peer's decrypted packets in order.
+func (peer *Peer) writeTUN(bufs [][]byte, offset int) (int, error) {
+	return peer.device.tun.writeTo(int(peer.flowID), bufs, offset)
+}
+
 func (peer *Peer) String() string {
 	// The awful goo that follows is identical to:
 	//
@@ -183,16 +276,20 @@ func (peer *Peer) Start() {
 	peer.state.Lock()
 	defer peer.state.Unlock()
 
-	if peer.isRunning.Load() {
+	if peer.runningState.isRunning.Load() {
 		return
 	}
 
 	device := peer.device
 	device.log.Verbosef("%v - Starting", peer)
 
+	// init inbound & outbound packet queues
+	peer.queue.outbound = make(chan *QueueOutboundElementsContainer, device.config.queueOutboundSize)
+	peer.queue.inbound = make(chan *QueueInboundElementsContainer, device.config.queueInboundSize)
+
 	// reset routine state
-	peer.stopping.Wait()
-	peer.stopping.Add(2)
+	peer.runningState.queueReaders.Wait()
+	peer.runningState.queueReaders.Add(2)
 
 	peer.handshake.mutex.Lock()
 	peer.handshake.lastSentHandshake = time.Now().Add(-(RekeyTimeout + time.Second))
@@ -202,20 +299,28 @@ func (peer *Peer) Start() {
 
 	peer.timersStart()
 
-	device.flushInboundQueue(peer.queue.inbound)
-	device.flushOutboundQueue(peer.queue.outbound)
-
 	// Use the device batch size, not the bind batch size, as the device size is
 	// the size of the batch pools.
 	batchSize := peer.device.BatchSize()
 	go peer.RoutineSequentialSender(batchSize)
 	go peer.RoutineSequentialReceiver(batchSize)
 
-	peer.isRunning.Store(true)
+	peer.runningState.isRunning.Store(true)
+
+	// A lazily-created peer that never completes a handshake otherwise never
+	// arms its reaping timer. Arm it here, while running under state.Lock, so
+	// it's reclaimed after RejectAfterTime*3 of no session and is guaranteed to
+	// be torn down by a matching Stop. A completed handshake re-Mods it.
+	if peer.deleteOnIdle {
+		peer.timers.zeroKeyMaterial.Mod(RejectAfterTime * 3)
+	}
 }
 
 func (peer *Peer) ZeroAndFlushAll() {
 	device := peer.device
+	if peer.timers.sessionExpired != nil {
+		peer.timers.sessionExpired.Del()
+	}
 
 	// clear key pairs
 
@@ -238,6 +343,11 @@ func (peer *Peer) ZeroAndFlushAll() {
 	handshake.mutex.Unlock()
 
 	peer.FlushStagedPackets()
+
+	peer.sessionState.Lock()
+	peer.sessionState.sessionExpires = time.Time{}
+	peer.noteSessionStateLocked(PeerSessionNone)
+	peer.sessionState.Unlock()
 }
 
 func (peer *Peer) ExpireCurrentKeypairs() {
@@ -257,36 +367,181 @@ func (peer *Peer) ExpireCurrentKeypairs() {
 		next.sendNonce.Store(RejectAfterMessages)
 	}
 	keypairs.Unlock()
+
+	peer.sessionState.Lock()
+	peer.sessionState.sessionExpires = time.Time{}
+	peer.noteSessionStateLocked(PeerSessionExpired)
+	peer.sessionState.Unlock()
+}
+
+// queueOutboundIfRunning queues elems for outbound encryption & transmission to
+// peer if peer is running, otherwise elems are returned to their respective
+// pools for later re-use.
+func (peer *Peer) queueOutboundIfRunning(elems *QueueOutboundElementsContainer) (queued bool) {
+	if queued = peer.doIfRunning(func() {
+		elems.filling.Add(1)
+		peer.queue.outbound <- elems
+		peer.device.queue.encryption.c <- elems
+	}); !queued {
+		for _, elem := range elems.elems {
+			peer.device.PutOutboundElement(elem)
+		}
+		peer.device.PutOutboundElementsContainer(elems)
+	}
+	return queued
+}
+
+// queueInboundIfRunning queues elems for inbound decryption from peer if peer
+// is running, otherwise elems are returned to their respective pools for later
+// re-use.
+func (peer *Peer) queueInboundIfRunning(elems *QueueInboundElementsContainer) (queued bool) {
+	if queued = peer.doIfRunning(func() {
+		elems.filling.Add(1)
+		peer.queue.inbound <- elems
+		peer.device.queue.decryption.c <- elems
+	}); !queued {
+		for _, elem := range elems.elems {
+			peer.device.PutInboundElement(elem)
+		}
+		peer.device.PutInboundElementsContainer(elems)
+	}
+	return queued
+}
+
+// doIfRunning executes doQueueCrit if the [Peer] is running, i.e. the [Peer] is
+// not stopped or stopping. It returns true if doQueueCrit ran, otherwise it
+// returns false.
+//
+// doIfRunning exists to enable safe, bare (no select{}) write operations against
+// [Peer.queue].* channels.
+//
+// doIfRunning makes no guarantee about the running generation of the [Peer], i.e.
+// a [Peer] may stop then start all while a goroutine is within doIfRunning, but
+// before doQueueCrit has run.
+//
+// doIfRunning is built on the assumption that the supplied critical sections are
+// relatively cheap, as the mutually excluded peer-stopping side
+// ([Peer.waitQueueActors]) busy-waits for actors to exit.
+func (peer *Peer) doIfRunning(doQueueCrit func()) (ranCrit bool) {
+	if !peer.runningState.isRunning.Load() {
+		return false
+	}
+	peer.runningState.queueWriters.Add(1)
+	defer peer.runningState.queueWriters.Add(-1)
+	if peer.runningState.isRunning.Load() {
+		doQueueCrit()
+		return true
+	}
+	return false
+}
+
+// waitQueueActors waits for [Peer.runningState.queueWriters] and [Peer.runningState.queueReaders]
+// to reach zero. It should be called after [Peer.runningState.isRunning] is
+// swapped false.
+func (peer *Peer) waitQueueActors() {
+	drainInboundOutbound := func() {
+		for {
+			select {
+			case container := <-peer.queue.inbound:
+				container.filling.Wait()
+				for _, elem := range container.elems {
+					peer.device.PutInboundElement(elem)
+				}
+				peer.device.PutInboundElementsContainer(container)
+			case container := <-peer.queue.outbound:
+				container.filling.Wait()
+				for _, elem := range container.elems {
+					peer.device.PutOutboundElement(elem)
+				}
+				peer.device.PutOutboundElementsContainer(container)
+			default:
+				return
+			}
+		}
+	}
+
+	// Busy-wait for queue writers to return.
+	for peer.runningState.queueWriters.Load() != 0 {
+		// We drain at every iteration to give our busy-wait some meaningful
+		// work that can unblock actors performing channel writes. The goal at
+		// this point is to get them to return as soon as possible. It's ok if
+		// there are outstanding elements on packet queues once queueWriters
+		// reaches zero, our reader shutdown below will handle those stragglers.
+		drainInboundOutbound()
+		// We don't want our busy-wait starving actors, so explicitly yield back
+		// to the runtime, to bias towards actor goroutines getting CPU time.
+		runtime.Gosched()
+	}
+
+	// Signal to RoutineSequentialSender and RoutineSequentialReceiver to
+	// exit. Let them deal with any outstanding elements on the associated
+	// channels.
+	close(peer.queue.inbound)
+	close(peer.queue.outbound)
+	peer.runningState.queueReaders.Wait()
 }
 
 func (peer *Peer) Stop() {
 	peer.state.Lock()
 	defer peer.state.Unlock()
 
-	if !peer.isRunning.Swap(false) {
+	if !peer.runningState.isRunning.Swap(false) {
 		return
 	}
 
+	peer.waitQueueActors()
+
 	peer.device.log.Verbosef("%v - Stopping", peer)
 
 	peer.timersStop()
-	// Signal that RoutineSequentialSender and RoutineSequentialReceiver should exit.
-	peer.queue.inbound.c <- nil
-	peer.queue.outbound.c <- nil
-	peer.stopping.Wait()
 	peer.device.queue.encryption.wg.Done() // no more writes to encryption queue from us
 
 	peer.ZeroAndFlushAll()
 }
 
-func (peer *Peer) SetEndpointFromPacket(endpoint conn.Endpoint) {
-	peer.endpoint.Lock()
-	defer peer.endpoint.Unlock()
-	if peer.endpoint.disableRoaming {
+func (peer *Peer) noteSessionState(state PeerSessionState) {
+	peer.sessionState.Lock()
+	defer peer.sessionState.Unlock()
+	peer.noteSessionStateLocked(state)
+}
+
+// noteSessionStateLocked records a session state transition and delivers the
+// callback. The caller must hold peer.sessionState.Mutex during the
+// state determination and transition.
+func (peer *Peer) noteSessionStateLocked(state PeerSessionState) {
+	if peer.sessionState.current == state {
 		return
 	}
-	peer.endpoint.clearSrcOnTx = false
-	peer.endpoint.val = endpoint
+	peer.sessionState.current = state
+	if f := peer.device.peerStateFn.Load(); f != nil {
+		(*f)(peer.handshake.remoteStatic, state)
+	}
+}
+
+func (peer *Peer) noteSessionHandshakeStarted() {
+	peer.sessionState.Lock()
+	defer peer.sessionState.Unlock()
+	if peer.sessionState.current == PeerSessionEstablished {
+		return
+	}
+	peer.noteSessionStateLocked(PeerSessionHandshake)
+}
+
+func (peer *Peer) noteSessionHandshakeStopped() {
+	peer.sessionState.Lock()
+	defer peer.sessionState.Unlock()
+	state := PeerSessionNone
+	if peer.hasKeyMaterial() {
+		state = PeerSessionExpired
+	}
+	peer.noteSessionStateLocked(state)
+}
+
+func (peer *Peer) hasKeyMaterial() bool {
+	keypairs := &peer.keypairs
+	keypairs.RLock()
+	defer keypairs.RUnlock()
+	return keypairs.previous != nil || keypairs.current != nil || keypairs.next.Load() != nil
 }
 
 func (peer *Peer) markEndpointSrcForClearing() {