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