Update

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

@@ -0,0 +1,41 @@
+//go:build amd64 && gc && !purego
+
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+package device
+
+import (
+	"crypto/cipher"
+	"os"
+
+	"golang.org/x/crypto/chacha20poly1305"
+	"golang.org/x/sys/cpu"
+
+	asmAEAD "github.com/tailscale/wireguard-go/tsasm/amd64/chacha20poly1305"
+)
+
+/*
+shouldUseTSAsm reports whether the data-path AEAD should come from tsasm/amd64. x/crypto gates its remaining amd64 kernel on HasSSSE3 && HasAVX2 && HasBMI2 and otherwise runs generic Go, so only CPUs with SSSE3 and no AVX2 are routed here. Pure so a test can cover every combination; the CPU variables cannot be written.
+*/
+func shouldUseTSAsm(hasSSSE3, hasAVX2, hasBMI2 bool) bool {
+	return hasSSSE3 && !(hasAVX2 && hasBMI2)
+}
+
+var useTSAsm = shouldUseTSAsm(asmAEAD.Available(), cpu.X86.HasAVX2, cpu.X86.HasBMI2)
+
+/*
+chacha20poly1305New returns a ChaCha20-Poly1305 AEAD. On amd64 CPUs with SSSE3 but no AVX2 it uses the assembly kernel from tsasm/amd64/chacha20poly1305.
+
+The cookie path (which uses the extended-nonce variant via chacha20poly1305.NewX) is left on the x/crypto path because it is not on the per-packet hot path.
+
+As an escape hatch for hardware regressions or asm bugs, setting the environment variable TS_WG_ASM=0 forces the x/crypto implementation instead.
+*/
+func chacha20poly1305New(key []byte) (cipher.AEAD, error) {
+	if !useTSAsm || os.Getenv("TS_WG_ASM") == "0" {
+		return chacha20poly1305.New(key)
+	}
+	return asmAEAD.New(key)
+}

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

@@ -0,0 +1,33 @@
+//go:build arm
+
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+package device
+
+import (
+	"crypto/cipher"
+	"os"
+
+	"golang.org/x/crypto/chacha20poly1305"
+
+	asmAEAD "github.com/tailscale/wireguard-go/tsasm/arm/chacha20poly1305"
+)
+
+// chacha20poly1305New returns a ChaCha20-Poly1305 AEAD. On GOARCH=arm
+// (32-bit, any GOOS) it uses the assembly kernels from tsasm/arm/.
+// The cookie path (which uses the extended-nonce variant via
+// chacha20poly1305.NewX) is left on the x/crypto path because it is
+// not on the per-packet hot path.
+//
+// As an escape hatch for hardware regressions or asm bugs, setting
+// the environment variable TS_WG_ASM=0 forces the pure-Go x/crypto
+// implementation instead.
+func chacha20poly1305New(key []byte) (cipher.AEAD, error) {
+	if os.Getenv("TS_WG_ASM") == "0" {
+		return chacha20poly1305.New(key)
+	}
+	return asmAEAD.New(key)
+}

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

@@ -0,0 +1,18 @@
+//go:build !arm && !(amd64 && gc && !purego)
+
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+package device
+
+import (
+	"crypto/cipher"
+
+	"golang.org/x/crypto/chacha20poly1305"
+)
+
+func chacha20poly1305New(key []byte) (cipher.AEAD, error) {
+	return chacha20poly1305.New(key)
+}

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

@@ -55,6 +55,25 @@ func commonBits(ip1, ip2 []byte) uint8 {
 	}
 }
 
+func commonBits4(ip1 []byte, ip2 [4]byte) uint8 {
+	a := binary.BigEndian.Uint32(ip1)
+	b := binary.BigEndian.Uint32(ip2[:])
+	return uint8(bits.LeadingZeros32(a ^ b))
+}
+
+func commonBits6(ip1 []byte, ip2 [16]byte) uint8 {
+	a := binary.BigEndian.Uint64(ip1)
+	b := binary.BigEndian.Uint64(ip2[:])
+	x := a ^ b
+	if x != 0 {
+		return uint8(bits.LeadingZeros64(x))
+	}
+	a = binary.BigEndian.Uint64(ip1[8:])
+	b = binary.BigEndian.Uint64(ip2[8:])
+	x = a ^ b
+	return 64 + uint8(bits.LeadingZeros64(x))
+}
+
 func (node *trieEntry) addToPeerEntries() {
 	node.perPeerElem = node.peer.trieEntries.PushBack(node)
 }
@@ -188,7 +207,37 @@ func (trie parentIndirection) insert(ip []byte, cidr uint8, peer *Peer) {
 	}
 }
 
