device: refactor container locking for lock-order clarity

Device-side portion of upstream tailscale/wireguard-go e3ac4a0
(device, cmd/check-lockorder: add static analysis tool for lock
ordering); the analyzer itself is not carried in this fork.

device/allowedips.go

@@ -329,7 +329,6 @@ func (table *AllowedIPs) Remove(prefix netip.Prefix, peer *Peer) {
 	node.remove()
 }
 
-
 // setPeerPrefixes atomically removes all of peer's existing prefixes and adds
 // the provided ones.
 func (table *AllowedIPs) setPeerPrefixes(peer *Peer, prefixes []netip.Prefix) {

device/channels.go

@@ -97,7 +97,7 @@ func (device *Device) flushInboundQueue(q *autodrainingInboundQueue) {
 	for {
 		select {
 		case elemsContainer := <-q.c:
-			elemsContainer.Lock()
+			elemsContainer.filling.Wait()
 			for _, elem := range elemsContainer.elems {
 				device.PutMessageBuffer(elem.buffer)
 				device.PutInboundElement(elem)
@@ -136,7 +136,7 @@ func (device *Device) flushOutboundQueue(q *autodrainingOutboundQueue) {
 	for {
 		select {
 		case elemsContainer := <-q.c:
-			elemsContainer.Lock()
+			elemsContainer.filling.Wait()
 			for _, elem := range elemsContainer.elems {
 				device.PutOutboundBuffer(elem.buffer)
 				device.PutOutboundElement(elem)

device/lock-ordering.md

@@ -1,27 +0,0 @@
-# Lock Ordering in wireguard-go/device
-
-## Lock hierarchy
-
-Locks must be acquired in the order listed below. A goroutine holding a
-lock with a higher number must never attempt to acquire a lock with a
-lower number.
-
-```
-Level 0  device.state.Mutex
-Level 1  device.ipcMutex                  (sync.RWMutex)
-Level 2  device.net.RWMutex
-Level 3  device.staticIdentity.RWMutex
-Level 4  device.peers.RWMutex
-Level 5  peer.state.Mutex
-Level 6  peer.handshake.mutex             (sync.RWMutex)
-Level 7  peer.keypairs.RWMutex
-Level 8  device.allowedips.mu             (sync.RWMutex)
-Level 9  device.indexTable.RWMutex
-Level 10 peer.endpoint.Mutex
-Level 11 device.cookieChecker.RWMutex
-Level 12 peer.cookieGenerator.RWMutex
-Level 13 Timer.modifyingLock / Timer.runningLock
-```
-
-Not every pair of locks appears in practice; the ordering above is the
-transitive closure of the pairs that do.

device/pools.go

@@ -74,7 +74,6 @@ func (device *Device) PopulatePools() {
 
 func (device *Device) GetInboundElementsContainer() *QueueInboundElementsContainer {
 	c := device.pool.inboundElementsContainer.Get().(*QueueInboundElementsContainer)
-	c.Mutex = sync.Mutex{}
 	return c
 }
 
@@ -88,7 +87,6 @@ func (device *Device) PutInboundElementsContainer(c *QueueInboundElementsContain
 
 func (device *Device) GetOutboundElementsContainer() *QueueOutboundElementsContainer {
 	c := device.pool.outboundElementsContainer.Get().(*QueueOutboundElementsContainer)
-	c.Mutex = sync.Mutex{}
 	return c
 }
 

device/receive.go

@@ -8,6 +8,7 @@ package device
 import (
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"net"
 	"net/netip"
 	"sync"
@@ -35,7 +36,12 @@ type QueueInboundElement struct {
 }
 
 type QueueInboundElementsContainer struct {
-	sync.Mutex
+	// 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
 }
 
@@ -178,7 +184,6 @@ func (device *Device) RoutineReceiveIncoming(maxBatchSize int, recv conn.Receive
 				elemsForPeer, ok := elemsByPeer[peer]
 				if !ok {
 					elemsForPeer = device.GetInboundElementsContainer()
-					elemsForPeer.Lock()
 					elemsByPeer[peer] = elemsForPeer
 				}
 				elemsForPeer.elems = append(elemsForPeer.elems, elem)
@@ -222,6 +227,7 @@ func (device *Device) RoutineReceiveIncoming(maxBatchSize int, recv conn.Receive
 		}
 		for peer, elemsContainer := range elemsByPeer {
 			if peer.isRunning.Load() {
+				elemsContainer.filling.Add(1)
 				peer.queue.inbound.c <- elemsContainer
 				device.queue.decryption.c <- elemsContainer
 			} else {
@@ -263,7 +269,7 @@ func (device *Device) RoutineDecryption(id int) {
 				elem.packet = nil
 			}
 		}
-		elemsContainer.Unlock()
+		elemsContainer.filling.Done()
 	}
 }
 
@@ -440,11 +446,40 @@ func (peer *Peer) RoutineSequentialReceiver(maxBatchSize int) {
 		if elemsContainer == nil {
 			return
 		}
-		elemsContainer.Lock()
+		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 elemsContainer.elems {
+	for i, elem := range elems {
 		if elem.packet == nil {
 			// decryption failed
 			continue
@@ -512,12 +547,12 @@ func (peer *Peer) RoutineSequentialReceiver(maxBatchSize int) {
 			continue
 		}
 
-			bufs = append(bufs, elem.buffer[:MessageTransportOffsetContent+len(elem.packet)])
+		scratch = append(scratch, elem.buffer[:MessageTransportOffsetContent+len(elem.packet)])
 	}
 
 	peer.rxBytes.Add(rxBytesLen)
 	if validTailPacket >= 0 {
-			peer.SetEndpointFromPacket(elemsContainer.elems[validTailPacket].endpoint)
+		peer.SetEndpointFromPacket(elems[validTailPacket].endpoint)
 		peer.keepKeyFreshReceiving()
 		peer.timersAnyAuthenticatedPacketTraversal()
 		peer.timersAnyAuthenticatedPacketReceived()
@@ -525,17 +560,14 @@ func (peer *Peer) RoutineSequentialReceiver(maxBatchSize int) {
 	if dataPacketReceived {
 		peer.timersDataReceived()
 	}
-		if len(bufs) > 0 {
+	if len(scratch) > 0 {
-			_, err := device.tun.device.Write(bufs, MessageTransportOffsetContent)
+		_, err := device.tun.device.Write(scratch, MessageTransportOffsetContent)
 		if err != nil && !device.isClosed() {
 			device.log.Errorf("Failed to write packets to TUN device: %v", err)
 		}
 	}
-		for _, elem := range elemsContainer.elems {
+	for _, elem := range elems {
 		device.PutMessageBuffer(elem.buffer)
 		device.PutInboundElement(elem)
 	}
-		bufs = bufs[:0]
-		device.PutInboundElementsContainer(elemsContainer)
-	}
 }

device/send.go

@@ -8,6 +8,7 @@ package device
 import (
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"net"
 	"net/netip"
 	"os"
@@ -58,7 +59,12 @@ type QueueOutboundElement struct {
 }
 
 type QueueOutboundElementsContainer struct {
-	sync.Mutex
+	// 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
 }
 
@@ -486,7 +492,6 @@ top:
 
 				elem.keypair = keypair
 			}
-			elemsContainer.Lock()
 			elemsContainer.elems = elemsContainer.elems[:i]
 
 			if elemsContainerOOO != nil {
@@ -502,6 +507,7 @@ top:
 
 			// add to parallel and sequential queue
 			if peer.isRunning.Load() {
+				elemsContainer.filling.Add(1)
 				peer.queue.outbound.c <- elemsContainer
 				peer.device.queue.encryption.c <- elemsContainer
 			} else {
@@ -595,7 +601,7 @@ func (device *Device) RoutineEncryption(id int) {
 			// re-slice packet to include encapsulating transport space
 			elem.packet = elem.buffer[:MessageEncapsulatingTransportSize+len(elem.packet)]
 		}
-		elemsContainer.Unlock()
+		elemsContainer.filling.Done()
 	}
 }
 
@@ -610,10 +616,38 @@ func (peer *Peer) RoutineSequentialSender(maxBatchSize int) {
 	bufs := make([][]byte, 0, maxBatchSize)
 
 	for elemsContainer := range peer.queue.outbound.c {
-		bufs = bufs[:0]
 		if elemsContainer == nil {
 			return
 		}
+		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.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
@@ -621,28 +655,26 @@ func (peer *Peer) RoutineSequentialSender(maxBatchSize int) {
 		// 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()
 		peer.queuedOutboundPackets.Add(-int32(len(elemsContainer.elems)))
 		for _, elem := range elemsContainer.elems {
 			device.PutOutboundBuffer(elem.buffer)
 			device.PutOutboundElement(elem)
 		}
-			device.PutOutboundElementsContainer(elemsContainer)
+		return
-			continue
 	}
+
 	dataSent := false
-		elemsContainer.Lock()
 	for _, elem := range elemsContainer.elems {
-			if len(elem.packet) != MessageKeepaliveSize {
+		if len(elem.packet[MessageEncapsulatingTransportSize:]) != MessageKeepaliveSize {
 			dataSent = true
 		}
-			bufs = append(bufs, elem.packet)
+		scratch = append(scratch, elem.packet)
 	}
 
 	peer.timersAnyAuthenticatedPacketTraversal()
 	peer.timersAnyAuthenticatedPacketSent()
 
-		err := peer.SendBuffers(bufs)
+	err := peer.SendBuffers(scratch)
 	if dataSent {
 		peer.timersDataSent()
 	}
@@ -651,7 +683,6 @@ func (peer *Peer) RoutineSequentialSender(maxBatchSize int) {
 		device.PutOutboundBuffer(elem.buffer)
 		device.PutOutboundElement(elem)
 	}
-		device.PutOutboundElementsContainer(elemsContainer)
 	if err != nil {
 		var errGSO conn.ErrUDPGSODisabled
 		if errors.As(err, &errGSO) {
@@ -661,9 +692,8 @@ func (peer *Peer) RoutineSequentialSender(maxBatchSize int) {
 	}
 	if err != nil {
 		device.log.Errorf("%v - Failed to send data packets: %v", peer, err)
-			continue
+		return
 	}
 
 	peer.keepKeyFreshSending()
-	}
 }