device: add priority message transmission around session establishment

device/device.go

@@ -60,7 +60,8 @@ type Device struct {
 		lookupFunc   PeerLookupFunc // or nil if unused
 	}
 
-	peerStateFn atomic.Pointer[PeerSessionStateFunc] // observes peer session state changes, nil if unset
+	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
@@ -547,6 +548,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/device_test.go

@@ -192,6 +192,86 @@ func genTestPair(tb testing.TB, realSocket bool) (pair testPair) {
 	return
 }
 
+// TestPriorityMessageOnEstablishment verifies that a message returned by a
+// [PeerPriorityMessageFunc] is transmitted to and received by the remote peer
+// when a session keypair is established, from both callback trigger sites.
+func TestPriorityMessageOnEstablishment(t *testing.T) {
+	cases := []struct {
+		name                               string
+		prioritySenderIsHandshakeInitiator bool
+	}{
+		// pair[1] is always the handshake initiator (see below), and we vary
+		// which side has the callback set (prioritySender).
+		{"received-with-keypair-trigger", false},
+		{"handshake-resp-trigger", true},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			goroutineLeakCheck(t)
+			pair := genTestPair(t, true)
+			prioritySender, priorityReceiver := pair[0], pair[1]
+			if tc.prioritySenderIsHandshakeInitiator {
+				prioritySender, priorityReceiver = pair[1], pair[0]
+			}
+
+			// Flip the bits in the priority message dest IP so we can distinguish
+			// it from a regular tuntest.Ping, which can also flow alongside
+			// depending on which side has the callback set.
+			priorityUniqueDst := priorityReceiver.ip.As4()
+			for i := range priorityUniqueDst {
+				priorityUniqueDst[i] = ^priorityUniqueDst[i]
+			}
+			priorityMsg := tuntest.Ping(netip.AddrFrom4(priorityUniqueDst), prioritySender.ip)
+
+			var gotKey atomic.Pointer[NoisePublicKey]
+			prioritySender.dev.SetPriorityMessageOnEstablishmentFunc(func(peer NoisePublicKey) []byte {
+				k := peer
+				gotKey.CompareAndSwap(nil, &k)
+				return priorityMsg
+			})
+
+			// pair[1] is always the handshake initiator
+			tunInjectedMsg := tuntest.Ping(pair[0].ip, pair[1].ip)
+			pair[1].tun.Outbound <- tunInjectedMsg
+
+			// Drain the priorityReceiver's TUN. The first message received
+			// should be the priorityMsg.
+			select {
+			case got := <-priorityReceiver.tun.Inbound:
+				if !bytes.Equal(got, priorityMsg) {
+					t.Fatal("first msg on receiver TUN is not priorityMsg")
+				}
+			case <-time.After(5 * time.Second):
+				t.Fatal("timeout waiting for priorityMsg")
+			}
+
+			// If the priority message receiver is pair[0] (handshake responder),
+			// it should also receive the tunInjectedMsg after the priorityMsg.
+			if priorityReceiver == pair[0] {
+				select {
+				case got := <-priorityReceiver.tun.Inbound:
+					if !bytes.Equal(got, tunInjectedMsg) {
+						t.Error("second msg on receiver TUN is not tunInjectedMsg")
+					}
+				case <-time.After(5 * time.Second):
+					t.Error("timeout waiting for tunInjectedMsg")
+				}
+			}
+
+			// The callback should observe the remote peer's public key.
+			priorityReceiver.dev.staticIdentity.RLock()
+			wantKey := priorityReceiver.dev.staticIdentity.publicKey
+			priorityReceiver.dev.staticIdentity.RUnlock()
+			if k := gotKey.Load(); k == nil {
+				t.Error("priority message callback was not invoked")
+			} else if *k != wantKey {
+				t.Errorf("callback peer key = %x, want %x", *k, wantKey)
+			}
+		})
+	}
+}
+
 func TestTwoDevicePing(t *testing.T) {
 	goroutineLeakCheck(t)
 	pair := genTestPair(t, true)

device/receive.go

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

device/send.go

@@ -106,6 +106,69 @@ 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.PutMessageBuffer(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.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)