-func (node *trieEntry) lookup(ip []byte) *Peer {
+func (node *trieEntry) lookup4(ip [4]byte) *Peer {
+	var found *Peer
+	for node != nil && commonBits4(node.bits, ip) >= node.cidr {
+		if node.peer != nil {
+			found = node.peer
+		}
+		if node.bitAtByte == 4 {
+			break
+		}
+		bit := (ip[node.bitAtByte] >> node.bitAtShift) & 1
+		node = node.child[bit]
+	}
+	return found
+}
+
+func (node *trieEntry) lookup6(ip [16]byte) *Peer {
+	var found *Peer
+	for node != nil && commonBits6(node.bits, ip) >= node.cidr {
+		if node.peer != nil {
+			found = node.peer
+		}
+		if node.bitAtByte == 16 {
+			break
+		}
+		bit := (ip[node.bitAtByte] >> node.bitAtShift) & 1
+		node = node.child[bit]
+	}
+	return found
+}
+
+func (node *trieEntry) lookup(ip net.IP) *Peer {
 	var found *Peer
 	size := uint8(len(ip))
 	for node != nil && commonBits(node.bits, ip) >= node.cidr {
@@ -205,14 +254,17 @@ func (node *trieEntry) lookup(ip []byte) *Peer {
 }
 
 type AllowedIPs struct {
-	IPv4  *trieEntry
-	IPv6  *trieEntry
-	mutex sync.RWMutex
+	mu   sync.RWMutex
+	ipv4 *trieEntry
+	ipv6 *trieEntry
+
+	peerByIPPacketFunc PeerByIPPacketFunc // if non-nil, called to look up peers by IP
+	device             *Device            // back-reference to parent device; non-nil only if peerByIPPacketFunc is set
 }
 
 func (table *AllowedIPs) EntriesForPeer(peer *Peer, cb func(prefix netip.Prefix) bool) {
-	table.mutex.RLock()
-	defer table.mutex.RUnlock()
+	table.mu.RLock()
+	defer table.mu.RUnlock()
 
 	for elem := peer.trieEntries.Front(); elem != nil; elem = elem.Next() {
 		node := elem.Value.(*trieEntry)
@@ -223,10 +275,25 @@ func (table *AllowedIPs) EntriesForPeer(peer *Peer, cb func(prefix netip.Prefix)
 	}
 }
 
-func (table *AllowedIPs) RemoveByPeer(peer *Peer) {
-	table.mutex.Lock()
-	defer table.mutex.Unlock()
+// 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()
@@ -266,29 +333,139 @@ func (table *AllowedIPs) RemoveByPeer(peer *Peer) {
 }
 
 func (table *AllowedIPs) Insert(prefix netip.Prefix, peer *Peer) {
-	table.mutex.Lock()
-	defer table.mutex.Unlock()
+	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)
+		parentIndirection{&table.ipv6, 2}.insert(ip[:], uint8(prefix.Bits()), peer)
 	} else if prefix.Addr().Is4() {
 		ip := prefix.Addr().As4()
-		parentIndirection{&table.IPv4, 2}.insert(ip[:], uint8(prefix.Bits()), peer)
+		parentIndirection{&table.ipv4, 2}.insert(ip[:], uint8(prefix.Bits()), peer)
 	} else {
 		panic(errors.New("inserting unknown address type"))
 	}
 }
 
-func (table *AllowedIPs) Lookup(ip []byte) *Peer {
-	table.mutex.RLock()
-	defer table.mutex.RUnlock()
-	switch len(ip) {
-	case net.IPv6len:
-		return table.IPv6.lookup(ip)
-	case net.IPv4len:
-		return table.IPv4.lookup(ip)
+// LookupFromPacket looks up the peer to which an outbound IP packet should be
+// sent. It lives on [AllowedIPs] for legacy/structural reasons: historically
+// WireGuard's only peer-selection mechanism was the AllowedIPs trie, and the
+// send path already had a reference to the table. When a [PeerByIPPacketFunc]
+// has been registered via [Device.SetPeerByIPPacketFunc], that callback is used
+// instead of the trie and the AllowedIPs table is not consulted at all.
+//
+// When no callback is registered, only dst is used (standard WireGuard
+// AllowedIPs trie lookup). When a callback is registered, all three
+// parameters are forwarded to it; see [PeerByIPPacketFunc] for details.
+func (table *AllowedIPs) LookupFromPacket(src, dst netip.Addr, ipPkt []byte) *Peer {
+	table.mu.RLock()
+	if f := table.peerByIPPacketFunc; f != nil {
+		device := table.device
+		table.mu.RUnlock()
+
+		if pubk, ok := f(src, dst, ipPkt); ok {
+			return device.LookupPeer(pubk)
+		}
+		return nil
+	}
+	defer table.mu.RUnlock()
+
+	switch {
+	case dst.Is6():
+		return table.ipv6.lookup6(dst.As16())
+	case dst.Is4():
+		return table.ipv4.lookup4(dst.As4())
 	default:
 		panic(errors.New("looking up unknown address type"))
 	}
 }
+
+// Deprecated: Lookup is only used by legacy tests. It does not call
+// [PeerByIPPacketFunc]; use [AllowedIPs.LookupFromPacket] for production lookups.
+func (table *AllowedIPs) Lookup(ip []byte) *Peer {
+	table.mu.RLock()
+	defer table.mu.RUnlock()
+	return table.lookupLocked(ip)
+}
+
+// lookupLocked looks up the peer associated with the given IP address.
+// It assumes the caller holds the read lock (or doesn't hold it, but also
+// doesn't concurrently mutate AllowedIP).
+//
+// It returns nil if no peer is associated with the given IP address.
+func (table *AllowedIPs) lookupLocked(ip []byte) *Peer {
+	switch len(ip) {
+	case net.IPv6len:
+		return table.ipv6.lookup(ip)
+	case net.IPv4len:
+		return table.ipv4.lookup(ip)
+	default:
+		panic(errors.New("looking up unknown address type"))
+	}
+}
+
+// AllowedPeerSourceIP reports whether the given source IP address is allowed
+// for the given peer.
+func (peer *Peer) AllowedPeerSourceIP(src netip.Addr) bool {
+	if f := peer.state.testAllowedIP.Load(); f != nil {
+		return (*f)(src)
+	}
+
+	table := &peer.device.allowedips
+	table.mu.RLock()
+	defer table.mu.RUnlock()
+	switch {
+	case src.Is6():
+		return table.ipv6.lookup6(src.As16()) == peer
+	case src.Is4():
+		return table.ipv4.lookup4(src.As4()) == peer
+	}
+	return false
+}
+
+// cidrLinearSearchMax is the largest number of CIDRs for which
+// mkIPInCIDRsTestFunc does a linear search instead of building a trie.
+const cidrLinearSearchMax = 4
+
+// mkIPInCIDRsTestFunc returns a function that tests whether an IP address is
+// contained in any of the given CIDRs.
+func mkIPInCIDRsTestFunc(cidrs []netip.Prefix) func(netip.Addr) bool {
+	if len(cidrs) == 0 {
+		return func(netip.Addr) bool { return false }
+	}
+	if len(cidrs) == 1 {
+		return func(addr netip.Addr) bool { return cidrs[0].Contains(addr) }
+	}
+	if len(cidrs) <= cidrLinearSearchMax {
+		// For small numbers of CIDRs, just do a linear search. The trie construction
+		// is more expensive than the linear search, and the test function is faster
+		// than the trie lookup, so this is a net win.
+		return func(addr netip.Addr) bool {
+			for _, c := range cidrs {
+				if c.Contains(addr) {
+					return true
+				}
+			}
+			return false
+		}
+	}
+	// Make a trie for faster lookups. Use a local dummy Peer so its
+	// trieEntries and the whole trie can be collected once the returned
+	// closure is gone.
+	fakePeer := new(Peer)
+	var a AllowedIPs
+	for _, c := range cidrs {
+		a.Insert(c, fakePeer)
+	}
+	return func(addr netip.Addr) bool {
+		switch {
+		case addr.Is4():
+			return a.ipv4.lookup4(addr.As4()) == fakePeer
+		default:
+			return a.ipv6.lookup6(addr.As16()) == fakePeer
+		}
+	}
+}

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

@@ -6,7 +6,6 @@
 package device
 
 import (
-	"runtime"
 	"sync"
 )
 
@@ -23,9 +22,9 @@ type outboundQueue struct {
 	wg sync.WaitGroup
 }
 
-func newOutboundQueue() *outboundQueue {
+func newOutboundQueue(capacity int) *outboundQueue {
 	q := &outboundQueue{
-		c: make(chan *QueueOutboundElementsContainer, QueueOutboundSize),
+		c: make(chan *QueueOutboundElementsContainer, capacity),
 	}
 	q.wg.Add(1)
 	go func() {
@@ -41,9 +40,9 @@ type inboundQueue struct {
 	wg sync.WaitGroup
 }
 
-func newInboundQueue() *inboundQueue {
+func newInboundQueue(capacity int) *inboundQueue {
 	q := &inboundQueue{
-		c: make(chan *QueueInboundElementsContainer, QueueInboundSize),
+		c: make(chan *QueueInboundElementsContainer, capacity),
 	}
 	q.wg.Add(1)
 	go func() {
@@ -59,9 +58,9 @@ type handshakeQueue struct {
 	wg sync.WaitGroup
 }
 
-func newHandshakeQueue() *handshakeQueue {
+func newHandshakeQueue(capacity int) *handshakeQueue {
 	q := &handshakeQueue{
-		c: make(chan QueueHandshakeElement, QueueHandshakeSize),
+		c: make(chan QueueHandshakeElement, capacity),
 	}
 	q.wg.Add(1)
 	go func() {
@@ -70,68 +69,3 @@ func newHandshakeQueue() *handshakeQueue {
 	}()
 	return q
 }
-
-type autodrainingInboundQueue struct {
-	c chan *QueueInboundElementsContainer
-}
-
-// newAutodrainingInboundQueue returns a channel that will be drained when it gets GC'd.
-// It is useful in cases in which is it hard to manage the lifetime of the channel.
-// The returned channel must not be closed. Senders should signal shutdown using
-// some other means, such as sending a sentinel nil values.
-func newAutodrainingInboundQueue(device *Device) *autodrainingInboundQueue {
-	q := &autodrainingInboundQueue{
-		c: make(chan *QueueInboundElementsContainer, QueueInboundSize),
-	}
-	runtime.SetFinalizer(q, device.flushInboundQueue)
-	return q
-}
-
-func (device *Device) flushInboundQueue(q *autodrainingInboundQueue) {
-	for {
-		select {
-		case elemsContainer := <-q.c:
-			elemsContainer.Lock()
-			for _, elem := range elemsContainer.elems {
-				device.PutMessageBuffer(elem.buffer)
-				device.PutInboundElement(elem)
-			}
-			device.PutInboundElementsContainer(elemsContainer)
-		default:
-			return
-		}
-	}
-}
-
-type autodrainingOutboundQueue struct {
-	c chan *QueueOutboundElementsContainer
-}
-
-// newAutodrainingOutboundQueue returns a channel that will be drained when it gets GC'd.
-// It is useful in cases in which is it hard to manage the lifetime of the channel.
-// The returned channel must not be closed. Senders should signal shutdown using
-// some other means, such as sending a sentinel nil values.
-// All sends to the channel must be best-effort, because there may be no receivers.
-func newAutodrainingOutboundQueue(device *Device) *autodrainingOutboundQueue {
-	q := &autodrainingOutboundQueue{
-		c: make(chan *QueueOutboundElementsContainer, QueueOutboundSize),
-	}
-	runtime.SetFinalizer(q, device.flushOutboundQueue)
-	return q
-}
-
-func (device *Device) flushOutboundQueue(q *autodrainingOutboundQueue) {
-	for {
-		select {
-		case elemsContainer := <-q.c:
-			elemsContainer.Lock()
-			for _, elem := range elemsContainer.elems {
-				device.PutMessageBuffer(elem.buffer)
-				device.PutOutboundElement(elem)
-			}
-			device.PutOutboundElementsContainer(elemsContainer)
-		default:
-			return
-		}
-	}
-}

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.

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

@@ -1,19 +0,0 @@
-/* SPDX-License-Identifier: MIT
- *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
- */
-
-package device
-
-// DisableSomeRoamingForBrokenMobileSemantics should ideally be called before peers are created,
-// though it will try to deal with it, and race maybe, if called after.
-func (device *Device) DisableSomeRoamingForBrokenMobileSemantics() {
-	device.net.brokenRoaming = true
-	device.peers.RLock()
-	for _, peer := range device.peers.keyMap {
-		peer.endpoint.Lock()
-		peer.endpoint.disableRoaming = peer.endpoint.val != nil
-		peer.endpoint.Unlock()
-	}
-	device.peers.RUnlock()
-}

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

@@ -308,7 +308,7 @@ func (device *Device) CreateMessageInitiation(peer *Peer) (*MessageInitiation, e
 		handshake.chainKey[:],
 		ss[:],
 	)
-	aead, _ := chacha20poly1305.New(key[:])
+	aead, _ := chacha20poly1305New(key[:])
 	aead.Seal(msg.Static[:0], ZeroNonce[:], device.staticIdentity.publicKey[:], handshake.hash[:])
 	handshake.mixHash(msg.Static[:])
 
@@ -323,7 +323,7 @@ func (device *Device) CreateMessageInitiation(peer *Peer) (*MessageInitiation, e
 		handshake.precomputedStaticStatic[:],
 	)
 	timestamp := tai64n.Now()
-	aead, _ = chacha20poly1305.New(key[:])
+	aead, _ = chacha20poly1305New(key[:])
 	aead.Seal(msg.Timestamp[:0], ZeroNonce[:], timestamp[:], handshake.hash[:])
 
 	// assign index
@@ -349,22 +349,27 @@ func (device *Device) ConsumeMessageInitiation(msg *MessageInitiation, endpoint
 		return nil
 	}
 
+	// Snapshot staticIdentity so we don't hold the RLock across LookupPeer,
+	// which may call NewPeer (reentrant RLock deadlocks against a pending
+	// SetPrivateKey writer; see lock-ordering.md).
 	device.staticIdentity.RLock()
-	defer device.staticIdentity.RUnlock()
+	publicKey := device.staticIdentity.publicKey
+	privateKey := device.staticIdentity.privateKey
+	device.staticIdentity.RUnlock()
 
-	mixHash(&hash, &InitialHash, device.staticIdentity.publicKey[:])
+	mixHash(&hash, &InitialHash, publicKey[:])
 	mixHash(&hash, &hash, msg.Ephemeral[:])
 	mixKey(&chainKey, &InitialChainKey, msg.Ephemeral[:])
 
 	// decrypt static key
 	var peerPK NoisePublicKey
 	var key [chacha20poly1305.KeySize]byte
-	ss, err := device.staticIdentity.privateKey.sharedSecret(msg.Ephemeral)
+	ss, err := privateKey.sharedSecret(msg.Ephemeral)
 	if err != nil {
 		return nil
 	}
 	KDF2(&chainKey, &key, chainKey[:], ss[:])
-	aead, _ := chacha20poly1305.New(key[:])
+	aead, _ := chacha20poly1305New(key[:])
 	_, err = aead.Open(peerPK[:0], ZeroNonce[:], msg.Static[:], hash[:])
 	if err != nil {
 		return nil
@@ -379,7 +384,7 @@ func (device *Device) ConsumeMessageInitiation(msg *MessageInitiation, endpoint
 	}
 
 	peer := device.LookupPeer(peerPK)
-	if peer == nil || !peer.isRunning.Load() {
+	if peer == nil || !peer.runningState.isRunning.Load() {
 		return nil
 	}
 
@@ -401,7 +406,7 @@ func (device *Device) ConsumeMessageInitiation(msg *MessageInitiation, endpoint
 		chainKey[:],
 		handshake.precomputedStaticStatic[:],
 	)
-	aead, _ = chacha20poly1305.New(key[:])
+	aead, _ = chacha20poly1305New(key[:])
 	_, err = aead.Open(timestamp[:0], ZeroNonce[:], msg.Timestamp[:], hash[:])
 	if err != nil {
 		handshake.mutex.RUnlock()
@@ -507,7 +512,7 @@ func (device *Device) CreateMessageResponse(peer *Peer) (*MessageResponse, error
 
 	handshake.mixHash(tau[:])
 
-	aead, _ := chacha20poly1305.New(key[:])
+	aead, _ := chacha20poly1305New(key[:])
 	aead.Seal(msg.Empty[:0], ZeroNonce[:], nil, handshake.hash[:])
 	handshake.mixHash(msg.Empty[:])
 
@@ -534,6 +539,14 @@ func (device *Device) ConsumeMessageResponse(msg *MessageResponse) *Peer {
 		chainKey [blake2s.Size]byte
 	)
 
+	// Snapshot the static private key before acquiring handshake.mutex so
+	// that handshake.mutex is never held while acquiring staticIdentity
+	// (which would invert the staticIdentity < handshake.mutex hierarchy;
+	// see lock-ordering.md).
+	device.staticIdentity.RLock()
+	privateKey := device.staticIdentity.privateKey
+	device.staticIdentity.RUnlock()
+
 	ok := func() bool {
 		// lock handshake state
 
@@ -544,11 +557,6 @@ func (device *Device) ConsumeMessageResponse(msg *MessageResponse) *Peer {
 			return false
 		}
 
-		// lock private key for reading
-
-		device.staticIdentity.RLock()
-		defer device.staticIdentity.RUnlock()
-
 		// finish 3-way DH
 
 		mixHash(&hash, &handshake.hash, msg.Ephemeral[:])
@@ -561,7 +569,7 @@ func (device *Device) ConsumeMessageResponse(msg *MessageResponse) *Peer {
 		mixKey(&chainKey, &chainKey, ss[:])
 		setZero(ss[:])
 
-		ss, err = device.staticIdentity.privateKey.sharedSecret(msg.Ephemeral)
+		ss, err = privateKey.sharedSecret(msg.Ephemeral)
 		if err != nil {
 			return false
 		}
@@ -583,7 +591,7 @@ func (device *Device) ConsumeMessageResponse(msg *MessageResponse) *Peer {
 
 		// authenticate transcript
 
-		aead, _ := chacha20poly1305.New(key[:])
+		aead, _ := chacha20poly1305New(key[:])
 		_, err = aead.Open(nil, ZeroNonce[:], msg.Empty[:], hash[:])
 		if err != nil {
 			return false
@@ -658,8 +666,8 @@ func (peer *Peer) BeginSymmetricSession() error {
 	// create AEAD instances
 
 	keypair := new(Keypair)
-	keypair.send, _ = chacha20poly1305.New(sendKey[:])
-	keypair.receive, _ = chacha20poly1305.New(recvKey[:])
+	keypair.send, _ = chacha20poly1305New(sendKey[:])
+	keypair.receive, _ = chacha20poly1305New(recvKey[:])
 
 	setZero(sendKey[:])
 	setZero(recvKey[:])

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

@@ -0,0 +1,66 @@
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+package device
+
+import (
+	"fmt"
+	"sync/atomic"
+
+	"github.com/tailscale/wireguard-go/tun"
+	"golang.org/x/crypto/poly1305"
+)
+
+const (
+	outboundPlaintextHeadroom = MessageEncapsulatingTransportSize + MessageTransportHeaderSize
+
+	outboundPlaintextTailroom = PaddingMultiple - 1 + poly1305.TagSize
+
+	outboundPlaintextSpacing = max(outboundPlaintextHeadroom+outboundPlaintextTailroom, tun.ReadPacketSpacing)
+
+	// singlePacketSlabSize is the size to be used for [packetBuf.slab] when
+	// neither [tun.Device] nor [conn.Bind] support batched reads.
+	singlePacketSlabSize = MaxContentSize + 2*tun.ReadPacketSpacing
+
+	// batchingSlabSize is the size to be used for [packetBuf.slab] when either
+	// [tun.Device] or [conn.Bind] support batched reads. This covers the
+	// [tun.GSOSplit] worst case for minimum-sized MTU/MSS, which is documented
+	// in [tun.TestGSOSplitMinimumGSOSize].
+	batchingSlabSize = 2 * (1<<16 - 1)
+)
+
+// The tun package cannot import device, so we assert [tun.ReadPacketSpacing] is
+// greater than or equal to [outboundPlaintextSpacing] instead.
+const _ = uint(tun.ReadPacketSpacing - outboundPlaintextSpacing)
+
+// packetBuf is a reference counted packet buffer.
+type packetBuf struct {
+	slab     []byte
+	refCount atomic.Int32
+	cleanup  func(*packetBuf)
+}
+
+// newPacketBuf returns a new [*packetBuf] of the provided size with its
+// reference count set to zero. cleanup runs for the returned [*packetBuf] when
+// its ref count moves 1->0.
+func newPacketBuf(size int, cleanup func(*packetBuf)) *packetBuf {
+	return &packetBuf{
+		slab:    make([]byte, size),
+		cleanup: cleanup,
+	}
+}
+
+func (p *packetBuf) incRef() {
+	p.refCount.Add(1)
+}
+
+func (p *packetBuf) decRef() {
+	c := p.refCount.Add(-1)
+	if c == 0 {
+		p.cleanup(p)
+	} else if c < 0 {
+		panic(fmt.Sprintf("packetBuf.decRef below zero (%d)", c))
+	}
+}

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() {

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

@@ -7,6 +7,8 @@ package device
 
 import (
 	"sync"
+
+	"github.com/tailscale/wireguard-go/tun"
 )
 
 type WaitPool struct {
@@ -46,29 +48,77 @@ func (p *WaitPool) Put(x any) {
 	p.cond.Signal()
 }
 
+const (
+	// smallPacketBufSize is the size to be used for [packetBuf.slab]'s pooled in
+	// [Device.pool.smallPacketBufs]. It's large enough to fit WireGuard keepalive
+	// messages, WireGuard-encrypted [MaxPriorityMessageContentSize], and
+	// WireGuard-encrypted internet MTU messages.
+	smallPacketBufSize = 2048
+
+	internetMTU = 1500
+	// assert that [smallPacketBufSize] is sufficient to hold [internetMTU]-sized
+	// plaintext read via [tun.Device.Read], for eventual encryption and
+	// transmission as a WireGuard transport message.
+	_ uint = smallPacketBufSize -
+		tun.ReadPacketSpacing -
+		internetMTU -
+		outboundPlaintextTailroom
+	// assert that [smallPacketBufSize] is sufficient to hold [MaxPriorityMessageContentSize]-sized
+	// plaintext for eventual encryption and transmission as a WireGuard transport
+	// message. This assertion also covers keepalive messages, which only have
+	// the auth tag, accounted for in [outboundPlaintextTailroom], following
+	// the WireGuard header.
+	_ uint = smallPacketBufSize -
+		outboundPlaintextHeadroom -
+		MaxPriorityMessageContentSize -
+		outboundPlaintextTailroom
+
+	// minPacketBufSizeForDistinctSmallPool represents the minimum [packetBuf.slab]
+	// size at which a distinct ([Device.pool.smallPacketBufs]) is created.
+	minPacketBufSizeForDistinctSmallPool = 2 * smallPacketBufSize
+)
+
 func (device *Device) PopulatePools() {
-	device.pool.inboundElementsContainer = NewWaitPool(PreallocatedBuffersPerPool, func() any {
+	device.pool.inboundElementsContainer = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
 		s := make([]*QueueInboundElement, 0, device.BatchSize())
 		return &QueueInboundElementsContainer{elems: s}
 	})
-	device.pool.outboundElementsContainer = NewWaitPool(PreallocatedBuffersPerPool, func() any {
+	device.pool.outboundElementsContainer = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
 		s := make([]*QueueOutboundElement, 0, device.BatchSize())
 		return &QueueOutboundElementsContainer{elems: s}
 	})
-	device.pool.messageBuffers = NewWaitPool(PreallocatedBuffersPerPool, func() any {
-		return new([MaxMessageSize]byte)
-	})
-	device.pool.inboundElements = NewWaitPool(PreallocatedBuffersPerPool, func() any {
+	device.pool.inboundElements = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
 		return new(QueueInboundElement)
 	})
-	device.pool.outboundElements = NewWaitPool(PreallocatedBuffersPerPool, func() any {
+	device.pool.outboundElements = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
 		return new(QueueOutboundElement)
 	})
+	packetBufSize := singlePacketSlabSize
+	if device.BatchSize() > 1 {
+		packetBufSize = batchingSlabSize
+	}
+	device.pool.packetBufs = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
+		return newPacketBuf(packetBufSize, func(buf *packetBuf) {
+			device.pool.packetBufs.Put(buf)
+		})
+	})
+	// A distinct smallPacketBufs pool makes TUN readers copy packets that fit
+	// into [smallPacketBufSize] in the interest of memory savings. Alias
+	// smallPacketBufs to the "full-sized" pool if there is little to be gained
+	// in memory savings, or if the device is configured to limit outstanding
+	// packet memory pool items, which we don't want to artificially inflate.
+	device.pool.smallPacketBufs = device.pool.packetBufs
+	if packetBufSize >= minPacketBufSizeForDistinctSmallPool && device.config.preallocatedBuffersPerPool == 0 {
+		device.pool.smallPacketBufs = NewWaitPool(device.config.preallocatedBuffersPerPool, func() any {
+			return newPacketBuf(smallPacketBufSize, func(buf *packetBuf) {
+				device.pool.smallPacketBufs.Put(buf)
+			})
+		})
+	}
 }
 
 func (device *Device) GetInboundElementsContainer() *QueueInboundElementsContainer {
 	c := device.pool.inboundElementsContainer.Get().(*QueueInboundElementsContainer)
-	c.Mutex = sync.Mutex{}
 	return c
 }
 
@@ -82,7 +132,6 @@ func (device *Device) PutInboundElementsContainer(c *QueueInboundElementsContain
 
 func (device *Device) GetOutboundElementsContainer() *QueueOutboundElementsContainer {
 	c := device.pool.outboundElementsContainer.Get().(*QueueOutboundElementsContainer)
-	c.Mutex = sync.Mutex{}
 	return c
 }
 
@@ -94,12 +143,25 @@ func (device *Device) PutOutboundElementsContainer(c *QueueOutboundElementsConta
 	device.pool.outboundElementsContainer.Put(c)
 }
 
-func (device *Device) GetMessageBuffer() *[MaxMessageSize]byte {
-	return device.pool.messageBuffers.Get().(*[MaxMessageSize]byte)
+// hasDistinctSmallPacketBufPool returns true if [Device.getSmallPacketBuf]
+// will fetch from a pool independent of [Device.getPacketBuf].
+func (device *Device) hasDistinctSmallPacketBufPool() bool {
+	return device.pool.smallPacketBufs != device.pool.packetBufs
 }
 
-func (device *Device) PutMessageBuffer(msg *[MaxMessageSize]byte) {
-	device.pool.messageBuffers.Put(msg)
+// getSmallPacketBuf returns a [*packetBuf] smaller or equal in size to what
+// [Device.getPacketBuf] returns depending on [Device] batching and [WaitPool]
+// configuration.
+func (device *Device) getSmallPacketBuf() *packetBuf {
+	b := device.pool.smallPacketBufs.Get().(*packetBuf)
+	b.incRef()
+	return b
+}
+
+func (device *Device) getPacketBuf() *packetBuf {
+	b := device.pool.packetBufs.Get().(*packetBuf)
+	b.incRef()
+	return b
 }
 
 func (device *Device) GetInboundElement() *QueueInboundElement {
@@ -107,15 +169,25 @@ func (device *Device) GetInboundElement() *QueueInboundElement {
 }
 
 func (device *Device) PutInboundElement(elem *QueueInboundElement) {
+	elem.buffer.decRef()
 	elem.clearPointers()
 	device.pool.inboundElements.Put(elem)
 }
 
+// GetOutboundElement returns a [*QueueOutboundElement] with all its fields
+// set to their respective zero values.
 func (device *Device) GetOutboundElement() *QueueOutboundElement {
-	return device.pool.outboundElements.Get().(*QueueOutboundElement)
+	elem := device.pool.outboundElements.Get().(*QueueOutboundElement)
+	elem.plaintextOffset = 0
+	elem.nonce = 0
+	// buffer, packet, keypair, and peer were cleared (if necessary) by [QueueOutboundElement.clearPointers].
+	return elem
 }
 
 func (device *Device) PutOutboundElement(elem *QueueOutboundElement) {
+	if elem.buffer != nil {
+		elem.buffer.decRef()
+	}
 	elem.clearPointers()
 	device.pool.outboundElements.Put(elem)
 }

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

@@ -10,10 +10,10 @@ import "github.com/tailscale/wireguard-go/conn"
 /* Reduce memory consumption for Android */
 
 const (
-	QueueStagedSize            = conn.IdealBatchSize
-	QueueOutboundSize          = 1024
-	QueueInboundSize           = 1024
-	QueueHandshakeSize         = 1024
-	MaxSegmentSize             = 2200
-	PreallocatedBuffersPerPool = 4096
+	DefaultQueueStagedSize            = conn.IdealBatchSize
+	DefaultQueueOutboundSize          = 1024
+	DefaultQueueInboundSize           = 1024
+	DefaultQueueHandshakeSize         = 1024
+	MaxSegmentSize                    = 2200
+	DefaultPreallocatedBuffersPerPool = 4096
 )

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

@@ -10,10 +10,10 @@ package device
 import "github.com/tailscale/wireguard-go/conn"
 
 const (
-	QueueStagedSize            = conn.IdealBatchSize
-	QueueOutboundSize          = 1024
-	QueueInboundSize           = 1024
-	QueueHandshakeSize         = 1024
-	MaxSegmentSize             = (1 << 16) - 1 // largest possible UDP datagram
-	PreallocatedBuffersPerPool = 0             // Disable and allow for infinite memory growth
+	DefaultQueueStagedSize            = conn.IdealBatchSize
+	DefaultQueueOutboundSize          = 1024
+	DefaultQueueInboundSize           = 1024
+	DefaultQueueHandshakeSize         = 1024
+	MaxSegmentSize                    = (1 << 16) - 1 // largest possible UDP datagram
+	DefaultPreallocatedBuffersPerPool = 0             // Disable and allow for infinite memory growth
 )

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

@@ -7,15 +7,15 @@
 
 package device
 
-// Fit within memory limits for iOS's Network Extension API, which has stricter requirements.
-// These are vars instead of consts, because heavier network extensions might want to reduce
-// them further.
-var (
-	QueueStagedSize                   = 128
-	QueueOutboundSize                 = 1024
-	QueueInboundSize                  = 1024
-	QueueHandshakeSize                = 1024
-	PreallocatedBuffersPerPool uint32 = 1024
+// Fit within memory limits for iOS's Network Extension API, which has stricter
+// requirements. Heavier network extensions can reduce these further using
+// [Device.Option]'s.
+const (
+	DefaultQueueStagedSize            = 128
+	DefaultQueueOutboundSize          = 1024
+	DefaultQueueInboundSize           = 1024
+	DefaultQueueHandshakeSize         = 1024
+	DefaultPreallocatedBuffersPerPool = 1024
 )
 
 const MaxSegmentSize = 1700

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

@@ -6,10 +6,10 @@
 package device
 
 const (
-	QueueStagedSize            = 128
-	QueueOutboundSize          = 1024
-	QueueInboundSize           = 1024
-	QueueHandshakeSize         = 1024
-	MaxSegmentSize             = 2048 - 32 // largest possible UDP datagram
-	PreallocatedBuffersPerPool = 0         // Disable and allow for infinite memory growth
+	DefaultQueueStagedSize            = 128
+	DefaultQueueOutboundSize          = 1024
+	DefaultQueueInboundSize           = 1024
+	DefaultQueueHandshakeSize         = 1024
+	MaxSegmentSize                    = 2048 - 32 // largest possible UDP datagram
+	DefaultPreallocatedBuffersPerPool = 0         // Disable and allow for infinite memory growth
 )

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

@@ -8,7 +8,9 @@ package device
 import (
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"net"
+	"net/netip"
 	"sync"
 	"time"
 
@@ -19,23 +21,28 @@ import (
 )
 
 type QueueHandshakeElement struct {
-	msgType  uint32
-	packet   []byte
-	endpoint conn.Endpoint
-	buffer   *[MaxMessageSize]byte
+	msgType    uint32
+	packet     []byte
+	packetMeta conn.ReceivedPacket
+	buffer     *packetBuf
 }
 
 type QueueInboundElement struct {
-	buffer   *[MaxMessageSize]byte
-	packet   []byte
-	counter  uint64
-	keypair  *Keypair
-	endpoint conn.Endpoint
+	buffer     *packetBuf
+	packet     []byte
+	packetMeta conn.ReceivedPacket
+	counter    uint64
+	keypair    *Keypair
 }
 
 type QueueInboundElementsContainer struct {
-	sync.Mutex
-	elems []*QueueInboundElement
+	// filling is a one-shot barrier signaling decryption→receive
+	// handoff. RoutineReceiveIncoming calls Add(1) before sending the
+	// container down the decryption and inbound queues; RoutineDecryption
+	// calls Done after decrypting; RoutineSequentialReceiver calls Wait
+	// before reading the decrypted packets.
+	filling sync.WaitGroup
+	elems   []*QueueInboundElement
 }
 
 // clearPointers clears elem fields that contain pointers.
@@ -45,8 +52,8 @@ type QueueInboundElementsContainer struct {
 func (elem *QueueInboundElement) clearPointers() {
 	elem.buffer = nil
 	elem.packet = nil
+	elem.packetMeta = conn.ReceivedPacket{}
 	elem.keypair = nil
-	elem.endpoint = nil
 }
 
 /* Called when a new authenticated message has been received
@@ -83,31 +90,29 @@ func (device *Device) RoutineReceiveIncoming(maxBatchSize int, recv conn.Receive
 	// receive datagrams until conn is closed
 
 	var (
-		bufsArrs    = make([]*[MaxMessageSize]byte, maxBatchSize)
-		bufs        = make([][]byte, maxBatchSize)
+		packets = make([]conn.ReceivedPacket, maxBatchSize)
+		buf     = device.getPacketBuf()
+		// bufDirty tracks whether the current buf has been pushed downstream to
+		// crypto and per-peer/handshake functions, or can be re-used across the
+		// next read cycle
+		bufDirty    bool
 		err         error
-		sizes       = make([]int, maxBatchSize)
 		count       int
-		endpoints   = make([]conn.Endpoint, maxBatchSize)
 		deathSpiral int
 		elemsByPeer = make(map[*Peer]*QueueInboundElementsContainer, maxBatchSize)
 	)
 
-	for i := range bufsArrs {
-		bufsArrs[i] = device.GetMessageBuffer()
-		bufs[i] = bufsArrs[i][:]
-	}
-
 	defer func() {
-		for i := 0; i < maxBatchSize; i++ {
-			if bufsArrs[i] != nil {
-				device.PutMessageBuffer(bufsArrs[i])
-			}
-		}
+		buf.decRef()
 	}()
 
 	for {
-		count, err = recv(bufs, sizes, endpoints)
+		if bufDirty {
+			buf.decRef()
+			buf = device.getPacketBuf()
+			bufDirty = false
+		}
+		count, err = recv(buf.slab, packets)
 		if err != nil {
 			if errors.Is(err, net.ErrClosed) {
 				return
@@ -126,14 +131,14 @@ func (device *Device) RoutineReceiveIncoming(maxBatchSize int, recv conn.Receive
 		deathSpiral = 0
 
 		// handle each packet in the batch
-		for i, size := range sizes[:count] {
-			if size < MinMessageSize {
+		for _, meta := range packets[:count] {
+			if meta.Size < MinMessageSize {
 				continue
 			}
 
 			// check size of packet
 
-			packet := bufsArrs[i][:size]
+			packet := meta.Bytes(buf.slab)
 			msgType := binary.LittleEndian.Uint32(packet[:4])
 
 			switch msgType {
@@ -168,21 +173,20 @@ func (device *Device) RoutineReceiveIncoming(maxBatchSize int, recv conn.Receive
 				// create work element
 				peer := value.peer
 				elem := device.GetInboundElement()
+				bufDirty = true
+				buf.incRef()
+				elem.buffer = buf
 				elem.packet = packet
-				elem.buffer = bufsArrs[i]
+				elem.packetMeta = meta
 				elem.keypair = keypair
-				elem.endpoint = endpoints[i]
 				elem.counter = 0
 
 				elemsForPeer, ok := elemsByPeer[peer]
 				if !ok {
 					elemsForPeer = device.GetInboundElementsContainer()
-					elemsForPeer.Lock()
 					elemsByPeer[peer] = elemsForPeer
 				}
 				elemsForPeer.elems = append(elemsForPeer.elems, elem)
-				bufsArrs[i] = device.GetMessageBuffer()
-				bufs[i] = bufsArrs[i][:]
 				continue
 
 			// otherwise it is a fixed size & handshake related packet
@@ -207,29 +211,21 @@ func (device *Device) RoutineReceiveIncoming(maxBatchSize int, recv conn.Receive
 				continue
 			}
 
+			buf.incRef()
 			select {
 			case device.queue.handshake.c <- QueueHandshakeElement{
-				msgType:  msgType,
-				buffer:   bufsArrs[i],
-				packet:   packet,
-				endpoint: endpoints[i],
+				msgType:    msgType,
+				buffer:     buf,
+				packet:     packet,
+				packetMeta: meta,
 			}:
-				bufsArrs[i] = device.GetMessageBuffer()
-				bufs[i] = bufsArrs[i][:]
+				bufDirty = true
 			default:
+				buf.decRef()
 			}
 		}
 		for peer, elemsContainer := range elemsByPeer {
-			if peer.isRunning.Load() {
-				peer.queue.inbound.c <- elemsContainer
-				device.queue.decryption.c <- elemsContainer
-			} else {
-				for _, elem := range elemsContainer.elems {
-					device.PutMessageBuffer(elem.buffer)
-					device.PutInboundElement(elem)
-				}
-				device.PutInboundElementsContainer(elemsContainer)
-			}
+			peer.queueInboundIfRunning(elemsContainer)
 			delete(elemsByPeer, peer)
 		}
 	}
@@ -262,7 +258,7 @@ func (device *Device) RoutineDecryption(id int) {
 				elem.packet = nil
 			}
 		}
-		elemsContainer.Unlock()
+		elemsContainer.filling.Done()
 	}
 }
 
@@ -302,8 +298,8 @@ func (device *Device) RoutineHandshake(id int) {
 
 			// consume reply
 
-			if peer := entry.peer; peer.isRunning.Load() {
-				device.log.Verbosef("Receiving cookie response from %s", elem.endpoint.DstToString())
+			if peer := entry.peer; peer.runningState.isRunning.Load() {
+				device.log.Verbosef("Receiving cookie response from %s", elem.packetMeta.Endpoint.DstToString())
 				if !peer.cookieGenerator.ConsumeReply(&reply) {
 					device.log.Verbosef("Could not decrypt invalid cookie response")
 				}
@@ -326,14 +322,14 @@ func (device *Device) RoutineHandshake(id int) {
 
 				// verify MAC2 field
 
-				if !device.cookieChecker.CheckMAC2(elem.packet, elem.endpoint.DstToBytes()) {
+				if !device.cookieChecker.CheckMAC2(elem.packet, elem.packetMeta.Endpoint.DstToBytes()) {
 					device.SendHandshakeCookie(&elem)
 					goto skip
 				}
 
 				// check ratelimiter
 
-				if !device.rate.limiter.Allow(elem.endpoint.DstIP()) {
+				if !device.rate.limiter.Allow(elem.packetMeta.Endpoint.DstIP()) {
 					goto skip
 				}
 			}
@@ -359,9 +355,9 @@ func (device *Device) RoutineHandshake(id int) {
 
 			// consume initiation
 
-			peer := device.ConsumeMessageInitiation(&msg, elem.endpoint)
+			peer := device.ConsumeMessageInitiation(&msg, elem.packetMeta.Endpoint)
 			if peer == nil {
-				device.log.Verbosef("Received invalid initiation message from %s", elem.endpoint.DstToString())
+				device.log.Verbosef("Received invalid initiation message from %s", elem.packetMeta.Endpoint.DstToString())
 				goto skip
 			}
 
@@ -370,9 +366,6 @@ func (device *Device) RoutineHandshake(id int) {
 			peer.timersAnyAuthenticatedPacketTraversal()
 			peer.timersAnyAuthenticatedPacketReceived()
 
-			// update endpoint
-			peer.SetEndpointFromPacket(elem.endpoint)
-
 			device.log.Verbosef("%v - Received handshake initiation", peer)
 			peer.rxBytes.Add(uint64(len(elem.packet)))
 
@@ -393,13 +386,10 @@ func (device *Device) RoutineHandshake(id int) {
 
 			peer := device.ConsumeMessageResponse(&msg)
 			if peer == nil {
-				device.log.Verbosef("Received invalid response message from %s", elem.endpoint.DstToString())
+				device.log.Verbosef("Received invalid response message from %s", elem.packetMeta.Endpoint.DstToString())
 				goto skip
 			}
 
-			// update endpoint
-			peer.SetEndpointFromPacket(elem.endpoint)
-
 			device.log.Verbosef("%v - Received handshake response", peer)
 			peer.rxBytes.Add(uint64(len(elem.packet)))
 
@@ -417,12 +407,14 @@ func (device *Device) RoutineHandshake(id int) {
 				goto skip
 			}
 
+			device.config.metrics.HandshakeInitiatorCompleted.Add(1)
 			peer.timersSessionDerived()
 			peer.timersHandshakeComplete()
+			peer.SendPriorityMessage()
 			peer.SendKeepalive()
 		}
 	skip:
-		device.PutMessageBuffer(elem.buffer)
+		elem.buffer.decRef()
 	}
 }
 
@@ -430,110 +422,136 @@ func (peer *Peer) RoutineSequentialReceiver(maxBatchSize int) {
 	device := peer.device
 	defer func() {
 		device.log.Verbosef("%v - Routine: sequential receiver - stopped", peer)
-		peer.stopping.Done()
+		peer.runningState.queueReaders.Done()
 	}()
 	device.log.Verbosef("%v - Routine: sequential receiver - started", peer)
 
 	bufs := make([][]byte, 0, maxBatchSize)
 
-	for elemsContainer := range peer.queue.inbound.c {
-		if elemsContainer == nil {
-			return
+	for elemsContainer := range peer.queue.inbound {
+		peer.processInboundContainer(elemsContainer, bufs[:0])
+	}
+}
+
+// processInboundContainer waits for the decryption routine to finish
+// filling elemsContainer, then writes the valid packets to the TUN
+// device and returns the container to the pool.
+//
+// scratch is a length-0 slice used to assemble the per-packet buffers
+// passed to tun.device.Write; its backing array is reused across calls.
+func (peer *Peer) processInboundContainer(elemsContainer *QueueInboundElementsContainer, scratch [][]byte) {
+	// Invariants from RoutineSequentialReceiver; all should be unreachable.
+	if len(scratch) != 0 || cap(scratch) == 0 {
+		panic(fmt.Sprintf("processInboundContainer: scratch must be empty with non-zero cap; got len=%d cap=%d",
+			len(scratch), cap(scratch)))
+	}
+	if cap(scratch) < len(elemsContainer.elems) {
+		panic(fmt.Sprintf("processInboundContainer: scratch cap %d < elems %d",
+			cap(scratch), len(elemsContainer.elems)))
+	}
+
+	device := peer.device
+	defer device.PutInboundElementsContainer(elemsContainer)
+
+	// Wait for RoutineDecryption to finish filling the container. After
+	// Wait returns we have happens-before with that goroutine and are the
+	// sole owner of the container until Put hands it back to the pool.
+	elemsContainer.filling.Wait()
+	elems := elemsContainer.elems
+
+	validTailPacket := -1
+	dataPacketReceived := false
+	rxBytesLen := uint64(0)
+	for i, elem := range elems {
+		if elem.packet == nil {
+			// decryption failed
+			continue
 		}
-		elemsContainer.Lock()
-		validTailPacket := -1
-		dataPacketReceived := false
-		rxBytesLen := uint64(0)
-		for i, elem := range elemsContainer.elems {
-			if elem.packet == nil {
-				// decryption failed
+
+		if !elem.keypair.replayFilter.ValidateCounter(elem.counter, RejectAfterMessages) {
+			device.config.metrics.MessageTransportRXDroppedReplay.Add(1)
+			continue
+		}
+
+		validTailPacket = i
+		if peer.ReceivedWithKeypair(elem.keypair) {
+			peer.device.config.metrics.HandshakeResponderCompleted.Add(1)
+			peer.timersHandshakeComplete()
+			peer.SendPriorityMessage()
+			peer.SendStagedPackets()
+		}
+		if ep, ok := elem.packetMeta.Endpoint.(conn.PeerAwareEndpoint); ok {
+			ep.FromPeer(peer.handshake.remoteStatic)
+		}
+		rxBytesLen += uint64(len(elem.packet) + MinMessageSize)
+
+		if len(elem.packet) == 0 {
+			device.log.Verbosef("%v - Receiving keepalive packet", peer)
+			continue
+		}
+		dataPacketReceived = true
+
+		switch elem.packet[0] >> 4 {
+		case 4:
+			if len(elem.packet) < ipv4.HeaderLen {
+				continue
+			}
+			field := elem.packet[IPv4offsetTotalLength : IPv4offsetTotalLength+2]
+			length := binary.BigEndian.Uint16(field)
+			if int(length) > len(elem.packet) || int(length) < ipv4.HeaderLen {
 				continue
 			}
-
-			if !elem.keypair.replayFilter.ValidateCounter(elem.counter, RejectAfterMessages) {
+			elem.packet = elem.packet[:length]
+			src := elem.packet[IPv4offsetSrc : IPv4offsetSrc+net.IPv4len]
+			srcAddr, _ := netip.AddrFromSlice(src)
+			if !peer.AllowedPeerSourceIP(srcAddr) {
+				device.log.Verbosef("IPv4 packet with disallowed source address from %v", peer)
 				continue
 			}
 
-			validTailPacket = i
-			if peer.ReceivedWithKeypair(elem.keypair) {
-				peer.SetEndpointFromPacket(elem.endpoint)
-				peer.timersHandshakeComplete()
-				peer.SendStagedPackets()
-			}
-			if ep, ok := elem.endpoint.(conn.PeerAwareEndpoint); ok {
-				ep.FromPeer(peer.handshake.remoteStatic)
+		case 6:
+			if len(elem.packet) < ipv6.HeaderLen {
+				continue
 			}
-			rxBytesLen += uint64(len(elem.packet) + MinMessageSize)
-
-			if len(elem.packet) == 0 {
-				device.log.Verbosef("%v - Receiving keepalive packet", peer)
+			field := elem.packet[IPv6offsetPayloadLength : IPv6offsetPayloadLength+2]
+			length := binary.BigEndian.Uint16(field)
+			length += ipv6.HeaderLen
+			if int(length) > len(elem.packet) {
 				continue
 			}
-			dataPacketReceived = true
-
-			switch elem.packet[0] >> 4 {
-			case 4:
-				if len(elem.packet) < ipv4.HeaderLen {
-					continue
-				}
-				field := elem.packet[IPv4offsetTotalLength : IPv4offsetTotalLength+2]
-				length := binary.BigEndian.Uint16(field)
-				if int(length) > len(elem.packet) || int(length) < ipv4.HeaderLen {
-					continue
-				}
-				elem.packet = elem.packet[:length]
-				src := elem.packet[IPv4offsetSrc : IPv4offsetSrc+net.IPv4len]
-				if device.allowedips.Lookup(src) != peer {
-					device.log.Verbosef("IPv4 packet with disallowed source address from %v", peer)
-					continue
-				}
-
-			case 6:
-				if len(elem.packet) < ipv6.HeaderLen {
-					continue
-				}
-				field := elem.packet[IPv6offsetPayloadLength : IPv6offsetPayloadLength+2]
-				length := binary.BigEndian.Uint16(field)
-				length += ipv6.HeaderLen
-				if int(length) > len(elem.packet) {
-					continue
-				}
-				elem.packet = elem.packet[:length]
-				src := elem.packet[IPv6offsetSrc : IPv6offsetSrc+net.IPv6len]
-				if device.allowedips.Lookup(src) != peer {
-					device.log.Verbosef("IPv6 packet with disallowed source address from %v", peer)
-					continue
-				}
-
-			default:
-				device.log.Verbosef("Packet with invalid IP version from %v", peer)
+			elem.packet = elem.packet[:length]
+			src := elem.packet[IPv6offsetSrc : IPv6offsetSrc+net.IPv6len]
+			srcAddr, _ := netip.AddrFromSlice(src)
+			if !peer.AllowedPeerSourceIP(srcAddr) {
+				device.log.Verbosef("IPv6 packet with disallowed source address from %v", peer)
 				continue
 			}
 
-			bufs = append(bufs, elem.buffer[:MessageTransportOffsetContent+len(elem.packet)])
+		default:
+			device.log.Verbosef("Packet with invalid IP version from %v", peer)
+			continue
 		}
+
+		elem.packetMeta.Size = MessageTransportOffsetContent + len(elem.packet)
+		scratch = append(scratch, elem.packetMeta.Bytes(elem.buffer.slab))
+	}
 
-		peer.rxBytes.Add(rxBytesLen)
-		if validTailPacket >= 0 {
-			peer.SetEndpointFromPacket(elemsContainer.elems[validTailPacket].endpoint)
-			peer.keepKeyFreshReceiving()
-			peer.timersAnyAuthenticatedPacketTraversal()
-			peer.timersAnyAuthenticatedPacketReceived()
-		}
-		if dataPacketReceived {
-			peer.timersDataReceived()
-		}
-		if len(bufs) > 0 {
-			_, err := device.tun.device.Write(bufs, MessageTransportOffsetContent)
-			if err != nil && !device.isClosed() {
-				device.log.Errorf("Failed to write packets to TUN device: %v", err)
-			}
-		}
-		for _, elem := range elemsContainer.elems {
-			device.PutMessageBuffer(elem.buffer)
-			device.PutInboundElement(elem)
+	peer.rxBytes.Add(rxBytesLen)
+	if validTailPacket >= 0 {
+		peer.keepKeyFreshReceiving()
+		peer.timersAnyAuthenticatedPacketTraversal()
+		peer.timersAnyAuthenticatedPacketReceived()
+	}
+	if dataPacketReceived {
+		peer.timersDataReceived()
+	}
+	if len(scratch) > 0 {
+		_, err := peer.writeTUN(scratch, MessageTransportOffsetContent)
+		if err != nil && !device.isClosed() {
+			device.log.Errorf("Failed to write packets to TUN device: %v", err)
 		}
-		bufs = bufs[:0]
-		device.PutInboundElementsContainer(elemsContainer)
+	}
+	for _, elem := range elems {
+		device.PutInboundElement(elem)
 	}
 }

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

@@ -8,7 +8,9 @@ package device
 import (
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"net"
+	"net/netip"
 	"os"
 	"sync"
 	"time"
@@ -45,28 +47,26 @@ import (
  */
 
 type QueueOutboundElement struct {
-	buffer *[MaxMessageSize]byte // slice holding the packet data
-	// packet is always a slice of "buffer". The starting offset in buffer
-	// is either:
-	//  a) MessageEncapsulatingTransportSize+MessageTransportHeaderSize (plaintext)
-	//  b) 0 (post-encryption)
-	packet  []byte
-	nonce   uint64   // nonce for encryption
-	keypair *Keypair // keypair for encryption
-	peer    *Peer    // related peer
+	buffer *packetBuf // shared packet buffer that contains packet
+	packet []byte     // packet is always a slice of buffer.slab
+	// plaintextOffset describes the pre-encryption offset of packet within
+	// buffer.slab. The encryption process grows packet on both ends to make
+	// room for any potential encapsulating transport, WireGuard header, and
+	// tail padding.
+	plaintextOffset int
+	nonce           uint64   // nonce for encryption
+	keypair         *Keypair // keypair for encryption
+	peer            *Peer    // related peer
 }
 
 type QueueOutboundElementsContainer struct {
-	sync.Mutex
-	elems []*QueueOutboundElement
-}
-
-func (device *Device) NewOutboundElement() *QueueOutboundElement {
-	elem := device.GetOutboundElement()
-	elem.buffer = device.GetMessageBuffer()
-	elem.nonce = 0
-	// keypair and peer were cleared (if necessary) by clearPointers.
-	return elem
+	// filling is a one-shot barrier signaling encryption→send handoff.
+	// SendStagedPackets calls Add(1) before sending the container down
+	// the encryption and outbound queues; RoutineEncryption calls Done
+	// after encrypting; RoutineSequentialSender calls Wait before
+	// reading the encrypted packets.
+	filling sync.WaitGroup
+	elems   []*QueueOutboundElement
 }
 
 // clearPointers clears elem fields that contain pointers.
@@ -80,18 +80,24 @@ func (elem *QueueOutboundElement) clearPointers() {
 	elem.peer = nil
 }
 
-/* Queues a keepalive if no packets are queued for peer
- */
+// SendKeepalive queues a keepalive if no packets are queued for
+// peer.
 func (peer *Peer) SendKeepalive() {
-	if len(peer.queue.staged) == 0 && peer.isRunning.Load() {
-		elem := peer.device.NewOutboundElement()
+	if len(peer.queue.staged) == 0 {
+		elem := peer.device.GetOutboundElement()
+		elem.buffer = peer.device.getSmallPacketBuf()
+		elem.plaintextOffset = MessageEncapsulatingTransportSize + MessageTransportHeaderSize
 		elemsContainer := peer.device.GetOutboundElementsContainer()
 		elemsContainer.elems = append(elemsContainer.elems, elem)
-		select {
-		case peer.queue.staged <- elemsContainer:
-			peer.device.log.Verbosef("%v - Sending keepalive packet", peer)
-		default:
-			peer.device.PutMessageBuffer(elem.buffer)
+		if queued := peer.doIfRunning(func() {
+			select {
+			case peer.queue.staged <- elemsContainer:
+				peer.device.log.Verbosef("%v - Sending keepalive packet", peer)
+			default:
+				peer.device.PutOutboundElement(elem)
+				peer.device.PutOutboundElementsContainer(elemsContainer)
+			}
+		}); !queued {
 			peer.device.PutOutboundElement(elem)
 			peer.device.PutOutboundElementsContainer(elemsContainer)
 		}
@@ -99,6 +105,62 @@ func (peer *Peer) SendKeepalive() {
 	peer.SendStagedPackets()
 }
 
+// SendPriorityMessage invokes the [PeerPriorityMessageFunc] callback if one is
+// set, and queues the returned message for encryption and transmission if the
+// current keypair is valid.
+func (peer *Peer) SendPriorityMessage() {
+	f := peer.device.priorityMsgFn.Load()
+	if f == nil {
+		return
+	}
+	keypair := peer.keypairs.Current()
+	if keypair == nil || keypair.sendNonce.Load() >= RejectAfterMessages || time.Since(keypair.created) >= RejectAfterTime {
+		// SendStagedPackets initializes a handshake when the keypair is invalid,
+		// but we explicitly avoid that here. A priority message is only intended
+		// to flow around symmetric session establishment, but it should never
+		// trigger a new session. Reaching this branch due to nonce exhaustion
+		// or keypair expiration is highly unlikely considering where
+		// SendPriorityMessage is called (at current keypair establishment).
+		return
+	}
+
+	// get plaintext message to send
+	msg := (*f)(peer.handshake.remoteStatic)
+	if len(msg) == 0 {
+		return
+	}
+	if len(msg) > MaxPriorityMessageContentSize {
+		peer.device.log.Verbosef("%v - Failed to queue priority message due to size", peer)
+		return
+	}
+
+	// get pooled elements
+	elem := peer.device.GetOutboundElement()
+	elemsContainer := peer.device.GetOutboundElementsContainer()
+	elemsContainer.elems = append(elemsContainer.elems, elem)
+	buf := peer.device.getSmallPacketBuf()
+
+	// initialize outbound element
+	const offset = MessageEncapsulatingTransportSize + MessageTransportHeaderSize
+
+	n := copy(buf.slab[offset:], msg)
+	elem.buffer = buf
+	elem.packet = buf.slab[offset : offset+n]
+	elem.plaintextOffset = offset
+	elem.peer = peer
+	elem.nonce = keypair.sendNonce.Add(1) - 1
+	if elem.nonce >= RejectAfterMessages {
+		keypair.sendNonce.Store(RejectAfterMessages)
+		peer.device.PutOutboundElement(elem)
+		peer.device.PutOutboundElementsContainer(elemsContainer)
+		return
+	}
+	elem.keypair = keypair
+
+	// add to parallel and sequential queue
+	peer.queueOutboundIfRunning(elemsContainer)
+}
+
 func (peer *Peer) SendHandshakeInitiation(isRetry bool) error {
 	if !isRetry {
 		peer.timers.handshakeAttempts.Store(0)
@@ -116,6 +178,10 @@ func (peer *Peer) SendHandshakeInitiation(isRetry bool) error {
 		peer.handshake.mutex.Unlock()
 		return nil
 	}
+	// handshakeOnUserSend is flipped false after lastSentHandshake checks,
+	// enabling eventual transmission at a future call of this method, while
+	// still honoring [RekeyTimeout].
+	peer.handshakeOnUserSend.Store(false)
 	peer.handshake.lastSentHandshake = time.Now()
 	peer.handshake.mutex.Unlock()
 
@@ -135,6 +201,11 @@ func (peer *Peer) SendHandshakeInitiation(isRetry bool) error {
 	peer.timersAnyAuthenticatedPacketTraversal()
 	peer.timersAnyAuthenticatedPacketSent()
 
+	if isRetry {
+		peer.device.config.metrics.MessageInitiationTXAttemptRetry.Add(1)
+	} else {
+		peer.device.config.metrics.MessageInitiationTXAttemptInitial.Add(1)
+	}
 	err = peer.SendBuffers([][]byte{buf})
 	if err != nil {
 		peer.device.log.Errorf("%v - Failed to send handshake initiation: %v", peer, err)
@@ -172,6 +243,7 @@ func (peer *Peer) SendHandshakeResponse() error {
 	peer.timersAnyAuthenticatedPacketTraversal()
 	peer.timersAnyAuthenticatedPacketSent()
 
+	peer.device.config.metrics.MessageResponseTXAttempt.Add(1)
 	// TODO: allocation could be avoided
 	err = peer.SendBuffers([][]byte{buf})
 	if err != nil {
@@ -181,10 +253,10 @@ func (peer *Peer) SendHandshakeResponse() error {
 }
 
 func (device *Device) SendHandshakeCookie(initiatingElem *QueueHandshakeElement) error {
-	device.log.Verbosef("Sending cookie response for denied handshake message for %v", initiatingElem.endpoint.DstToString())
+	device.log.Verbosef("Sending cookie response for denied handshake message for %v", initiatingElem.packetMeta.Endpoint.DstToString())
 
 	sender := binary.LittleEndian.Uint32(initiatingElem.packet[4:8])
-	reply, err := device.cookieChecker.CreateReply(initiatingElem.packet, sender, initiatingElem.endpoint.DstToBytes())
+	reply, err := device.cookieChecker.CreateReply(initiatingElem.packet, sender, initiatingElem.packetMeta.Endpoint.DstToBytes())
 	if err != nil {
 		device.log.Errorf("Failed to create cookie reply: %v", err)
 		return err
@@ -193,8 +265,9 @@ func (device *Device) SendHandshakeCookie(initiatingElem *QueueHandshakeElement)
 	buf := make([]byte, MessageEncapsulatingTransportSize+MessageCookieReplySize)
 	packet := buf[MessageEncapsulatingTransportSize:]
 	_ = reply.marshal(packet)
+	device.config.metrics.MessageCookieReplyTXAttempt.Add(1)
 	// TODO: allocation could be avoided
-	device.net.bind.Send([][]byte{buf}, initiatingElem.endpoint, MessageEncapsulatingTransportSize)
+	device.net.bind.Send([][]byte{buf}, initiatingElem.packetMeta.Endpoint, MessageEncapsulatingTransportSize)
 
 	return nil
 }
@@ -204,73 +277,132 @@ func (peer *Peer) keepKeyFreshSending() {
 	if keypair == nil {
 		return
 	}
+	txHandshake := peer.handshakeOnUserSend.Load()
 	nonce := keypair.sendNonce.Load()
-	if nonce > RekeyAfterMessages || (keypair.isInitiator && time.Since(keypair.created) > RekeyAfterTime) {
+	if txHandshake || nonce > RekeyAfterMessages || (keypair.isInitiator && time.Since(keypair.created) > RekeyAfterTime) {
 		peer.SendHandshakeInitiation(false)
 	}
 }
 
-func (device *Device) RoutineReadFromTUN() {
+// copyPacketBufIfFits copies the prefix of src.slab ending at the final packet
+// into dst.slab if that prefix+[outboundPlaintextTailroom] fits. It reports
+// whether the copy occurred. Packet descriptors must be in offset ascending order.
+func copyPacketBufIfFits(dst, src *packetBuf, packets []tun.ReadPacket) bool {
+	if dst == nil || len(packets) < 1 {
+		return false
+	}
+	// validate tail descriptor as it relates to src
+	tailPkt := packets[len(packets)-1]
+	if tailPkt.Offset < 0 || tailPkt.Size < 1 {
+		return false
+	}
+	end := tailPkt.Offset + tailPkt.Size
+	if end < tailPkt.Offset || end > len(src.slab) ||
+		outboundPlaintextTailroom > len(src.slab)-end {
+		return false
+	}
+	// validate src fits in dst
+	if end+outboundPlaintextTailroom > len(dst.slab) {
+		return false
+	}
+	copy(dst.slab, src.slab[:end])
+	return true
+}
+
+func (device *Device) RoutineReadFromTUN(id int, queue tun.Reader) {
 	defer func() {
-		device.log.Verbosef("Routine: TUN reader - stopped")
+		device.log.Verbosef("Routine: TUN reader %d - stopped", id)
 		device.state.stopping.Done()
 		device.queue.encryption.wg.Done()
 	}()
 
-	device.log.Verbosef("Routine: TUN reader - started")
+	device.log.Verbosef("Routine: TUN reader %d - started", id)
 
 	var (
 		batchSize   = device.BatchSize()
+		packets     = make([]tun.ReadPacket, batchSize)
+		fullSizeBuf = device.getPacketBuf()
+		smallBuf    *packetBuf
 		readErr     error
 		elems       = make([]*QueueOutboundElement, batchSize)
-		bufs        = make([][]byte, batchSize)
 		elemsByPeer = make(map[*Peer]*QueueOutboundElementsContainer, batchSize)
 		count       = 0
-		sizes       = make([]int, batchSize)
-		offset      = MessageEncapsulatingTransportSize + MessageTransportHeaderSize
 	)
 
+	if device.hasDistinctSmallPacketBufPool() {
+		smallBuf = device.getSmallPacketBuf()
+	}
+
+	defer func() {
+		fullSizeBuf.decRef()
+		if smallBuf != nil {
+			smallBuf.decRef()
+		}
+	}()
+
 	for i := range elems {
-		elems[i] = device.NewOutboundElement()
-		bufs[i] = elems[i].buffer[:]
+		elems[i] = device.GetOutboundElement()
 	}
 
 	defer func() {
 		for _, elem := range elems {
 			if elem != nil {
-				device.PutMessageBuffer(elem.buffer)
 				device.PutOutboundElement(elem)
 			}
 		}
 	}()
 
+	var (
+		buf          *packetBuf // either fullSizeBuf or smallBuf
+		bufHandedOff bool
+	)
+
 	for {
+		if bufHandedOff {
+			if buf == smallBuf {
+				smallBuf.decRef()
+				smallBuf = device.getSmallPacketBuf()
+			} else {
+				fullSizeBuf.decRef()
+				fullSizeBuf = device.getPacketBuf()
+			}
+		}
+		buf = fullSizeBuf
+		bufHandedOff = false
+
 		// read packets
-		count, readErr = device.tun.device.Read(bufs, sizes, offset)
-		for i := 0; i < count; i++ {
-			if sizes[i] < 1 {
+		count, readErr = queue.Read(fullSizeBuf.slab, packets)
+		if copyPacketBufIfFits(smallBuf, fullSizeBuf, packets[:count]) {
+			buf = smallBuf
+		}
+
+		for i, meta := range packets[:count] {
+			if meta.Size < 1 || meta.Size > MaxContentSize {
 				continue
 			}
 
-			elem := elems[i]
-			elem.packet = bufs[i][offset : offset+sizes[i]]
+			// explicitly pin slice capacity, so downstream expansion of length
+			// never leak into an adjacent packet
+			packet := buf.slab[meta.Offset : meta.Offset+meta.Size : meta.Offset+meta.Size+outboundPlaintextTailroom]
 
 			// lookup peer
 			var peer *Peer
-			switch elem.packet[0] >> 4 {
+			switch packet[0] >> 4 {
 			case 4:
-				if len(elem.packet) < ipv4.HeaderLen {
+				if len(packet) < ipv4.HeaderLen {
 					continue
 				}
-				dst := elem.packet[IPv4offsetDst : IPv4offsetDst+net.IPv4len]
-				peer = device.allowedips.Lookup(dst)
+				src := netip.AddrFrom4([4]byte(packet[IPv4offsetSrc : IPv4offsetSrc+net.IPv4len]))
+				dst := netip.AddrFrom4([4]byte(packet[IPv4offsetDst : IPv4offsetDst+net.IPv4len]))
+				peer = device.allowedips.LookupFromPacket(src, dst, packet)
 
 			case 6:
-				if len(elem.packet) < ipv6.HeaderLen {
+				if len(packet) < ipv6.HeaderLen {
 					continue
 				}
-				dst := elem.packet[IPv6offsetDst : IPv6offsetDst+net.IPv6len]
-				peer = device.allowedips.Lookup(dst)
+				src := netip.AddrFrom16([16]byte(packet[IPv6offsetSrc : IPv6offsetSrc+net.IPv6len]))
+				dst := netip.AddrFrom16([16]byte(packet[IPv6offsetDst : IPv6offsetDst+net.IPv6len]))
+				peer = device.allowedips.LookupFromPacket(src, dst, packet)
 
 			default:
 				device.log.Verbosef("Received packet with unknown IP version")
@@ -279,27 +411,26 @@ func (device *Device) RoutineReadFromTUN() {
 			if peer == nil {
 				continue
 			}
+
+			elem := elems[i]
+			buf.incRef()
+			bufHandedOff = true
+			elem.buffer = buf
+			elem.packet = packet
+			elem.plaintextOffset = meta.Offset
+
 			elemsForPeer, ok := elemsByPeer[peer]
 			if !ok {
 				elemsForPeer = device.GetOutboundElementsContainer()
 				elemsByPeer[peer] = elemsForPeer
 			}
 			elemsForPeer.elems = append(elemsForPeer.elems, elem)
-			elems[i] = device.NewOutboundElement()
-			bufs[i] = elems[i].buffer[:]
+			elems[i] = device.GetOutboundElement()
 		}
 
 		for peer, elemsForPeer := range elemsByPeer {
-			if peer.isRunning.Load() {
-				peer.StagePackets(elemsForPeer)
-				peer.SendStagedPackets()
-			} else {
-				for _, elem := range elemsForPeer.elems {
-					device.PutMessageBuffer(elem.buffer)
-					device.PutOutboundElement(elem)
-				}
-				device.PutOutboundElementsContainer(elemsForPeer)
-			}
+			peer.StagePackets(elemsForPeer)
+			peer.SendStagedPackets()
 			delete(elemsByPeer, peer)
 		}
 
@@ -323,24 +454,31 @@ func (device *Device) RoutineReadFromTUN() {
 }
 
 func (peer *Peer) StagePackets(elems *QueueOutboundElementsContainer) {
-	for {
-		select {
-		case peer.queue.staged <- elems:
-			return
-		default:
-		}
-		select {
-		case tooOld := <-peer.queue.staged:
-			for _, elem := range tooOld.elems {
-				peer.device.PutMessageBuffer(elem.buffer)
-				peer.device.PutOutboundElement(elem)
+	if running := peer.doIfRunning(func() {
+		for {
+			select {
+			case peer.queue.staged <- elems:
+				return
+			default:
+			}
+			select {
+			case tooOld := <-peer.queue.staged:
+				for _, elem := range tooOld.elems {
+					peer.device.PutOutboundElement(elem)
+				}
+				peer.device.PutOutboundElementsContainer(tooOld)
+			default:
 			}
-			peer.device.PutOutboundElementsContainer(tooOld)
-		default:
 		}
+	}); !running {
+		for _, elem := range elems.elems {
+			peer.device.PutOutboundElement(elem)
+		}
+		peer.device.PutOutboundElementsContainer(elems)
 	}
 }
 
+// SendStagedPackets sends any staged packets to Peer.
 func (peer *Peer) SendStagedPackets() {
 top:
 	if len(peer.queue.staged) == 0 || !peer.device.isUp() {
@@ -375,7 +513,6 @@ top:
 
 				elem.keypair = keypair
 			}
-			elemsContainer.Lock()
 			elemsContainer.elems = elemsContainer.elems[:i]
 
 			if elemsContainerOOO != nil {
@@ -387,17 +524,7 @@ top:
 				goto top
 			}
 
-			// add to parallel and sequential queue
-			if peer.isRunning.Load() {
-				peer.queue.outbound.c <- elemsContainer
-				peer.device.queue.encryption.c <- elemsContainer
-			} else {
-				for _, elem := range elemsContainer.elems {
-					peer.device.PutMessageBuffer(elem.buffer)
-					peer.device.PutOutboundElement(elem)
-				}
-				peer.device.PutOutboundElementsContainer(elemsContainer)
-			}
+			peer.queueOutboundIfRunning(elemsContainer)
 
 			if elemsContainerOOO != nil {
 				goto top
@@ -413,7 +540,6 @@ func (peer *Peer) FlushStagedPackets() {
 		select {
 		case elemsContainer := <-peer.queue.staged:
 			for _, elem := range elemsContainer.elems {
-				peer.device.PutMessageBuffer(elem.buffer)
 				peer.device.PutOutboundElement(elem)
 			}
 			peer.device.PutOutboundElementsContainer(elemsContainer)
@@ -453,7 +579,7 @@ func (device *Device) RoutineEncryption(id int) {
 	for elemsContainer := range device.queue.encryption.c {
 		for _, elem := range elemsContainer.elems {
 			// populate header fields
-			header := elem.buffer[MessageEncapsulatingTransportSize : MessageEncapsulatingTransportSize+MessageTransportHeaderSize]
+			header := elem.buffer.slab[elem.plaintextOffset-MessageTransportHeaderSize : elem.plaintextOffset]
 
 			fieldType := header[0:4]
 			fieldReceiver := header[4:8]
@@ -465,6 +591,7 @@ func (device *Device) RoutineEncryption(id int) {
 
 			// pad content to multiple of 16
 			paddingSize := calculatePaddingSize(len(elem.packet), int(device.tun.mtu.Load()))
+
 			elem.packet = append(elem.packet, paddingZeros[:paddingSize]...)
 
 			// encrypt content and release to consumer
@@ -478,9 +605,11 @@ func (device *Device) RoutineEncryption(id int) {
 			)
 
 			// re-slice packet to include encapsulating transport space
-			elem.packet = elem.buffer[:MessageEncapsulatingTransportSize+len(elem.packet)]
+			start := elem.plaintextOffset - MessageTransportHeaderSize - MessageEncapsulatingTransportSize
+			end := elem.plaintextOffset - MessageTransportHeaderSize + len(elem.packet)
+			elem.packet = elem.buffer.slab[start:end]
 		}
-		elemsContainer.Unlock()
+		elemsContainer.filling.Done()
 	}
 }
 
@@ -488,65 +617,84 @@ func (peer *Peer) RoutineSequentialSender(maxBatchSize int) {
 	device := peer.device
 	defer func() {
 		defer device.log.Verbosef("%v - Routine: sequential sender - stopped", peer)
-		peer.stopping.Done()
+		peer.runningState.queueReaders.Done()
 	}()
 	device.log.Verbosef("%v - Routine: sequential sender - started", peer)
 
 	bufs := make([][]byte, 0, maxBatchSize)
 
-	for elemsContainer := range peer.queue.outbound.c {
-		bufs = bufs[:0]
-		if elemsContainer == nil {
-			return
-		}
-		if !peer.isRunning.Load() {
-			// peer has been stopped; return re-usable elems to the shared pool.
-			// This is an optimization only. It is possible for the peer to be stopped
-			// immediately after this check, in which case, elem will get processed.
-			// The timers and SendBuffers code are resilient to a few stragglers.
-			// TODO: rework peer shutdown order to ensure
-			// that we never accidentally keep timers alive longer than necessary.
-			elemsContainer.Lock()
-			for _, elem := range elemsContainer.elems {
-				device.PutMessageBuffer(elem.buffer)
-				device.PutOutboundElement(elem)
-			}
-			device.PutOutboundElementsContainer(elemsContainer)
-			continue
-		}
-		dataSent := false
-		elemsContainer.Lock()
-		for _, elem := range elemsContainer.elems {
-			if len(elem.packet[MessageEncapsulatingTransportSize:]) != MessageKeepaliveSize {
-				dataSent = true
-			}
-			bufs = append(bufs, elem.packet)
-		}
-
-		peer.timersAnyAuthenticatedPacketTraversal()
-		peer.timersAnyAuthenticatedPacketSent()
-
-		err := peer.SendBuffers(bufs)
-		if dataSent {
-			peer.timersDataSent()
-		}
-		for _, elem := range elemsContainer.elems {
-			device.PutMessageBuffer(elem.buffer)
-			device.PutOutboundElement(elem)
-		}
-		device.PutOutboundElementsContainer(elemsContainer)
-		if err != nil {
-			var errGSO conn.ErrUDPGSODisabled
-			if errors.As(err, &errGSO) {
-				device.log.Verbosef(err.Error())
-				err = errGSO.RetryErr
-			}
-		}
-		if err != nil {
-			device.log.Errorf("%v - Failed to send data packets: %v", peer, err)
-			continue
-		}
-
-		peer.keepKeyFreshSending()
+	for elemsContainer := range peer.queue.outbound {
+		peer.processOutboundContainer(elemsContainer, bufs[:0])
 	}
 }
+
+// processOutboundContainer waits for the encryption routine to finish
+// filling elemsContainer, then sends the batch (or drops it, if the peer
+// has been stopped) and returns the container to the pool.
+//
+// scratch is a length-0 slice used to assemble the per-packet buffers
+// passed to SendBuffers; its backing array is reused across calls.
+func (peer *Peer) processOutboundContainer(elemsContainer *QueueOutboundElementsContainer, scratch [][]byte) {
+	// Invariants from RoutineSequentialSender; all should be unreachable.
+	if len(scratch) != 0 || cap(scratch) == 0 {
+		panic(fmt.Sprintf("processOutboundContainer: scratch must be empty with non-zero cap; got len=%d cap=%d",
+			len(scratch), cap(scratch)))
+	}
+	if cap(scratch) < len(elemsContainer.elems) {
+		panic(fmt.Sprintf("processOutboundContainer: scratch cap %d < elems %d",
+			cap(scratch), len(elemsContainer.elems)))
+	}
+
+	device := peer.device
+	defer device.PutOutboundElementsContainer(elemsContainer)
+
+	// Wait for RoutineEncryption to finish filling the container. After
+	// Wait returns we have happens-before with that goroutine and are the
+	// sole owner of the container until Put hands it back to the pool.
+	elemsContainer.filling.Wait()
+
+	if !peer.runningState.isRunning.Load() {
+		// peer has been stopped; return re-usable elems to the shared pool.
+		// This is an optimization only. It is possible for the peer to be stopped
+		// immediately after this check, in which case, elem will get processed.
+		// The timers and SendBuffers code are resilient to a few stragglers.
+		// TODO: rework peer shutdown order to ensure
+		// that we never accidentally keep timers alive longer than necessary.
+		for _, elem := range elemsContainer.elems {
+			device.PutOutboundElement(elem)
+		}
+		return
+	}
+
+	dataSent := false
+	for _, elem := range elemsContainer.elems {
+		if len(elem.packet[MessageEncapsulatingTransportSize:]) != MessageKeepaliveSize {
+			dataSent = true
+		}
+		scratch = append(scratch, elem.packet)
+	}
+
+	peer.timersAnyAuthenticatedPacketTraversal()
+	peer.timersAnyAuthenticatedPacketSent()
+
+	err := peer.SendBuffers(scratch)
+	if dataSent {
+		peer.timersDataSent()
+	}
+	for _, elem := range elemsContainer.elems {
+		device.PutOutboundElement(elem)
+	}
+	if err != nil {
+		var errGSO conn.ErrUDPGSODisabled
+		if errors.As(err, &errGSO) {
+			device.log.Verbosef(err.Error())
+			err = errGSO.RetryErr
+		}
+	}
+	if err != nil {
+		device.log.Errorf("%v - Failed to send data packets: %v", peer, err)
+		return
+	}
+
+	peer.keepKeyFreshSending()
+}

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

@@ -73,7 +73,7 @@ func (timer *Timer) IsPending() bool {
 }
 
 func (peer *Peer) timersActive() bool {
-	return peer.isRunning.Load() && peer.device != nil && peer.device.isUp()
+	return peer.runningState.isRunning.Load() && peer.device != nil && peer.device.isUp()
 }
 
 func expiredRetransmitHandshake(peer *Peer) {
@@ -95,6 +95,7 @@ func expiredRetransmitHandshake(peer *Peer) {
 		if peer.timersActive() && !peer.timers.zeroKeyMaterial.IsPending() {
 			peer.timers.zeroKeyMaterial.Mod(RejectAfterTime * 3)
 		}
+		peer.noteSessionHandshakeStopped()
 	} else {
 		peer.timers.handshakeAttempts.Add(1)
 		peer.device.log.Verbosef("%s - Handshake did not complete after %d seconds, retrying (try %d)", peer, int(RekeyTimeout.Seconds()), peer.timers.handshakeAttempts.Load()+1)
@@ -126,6 +127,24 @@ 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 expiredSession(peer *Peer) {
+	peer.sessionState.Lock()
+	defer peer.sessionState.Unlock()
+	if peer.sessionState.sessionExpires.IsZero() || time.Now().Before(peer.sessionState.sessionExpires) {
+		return
+	}
+	peer.device.log.Verbosef("%s - Session expired after %d seconds", peer, int(RejectAfterTime.Seconds()))
+	peer.noteSessionStateLocked(PeerSessionExpired)
 }
 
 func expiredPersistentKeepalive(peer *Peer) {
@@ -171,6 +190,7 @@ func (peer *Peer) timersHandshakeInitiated() {
 	if peer.timersActive() {
 		peer.timers.retransmitHandshake.Mod(RekeyTimeout + time.Millisecond*time.Duration(fastrandn(RekeyTimeoutJitterMaxMs)))
 	}
+	peer.noteSessionHandshakeStarted()
 }
 
 /* Should be called after a handshake response message is received and processed or when getting key confirmation via the first data message. */
@@ -186,7 +206,14 @@ func (peer *Peer) timersHandshakeComplete() {
 /* Should be called after an ephemeral key is created, which is before sending a handshake response or after receiving a handshake response. */
 func (peer *Peer) timersSessionDerived() {
 	if peer.timersActive() {
+		peer.sessionState.Lock()
+		peer.sessionState.sessionExpires = time.Now().Add(RejectAfterTime)
+		peer.noteSessionStateLocked(PeerSessionEstablished)
+		peer.sessionState.Unlock()
+		peer.timers.sessionExpired.Mod(RejectAfterTime)
 		peer.timers.zeroKeyMaterial.Mod(RejectAfterTime * 3)
+	} else {
+		peer.noteSessionState(PeerSessionEstablished)
 	}
 }
 
@@ -202,6 +229,7 @@ func (peer *Peer) timersInit() {
 	peer.timers.retransmitHandshake = peer.NewTimer(expiredRetransmitHandshake)
 	peer.timers.sendKeepalive = peer.NewTimer(expiredSendKeepalive)
 	peer.timers.newHandshake = peer.NewTimer(expiredNewHandshake)
+	peer.timers.sessionExpired = peer.NewTimer(expiredSession)
 	peer.timers.zeroKeyMaterial = peer.NewTimer(expiredZeroKeyMaterial)
 	peer.timers.persistentKeepalive = peer.NewTimer(expiredPersistentKeepalive)
 }
@@ -216,6 +244,7 @@ func (peer *Peer) timersStop() {
 	peer.timers.retransmitHandshake.DelSync()
 	peer.timers.sendKeepalive.DelSync()
 	peer.timers.newHandshake.DelSync()
+	peer.timers.sessionExpired.DelSync()
 	peer.timers.zeroKeyMaterial.DelSync()
 	peer.timers.persistentKeepalive.DelSync()
 }

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

@@ -259,9 +259,6 @@ func (peer *ipcSetPeer) handlePostConfig() {
 	if peer.Peer == nil || peer.dummy {
 		return
 	}
-	if peer.created {
-		peer.endpoint.disableRoaming = peer.device.net.brokenRoaming && peer.endpoint.val != nil
-	}
 	if peer.device.isUp() {
 		peer.Start()
 		if peer.pkaOn {