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,8 +36,13 @@ type QueueInboundElement struct {
 }
 
 type QueueInboundElementsContainer struct {
-	sync.Mutex
+	// filling is a one-shot barrier signaling decryption→receive
-	elems []*QueueInboundElement
+	// 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.
@@ -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,102 +446,128 @@ func (peer *Peer) RoutineSequentialReceiver(maxBatchSize int) {
 		if elemsContainer == nil {
 			return
 		}
-		elemsContainer.Lock()
+		peer.processInboundContainer(elemsContainer, bufs[:0])
-		validTailPacket := -1
+	}
-		dataPacketReceived := false
+}
-		rxBytesLen := uint64(0)
+
-		for i, elem := range elemsContainer.elems {
+// processInboundContainer waits for the decryption routine to finish
-			if elem.packet == nil {
+// filling elemsContainer, then writes the valid packets to the TUN
-				// decryption failed
+// device and returns the container to the pool.
-				continue
+//
-			}
+// 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.
-			if !elem.keypair.replayFilter.ValidateCounter(elem.counter, RejectAfterMessages) {
+func (peer *Peer) processInboundContainer(elemsContainer *QueueInboundElementsContainer, scratch [][]byte) {
-				continue
+	// 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",
-			validTailPacket = i
+			len(scratch), cap(scratch)))
-			if peer.ReceivedWithKeypair(elem.keypair) {
+	}
-				peer.SetEndpointFromPacket(elem.endpoint)
+	if cap(scratch) < len(elemsContainer.elems) {
-				peer.timersHandshakeComplete()
+		panic(fmt.Sprintf("processInboundContainer: scratch cap %d < elems %d",
-				peer.SendStagedPackets()
+			cap(scratch), len(elemsContainer.elems)))
-			}
+	}
-			if ep, ok := elem.endpoint.(conn.PeerAwareEndpoint); ok {
+
-				ep.FromPeer(peer.handshake.remoteStatic)
+	device := peer.device
-			}
+	defer device.PutInboundElementsContainer(elemsContainer)
-			rxBytesLen += uint64(len(elem.packet) + MinMessageSize)
+
-
+	// Wait for RoutineDecryption to finish filling the container. After
-			if len(elem.packet) == 0 {
+	// Wait returns we have happens-before with that goroutine and are the
-				device.log.Verbosef("%v - Receiving keepalive packet", peer)
+	// sole owner of the container until Put hands it back to the pool.
-				continue
+	elemsContainer.filling.Wait()
-			}
+	elems := elemsContainer.elems
-			dataPacketReceived = true
+
-
+	validTailPacket := -1
-			switch elem.packet[0] >> 4 {
+	dataPacketReceived := false
-			case 4:
+	rxBytesLen := uint64(0)
-				if len(elem.packet) < ipv4.HeaderLen {
+	for i, elem := range elems {
-					continue
+		if elem.packet == nil {
-				}
+			// decryption failed
-				field := elem.packet[IPv4offsetTotalLength : IPv4offsetTotalLength+2]
+			continue
-				length := binary.BigEndian.Uint16(field)
+		}
-				if int(length) > len(elem.packet) || int(length) < ipv4.HeaderLen {
+
-					continue
+		if !elem.keypair.replayFilter.ValidateCounter(elem.counter, RejectAfterMessages) {
-				}
+			continue
-				elem.packet = elem.packet[:length]
+		}
-				src := elem.packet[IPv4offsetSrc : IPv4offsetSrc+net.IPv4len]
+
-				srcAddr, _ := netip.AddrFromSlice(src)
+		validTailPacket = i
-				if !peer.AllowedPeerSourceIP(srcAddr) {
+		if peer.ReceivedWithKeypair(elem.keypair) {
-					device.log.Verbosef("IPv4 packet with disallowed source address from %v", peer)
+			peer.SetEndpointFromPacket(elem.endpoint)
-					continue
+			peer.timersHandshakeComplete()
-				}
+			peer.SendStagedPackets()
-
+		}
-			case 6:
+		if ep, ok := elem.endpoint.(conn.PeerAwareEndpoint); ok {
-				if len(elem.packet) < ipv6.HeaderLen {
+			ep.FromPeer(peer.handshake.remoteStatic)
-					continue
+		}
-				}
+		rxBytesLen += uint64(len(elem.packet) + MinMessageSize)
-				field := elem.packet[IPv6offsetPayloadLength : IPv6offsetPayloadLength+2]
+
-				length := binary.BigEndian.Uint16(field)
+		if len(elem.packet) == 0 {
-				length += ipv6.HeaderLen
+			device.log.Verbosef("%v - Receiving keepalive packet", peer)
-				if int(length) > len(elem.packet) {
+			continue
-					continue
+		}
-				}
+		dataPacketReceived = true
-				elem.packet = elem.packet[:length]
+
-				src := elem.packet[IPv6offsetSrc : IPv6offsetSrc+net.IPv6len]
+		switch elem.packet[0] >> 4 {
-				srcAddr, _ := netip.AddrFromSlice(src)
+		case 4:
-				if !peer.AllowedPeerSourceIP(srcAddr) {
+			if len(elem.packet) < ipv4.HeaderLen {
-					device.log.Verbosef("IPv6 packet with disallowed source address from %v", peer)
+				continue
-					continue
+			}
-				}
+			field := elem.packet[IPv4offsetTotalLength : IPv4offsetTotalLength+2]
-
+			length := binary.BigEndian.Uint16(field)
-			default:
+			if int(length) > len(elem.packet) || int(length) < ipv4.HeaderLen {
-				device.log.Verbosef("Packet with invalid IP version from %v", peer)
+				continue
-				continue
+			}
-			}
+			elem.packet = elem.packet[:length]
-
+			src := elem.packet[IPv4offsetSrc : IPv4offsetSrc+net.IPv4len]
-			bufs = append(bufs, elem.buffer[:MessageTransportOffsetContent+len(elem.packet)])
+			srcAddr, _ := netip.AddrFromSlice(src)
-		}
+			if !peer.AllowedPeerSourceIP(srcAddr) {
-
+				device.log.Verbosef("IPv4 packet with disallowed source address from %v", peer)
-		peer.rxBytes.Add(rxBytesLen)
+				continue
-		if validTailPacket >= 0 {
+			}
-			peer.SetEndpointFromPacket(elemsContainer.elems[validTailPacket].endpoint)
+
-			peer.keepKeyFreshReceiving()
+		case 6:
-			peer.timersAnyAuthenticatedPacketTraversal()
+			if len(elem.packet) < ipv6.HeaderLen {
-			peer.timersAnyAuthenticatedPacketReceived()
+				continue
-		}
+			}
-		if dataPacketReceived {
+			field := elem.packet[IPv6offsetPayloadLength : IPv6offsetPayloadLength+2]
-			peer.timersDataReceived()
+			length := binary.BigEndian.Uint16(field)
-		}
+			length += ipv6.HeaderLen
-		if len(bufs) > 0 {
+			if int(length) > len(elem.packet) {
-			_, err := device.tun.device.Write(bufs, MessageTransportOffsetContent)
+				continue
-			if err != nil && !device.isClosed() {
+			}
-				device.log.Errorf("Failed to write packets to TUN device: %v", err)
+			elem.packet = elem.packet[:length]
-			}
+			src := elem.packet[IPv6offsetSrc : IPv6offsetSrc+net.IPv6len]
-		}
+			srcAddr, _ := netip.AddrFromSlice(src)
-		for _, elem := range elemsContainer.elems {
+			if !peer.AllowedPeerSourceIP(srcAddr) {
-			device.PutMessageBuffer(elem.buffer)
+				device.log.Verbosef("IPv6 packet with disallowed source address from %v", peer)
-			device.PutInboundElement(elem)
+				continue
-		}
+			}
-		bufs = bufs[:0]
+
-		device.PutInboundElementsContainer(elemsContainer)
+		default:
+			device.log.Verbosef("Packet with invalid IP version from %v", peer)
+			continue
+		}
+
+		scratch = append(scratch, elem.buffer[:MessageTransportOffsetContent+len(elem.packet)])
+	}
+
+	peer.rxBytes.Add(rxBytesLen)
+	if validTailPacket >= 0 {
+		peer.SetEndpointFromPacket(elems[validTailPacket].endpoint)
+		peer.keepKeyFreshReceiving()
+		peer.timersAnyAuthenticatedPacketTraversal()
+		peer.timersAnyAuthenticatedPacketReceived()
+	}
+	if dataPacketReceived {
+		peer.timersDataReceived()
+	}
+	if len(scratch) > 0 {
+		_, 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 elems {
+		device.PutMessageBuffer(elem.buffer)
+		device.PutInboundElement(elem)
 	}
 }

