Update

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)
 	}
 }