diff --git a/device/device.go b/device/device.go index a826c6d..4cb1afe 100644 --- a/device/device.go +++ b/device/device.go @@ -64,7 +64,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 @@ -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() diff --git a/device/receive.go b/device/receive.go index e11e30a..3d42639 100644 --- a/device/receive.go +++ b/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 { diff --git a/device/send.go b/device/send.go index fae5c41..f5561d7 100644 --- a/device/send.go +++ b/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)