device/send.go

@@ -8,6 +8,7 @@ package device
 import (
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"net"
 	"net/netip"
 	"os"
@@ -58,8 +59,13 @@ type QueueOutboundElement struct {
 }
 
 type QueueOutboundElementsContainer struct {
-	sync.Mutex
+	// filling is a one-shot barrier signaling encryption→send handoff.
-	elems []*QueueOutboundElement
+	// 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
 }
 
 func (device *Device) NewOutboundElement() *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,60 +616,84 @@ 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
 		}
-		if !peer.isRunning.Load() {
+		peer.processOutboundContainer(elemsContainer, bufs[:0])
-			// 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()
-			peer.queuedOutboundPackets.Add(-int32(len(elemsContainer.elems)))
-			for _, elem := range elemsContainer.elems {
-				device.PutOutboundBuffer(elem.buffer)
-				device.PutOutboundElement(elem)
-			}
-			device.PutOutboundElementsContainer(elemsContainer)
-			continue
-		}
-		dataSent := false
-		elemsContainer.Lock()
-		for _, elem := range elemsContainer.elems {
-			if len(elem.packet) != MessageKeepaliveSize {
-				dataSent = true
-			}
-			bufs = append(bufs, elem.packet)
-		}
 
-		peer.timersAnyAuthenticatedPacketTraversal()
+// processOutboundContainer waits for the encryption routine to finish
-		peer.timersAnyAuthenticatedPacketSent()
+// 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)))
+	}
 
-		err := peer.SendBuffers(bufs)
+	device := peer.device
-		if dataSent {
+	defer device.PutOutboundElementsContainer(elemsContainer)
-			peer.timersDataSent()
+
-		}
+	// 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
+		// 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.
 		peer.queuedOutboundPackets.Add(-int32(len(elemsContainer.elems)))
 		for _, elem := range elemsContainer.elems {
 			device.PutOutboundBuffer(elem.buffer)
 			device.PutOutboundElement(elem)
 		}
-		device.PutOutboundElementsContainer(elemsContainer)
+		return
-		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()
 	}
+
+	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()
+	}
+	peer.queuedOutboundPackets.Add(-int32(len(elemsContainer.elems)))
+	for _, elem := range elemsContainer.elems {
+		device.PutOutboundBuffer(elem.buffer)
+		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()
 }