Update

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

@@ -6,6 +6,8 @@
 package device
 
 import (
+	"errors"
+	"net/netip"
 	"runtime"
 	"sync"
 	"sync/atomic"
@@ -18,6 +20,8 @@ import (
 )
 
 type Device struct {
+	config config
+
 	state struct {
 		// state holds the device's state. It is accessed atomically.
 		// Use the device.deviceState method to read it.
@@ -44,7 +48,6 @@ type Device struct {
 		netlinkCancel *rwcancel.RWCancel
 		port          uint16 // listening port
 		fwmark        uint32 // mark value (0 = disabled)
-		brokenRoaming bool
 	}
 
 	staticIdentity struct {
@@ -56,8 +59,12 @@ type Device struct {
 	peers struct {
 		sync.RWMutex // protects keyMap
 		keyMap       map[NoisePublicKey]*Peer
+		lookupFunc   PeerLookupFunc // or nil if unused
 	}
 
+	peerStateFn   atomic.Pointer[PeerSessionStateFunc]    // observes peer session state changes, nil if unset
+	priorityMsgFn atomic.Pointer[PeerPriorityMessageFunc] // returns a priority message to be sent around session establishment, nil if unset
+
 	rate struct {
 		underLoadUntil atomic.Int64
 		limiter        ratelimiter.Ratelimiter
@@ -70,9 +77,12 @@ type Device struct {
 	pool struct {
 		inboundElementsContainer  *WaitPool
 		outboundElementsContainer *WaitPool
-		messageBuffers            *WaitPool
 		inboundElements           *WaitPool
 		outboundElements          *WaitPool
+		packetBufs                *WaitPool
+		// smallPacketBufs aliases packetBufs when no distinct small packet pool
+		// is used. See [Device.PopulatePools].
+		smallPacketBufs *WaitPool
 	}
 
 	queue struct {
@@ -82,8 +92,12 @@ type Device struct {
 	}
 
 	tun struct {
-		device tun.Device
-		mtu    atomic.Int32
+		// Fields in this block are read-only after instantiation.
+		device  tun.Device
+		queues  []tun.Queue // device's read queues, at least one, see [tun.QueuesOf]
+		writeTo func(flow int, bufs [][]byte, offset int) (int, error)
+
+		mtu atomic.Int32
 	}
 
 	ipcMutex sync.RWMutex
@@ -91,6 +105,152 @@ type Device struct {
 	log      *Logger
 }
 
+type config struct {
+	queueStagedSize            int
+	queueOutboundSize          int
+	queueInboundSize           int
+	queueHandshakeSize         int
+	preallocatedBuffersPerPool uint32
+	metrics                    Metrics
+}
+
+func defaultConfig() config {
+	return config{
+		queueStagedSize:            DefaultQueueStagedSize,
+		queueOutboundSize:          DefaultQueueOutboundSize,
+		queueInboundSize:           DefaultQueueInboundSize,
+		queueHandshakeSize:         DefaultQueueHandshakeSize,
+		preallocatedBuffersPerPool: DefaultPreallocatedBuffersPerPool,
+	}
+}
+
+// An Option configures a [Device].
+type Option interface {
+	apply(*config)
+}
+
+type optionFunc func(*config)
+
+func (f optionFunc) apply(config *config) {
+	f(config)
+}
+
+// WithQueueStagedSize sets the capacity of each peer's staged packet queue.
+// Staged packet queues must be buffered, so max(tun readers, size) is applied
+// to the user-supplied value. [DefaultQueueStagedSize] is the default.
+func WithQueueStagedSize(size int) Option {
+	return optionFunc(func(config *config) {
+		config.queueStagedSize = max(1, size)
+	})
+}
+
+// WithQueueOutboundSize sets the capacity of each peer's outbound packet queue.
+// It also sets the capacity of the device-wide outbound queue, whose total
+// capacity is size multiplied by the number of TUN queues.
+// [DefaultQueueOutboundSize] is the default.
+func WithQueueOutboundSize(size int) Option {
+	return optionFunc(func(config *config) {
+		config.queueOutboundSize = size
+	})
+}
+
+// WithQueueInboundSize sets the capacity of the device and per-peer inbound
+// packet queues. [DefaultQueueInboundSize] is the default.
+func WithQueueInboundSize(size int) Option {
+	return optionFunc(func(config *config) {
+		config.queueInboundSize = size
+	})
+}
+
+// WithQueueHandshakeSize sets the capacity of the device's handshake queue.
+// [DefaultQueueHandshakeSize] is the default.
+func WithQueueHandshakeSize(size int) Option {
+	return optionFunc(func(config *config) {
+		config.queueHandshakeSize = size
+	})
+}
+
+// WithPreallocatedBuffersPerPool sets the maximum number of packet memory
+// pool outstanding items. A value of zero is unlimited. Care must be taken to
+// not set a size that is too small, which can lead to immediate deadlock, or
+// increase the probability of deadlock once packets start flowing.
+// See tailscale/corp#46396. [DefaultPreallocatedBuffersPerPool] is the default.
+func WithPreallocatedBuffersPerPool(size uint32) Option {
+	return optionFunc(func(config *config) {
+		config.preallocatedBuffersPerPool = size
+	})
+}
+
+// Counter is a monotonically increasing counter used in [Metrics].
+type Counter interface {
+	// Add increments the Counter's value by n. n must not be negative. Add is
+	// called in performance-sensitive contexts, therefore it must be cheap. Add
+	// may be called concurrently.
+	Add(n int64)
+}
+
+// Metrics contains metrics that can be exported via [WithMetrics].
+type Metrics struct {
+	// MessageInitiationTXAttemptInitial counts non-retry handshake
+	// initiation messages passed to the network send path.
+	MessageInitiationTXAttemptInitial Counter
+	// MessageInitiationTXAttemptRetry counts retry handshake initiation
+	// messages passed to the network send path.
+	MessageInitiationTXAttemptRetry Counter
+	// MessageResponseTXAttempt counts handshake response messages passed to the
+	// network send path.
+	MessageResponseTXAttempt Counter
+	// MessageCookieReplyTXAttempt counts cookie reply messages passed to the
+	// network send path.
+	MessageCookieReplyTXAttempt Counter
+	// HandshakeInitiatorCompleted counts symmetric session establishment on the
+	// initiator side.
+	HandshakeInitiatorCompleted Counter
+	// HandshakeResponderCompleted counts responder-side handshake completions,
+	// observed when an authenticated transport packet confirms the new keypair.
+	HandshakeResponderCompleted Counter
+	// MessageTransportRXDroppedReplay counts inbound transport packets rejected by the
+	// replay filter.
+	MessageTransportRXDroppedReplay Counter
+}
+
+// noopCounter is a noop implementation of [Counter].
+type noopCounter struct {
+}
+
+func (noopCounter) Add(int64) {}
+
+func (m *Metrics) fillNils() {
+	if m.MessageInitiationTXAttemptInitial == nil {
+		m.MessageInitiationTXAttemptInitial = noopCounter{}
+	}
+	if m.MessageInitiationTXAttemptRetry == nil {
+		m.MessageInitiationTXAttemptRetry = noopCounter{}
+	}
+	if m.MessageResponseTXAttempt == nil {
+		m.MessageResponseTXAttempt = noopCounter{}
+	}
+	if m.MessageCookieReplyTXAttempt == nil {
+		m.MessageCookieReplyTXAttempt = noopCounter{}
+	}
+	if m.HandshakeInitiatorCompleted == nil {
+		m.HandshakeInitiatorCompleted = noopCounter{}
+	}
+	if m.HandshakeResponderCompleted == nil {
+		m.HandshakeResponderCompleted = noopCounter{}
+	}
+	if m.MessageTransportRXDroppedReplay == nil {
+		m.MessageTransportRXDroppedReplay = noopCounter{}
+	}
+}
+
+// WithMetrics sets the provided [Metrics] to be used at instrumentation points.
+func WithMetrics(metrics Metrics) Option {
+	return optionFunc(func(config *config) {
+		config.metrics = metrics
+	})
+}
+
 // deviceState represents the state of a Device.
 // There are three states: down, up, closed.
 // Transitions:
@@ -179,14 +339,22 @@ func (device *Device) upLocked() error {
 	device.ipcMutex.Lock()
 	defer device.ipcMutex.Unlock()
 
+	// Collect peers under RLock and then release before calling into them,
+	// because SendKeepalive can reach CreateMessageInitiation which acquires
+	// staticIdentity.RLock; holding peers.RLock across that path would
+	// invert the staticIdentity < peers hierarchy (see lock-ordering.md).
 	device.peers.RLock()
+	peers := make([]*Peer, 0, len(device.peers.keyMap))
 	for _, peer := range device.peers.keyMap {
+		peers = append(peers, peer)
+	}
+	device.peers.RUnlock()
+	for _, peer := range peers {
 		peer.Start()
 		if peer.persistentKeepaliveInterval.Load() > 0 {
 			peer.SendKeepalive()
 		}
 	}
-	device.peers.RUnlock()
 	return nil
 }
 
@@ -217,7 +385,9 @@ func (device *Device) Down() error {
 func (device *Device) IsUnderLoad() bool {
 	// check if currently under load
 	now := time.Now()
-	underLoad := len(device.queue.handshake.c) >= QueueHandshakeSize/8
+	// max(1, ...) is required on the right hand side, otherwise underLoad would
+	// always be true when queueHandshakeSize < 8.
+	underLoad := len(device.queue.handshake.c) >= max(1, device.config.queueHandshakeSize/8)
 	if underLoad {
 		device.rate.underLoadUntil.Store(now.Add(UnderLoadAfterTime).UnixNano())
 		return true
@@ -281,8 +451,12 @@ func (device *Device) SetPrivateKey(sk NoisePrivateKey) error {
 	return nil
 }
 
-func NewDevice(tunDevice tun.Device, bind conn.Bind, logger *Logger) *Device {
-	device := new(Device)
+func NewDevice(tunDevice tun.Device, bind conn.Bind, logger *Logger, opts ...Option) *Device {
+	device := &Device{config: defaultConfig()}
+	for _, opt := range opts {
+		opt.apply(&device.config)
+	}
+	device.config.metrics.fillNils()
 	device.state.state.Store(uint32(deviceStateDown))
 	device.closed = make(chan struct{})
 	device.log = logger
@@ -294,17 +468,31 @@ func NewDevice(tunDevice tun.Device, bind conn.Bind, logger *Logger) *Device {
 		mtu = DefaultMTU
 	}
 	device.tun.mtu.Store(int32(mtu))
+	device.tun.queues = tun.QueuesOf(tunDevice)
+	device.tun.writeTo = tun.WriteToOf(tunDevice)
 	device.peers.keyMap = make(map[NoisePublicKey]*Peer)
 	device.rate.limiter.Init()
 	device.indexTable.Init()
 
+	if want := len(device.tun.queues); want > device.config.queueStagedSize {
+		device.log.Errorf(
+			"Raising staged queue size to fit concurrent tun readers: %d -> %d",
+			device.config.queueStagedSize,
+			want,
+		)
+		device.config.queueStagedSize = want
+	}
+
 	device.PopulatePools()
 
 	// create queues
 
-	device.queue.handshake = newHandshakeQueue()
-	device.queue.encryption = newOutboundQueue()
-	device.queue.decryption = newInboundQueue()
+	device.queue.handshake = newHandshakeQueue(device.config.queueHandshakeSize)
+	// Scale to the number of producers
+	device.queue.encryption = newOutboundQueue(
+		device.config.queueOutboundSize * len(device.tun.queues),
+	)
+	device.queue.decryption = newInboundQueue(device.config.queueInboundSize)
 
 	// start workers
 
@@ -317,9 +505,11 @@ func NewDevice(tunDevice tun.Device, bind conn.Bind, logger *Logger) *Device {
 		go device.RoutineHandshake(i + 1)
 	}
 
-	device.state.stopping.Add(1)      // RoutineReadFromTUN
-	device.queue.encryption.wg.Add(1) // RoutineReadFromTUN
-	go device.RoutineReadFromTUN()
+	device.state.stopping.Add(len(device.tun.queues))      // RoutineReadFromTUN
+	device.queue.encryption.wg.Add(len(device.tun.queues)) // RoutineReadFromTUN
+	for i, q := range device.tun.queues {
+		go device.RoutineReadFromTUN(i, q)
+	}
 	go device.RoutineTUNEventReader()
 
 	return device
@@ -338,13 +528,69 @@ 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]
+	conf, ok := lookupFunc(pk)
+	if !ok || conf == 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.SetPresharedKey(conf.PresharedKey)
+	p.SetAllowedIPs(conf.AllowedIPs)
+	p.deleteOnIdle = true
+	if conf.Endpoint != nil {
+		p.endpoint.Lock()
+		p.endpoint.val = conf.Endpoint
+		p.endpoint.Unlock()
+	}
+	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()
@@ -367,6 +613,175 @@ 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
+}
+
+// NewPeerConfig are the configuration parameters for a new peer created via a
+// [PeerLookupFunc] func.
+type NewPeerConfig struct {
+	// AllowedIPs is the initial set of allowed IPs for the new peer.
+	AllowedIPs []netip.Prefix
+
+	// PresharedKey is the initial pre-shared key for the new peer. The zero
+	// value disables the optional WireGuard pre-shared-key layer.
+	PresharedKey NoisePresharedKey
+
+	// Endpoint, if non-nil, sets the endpoint for newly created peers.
+	// The endpoint is pinned for the lifetime of the peer.
+	Endpoint conn.Endpoint
+}
+
+// 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) (_ *NewPeerConfig, ok bool)
+
+// PeerByIPPacketFunc is the type of function used to look up a peer to send to
+// for a given src/dst IP pair. The ipPkt parameter is the raw IP packet being
+// routed; callers needing transport-layer ports or other header fields may parse
+// them from ipPkt, but must handle IP fragmentation (ports may be absent on
+// non-first fragments) and protocols that do not use ports (e.g. ICMP).
+//
+// Except for experimental use cases, dst is the only address
+// that should be relied upon when looking up a peer.
+//
+// If it returns ok=false, the peer is not known.
+//
+// See [Device.SetPeerByIPPacketFunc] and [Device.SetPeerLookupFunc].
+type PeerByIPPacketFunc func(src, dst netip.Addr, ipPkt []byte) (_ NoisePublicKey, ok bool)
+
+// PeerSessionState is the current WireGuard session state for a peer.
+type PeerSessionState uint8
+
+const (
+	// PeerSessionNone means there is no handshake in progress and no session key
+	// material retained for this peer.
+	PeerSessionNone PeerSessionState = iota
+
+	// PeerSessionHandshake means a handshake is in progress for this peer, but
+	// there is not currently a usable WireGuard session.
+	PeerSessionHandshake
+
+	// PeerSessionEstablished means the peer has a completed WireGuard session
+	// with usable session key material.
+	PeerSessionEstablished
+
+	// PeerSessionExpired means the peer's session key material is no longer
+	// considered usable, but final key cleanup or lazy peer removal may not have
+	// happened yet.
+	PeerSessionExpired
+)
+
+// PeerSessionStateFunc is called when a peer's WireGuard session state changes.
+//
+// Calls are serialized per peer and delivered in that peer's transition order. The
+// callback must be cheap and must not call back into Device.
+type PeerSessionStateFunc func(peer NoisePublicKey, state PeerSessionState)
+
+// 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
+}
+
+// SetPeerByIPPacketFunc sets the function used to look up peers by IP address
+// when sending packets to unknown peers.
+func (device *Device) SetPeerByIPPacketFunc(f PeerByIPPacketFunc) {
+	device.allowedips.mu.Lock()
+	defer device.allowedips.mu.Unlock()
+	device.allowedips.peerByIPPacketFunc = f
+	device.allowedips.device = device
+}
+
+// SetSessionStateFunc sets the function used to observe peer WireGuard session
+// state changes.
+//
+// It does not replay current state. Callers that need a complete view should set
+// it before peers are started or lazily created, and maintain any snapshots,
+// sequence numbers, and pubsub state outside wireguard-go.
+//
+// The callback must be concurrent-safe and must not call back into Device.
+func (device *Device) SetSessionStateFunc(f PeerSessionStateFunc) {
+	if f == nil {
+		device.peerStateFn.Store(nil)
+		return
+	}
+	device.peerStateFn.Store(&f)
+}
+
+// MaxPriorityMessageContentSize is the maximum size of a message returned by a
+// [PeerPriorityMessageFunc]. It's a power of 2 that leaves significant space
+// when accounting for all WireGuard overhead and encapsulating network protocol
+// headers. Future adjustments to this value should consider all these overheads
+// and any [conn.Bind] implementation limitations.
+const MaxPriorityMessageContentSize = 512
+
+// PeerPriorityMessageFunc is called when a peer's WireGuard session keypair is
+// established (or re-keyed) for forward data transmission.
+//
+// The returned message is transmitted to the peer in priority fashion. Priority
+// means it cannot be evicted from the staged packet queue by non-priority
+// (read from [tun.Device]) packets. It avoids the staged queue altogether.
+//
+// The callback must be cheap and must not call back into [Device]. A zero length
+// message or a message whose length exceeds [MaxPriorityMessageContentSize] will
+// be silently dropped. Message should start with an IPv4 or IPv6 header as it
+// is subject to allowed IPs lookup on the receiver, same as any other transport
+// message.
+type PeerPriorityMessageFunc func(peer NoisePublicKey) (msg []byte)
+
+// SetPriorityMessageOnEstablishmentFunc sets a function to be used for sending
+// a priority message around session establishment. See [PeerPriorityMessageFunc]
+// docs for more details. A nil value clears any previously set value.
+func (device *Device) SetPriorityMessageOnEstablishmentFunc(f PeerPriorityMessageFunc) {
+	if f == nil {
+		device.priorityMsgFn.Store(nil)
+		return
+	}
+	device.priorityMsgFn.Store(&f)
+}
+
+// ScheduleHandshakeOnUserSend marks the peer for handshake initiation immediately
+// following its next outbound transport message. It is a no-op if the peer is
+// unknown or currently holds no key material.
+//
+// The initiation remains subject to [RekeyTimeout] elapsing since the last
+// handshake message was sent. A request blocked by that window stays armed for
+// a later transport message rather than being dropped.
+func (device *Device) ScheduleHandshakeOnUserSend(peer NoisePublicKey) {
+	device.peers.RLock()
+	defer device.peers.RUnlock()
+	p, ok := device.peers.keyMap[peer]
+	if !ok {
+		return
+	}
+	if !p.hasKeyMaterial() {
+		return
+	}
+	p.handshakeOnUserSend.Store(true)
+}
+
 func (device *Device) Close() {
 	device.state.Lock()
 	defer device.state.Unlock()
@@ -408,16 +823,25 @@ func (device *Device) SendKeepalivesToPeersWithCurrentKeypair() {
 		return
 	}
 
+	// Collect the set of peers to keepalive under peers.RLock, then release
+	// before invoking SendKeepalive. SendKeepalive can reach
+	// CreateMessageInitiation which acquires staticIdentity.RLock; holding
+	// peers.RLock across that path would invert the
+	// staticIdentity < peers hierarchy (see lock-ordering.md).
+	var peers []*Peer
 	device.peers.RLock()
 	for _, peer := range device.peers.keyMap {
 		peer.keypairs.RLock()
 		sendKeepalive := peer.keypairs.current != nil && !peer.keypairs.current.created.Add(RejectAfterTime).Before(time.Now())
 		peer.keypairs.RUnlock()
 		if sendKeepalive {
-			peer.SendKeepalive()
+			peers = append(peers, peer)
 		}
 	}
 	device.peers.RUnlock()
+	for _, peer := range peers {
+		peer.SendKeepalive()
+	}
 }
 
 // closeBindLocked closes the device's net.bind.