device: add priority message transmission around session establishment

Add SetPriorityMessageOnEstablishmentFunc, which registers a
PeerPriorityMessageFunc callback invoked when a peer's session keypair
is established or re-keyed for forward data transmission. The bytes it
returns are transmitted to the peer as a transport message.

The message is "priority" in two senses: it bypasses the staged packet
queue entirely, so it cannot be evicted by TUN-sourced packets, and it
is enqueued ahead of the keepalive/staged packets that follow keypair
establishment.

Updates tailscale/tailscale#20081

Signed-off-by: Jordan Whited <jordan@tailscale.com>

device/device.go

@@ -65,6 +65,7 @@ type Device struct {
 	}
 
 	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
@@ -553,6 +554,38 @@ func (device *Device) SetSessionStateFunc(f PeerSessionStateFunc) {
 	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)
+}
+
 func (device *Device) Close() {
 	device.state.Lock()
 	defer device.state.Unlock()

device/receive.go

@@ -425,6 +425,7 @@ func (device *Device) RoutineHandshake(id int) {
 
 			peer.timersSessionDerived()
 			peer.timersHandshakeComplete()
+			peer.SendPriorityMessage()
 			peer.SendKeepalive()
 		}
 	skip:
@@ -493,6 +494,7 @@ func (peer *Peer) processInboundContainer(elemsContainer *QueueInboundElementsCo
 		if peer.ReceivedWithKeypair(elem.keypair) {
 			peer.SetEndpointFromPacket(elem.endpoint)
 			peer.timersHandshakeComplete()
+			peer.SendPriorityMessage()
 			peer.SendStagedPackets()
 		}
 		if ep, ok := elem.endpoint.(conn.PeerAwareEndpoint); ok {

device/send.go

@@ -107,6 +107,70 @@ 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.NewOutboundElement()
+	elemsContainer := peer.device.GetOutboundElementsContainer()
+	elemsContainer.elems = append(elemsContainer.elems, elem)
+	packetQueued := false
+	defer func() {
+		if !packetQueued {
+			peer.device.PutOutboundBuffer(elem.buffer)
+			peer.device.PutOutboundElement(elem)
+			peer.device.PutOutboundElementsContainer(elemsContainer)
+		}
+	}()
+
+	// initialize outbound element
+	const offset = MessageEncapsulatingTransportSize + MessageTransportHeaderSize
+	n := copy(elem.buffer[offset:], msg)
+	elem.packet = elem.buffer[offset : offset+n]
+	elem.peer = peer
+	elem.nonce = keypair.sendNonce.Add(1) - 1
+	if elem.nonce >= RejectAfterMessages {
+		keypair.sendNonce.Store(RejectAfterMessages)
+		return
+	}
+	elem.keypair = keypair
+
+	// add to parallel and sequential queue
+	if peer.isRunning.Load() {
+		elemsContainer.filling.Add(1)
+		peer.queuedOutboundPackets.Add(1)
+		peer.queue.outbound.c <- elemsContainer
+		peer.device.queue.encryption.c <- elemsContainer
+		packetQueued = true
+	}
+}
+
 func (peer *Peer) SendHandshakeInitiation(isRetry bool) error {
 	if !isRetry {
 		peer.timers.handshakeAttempts.Store(0)