diff --git a/device/device.go b/device/device.go index a2643e0d5..9e000a7b9 100644 --- a/device/device.go +++ b/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() diff --git a/device/device_test.go b/device/device_test.go index e44342170..05ada280c 100644 --- a/device/device_test.go +++ b/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) diff --git a/device/receive.go b/device/receive.go index d2cc46cb2..2ce71e9a4 100644 --- a/device/receive.go +++ b/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 { diff --git a/device/send.go b/device/send.go index 497c39930..e0239789f 100644 --- a/device/send.go +++ b/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)