diff --git a/conn/conn.go b/conn/conn.go index 139b06e..87b7510 100644 --- a/conn/conn.go +++ b/conn/conn.go @@ -86,6 +86,40 @@ type Endpoint interface { SrcIP() netip.Addr } +// InitiationAwareEndpoint is an optional [Endpoint] specialization for +// integrations that want to know when a WireGuard handshake initiation +// message has been received, enabling just-in-time peer configuration before +// attempted decryption. +// +// It's most useful when used in combination with [PeerAwareEndpoint], enabling +// JIT peer configuration and post-decryption peer verification from a single +// implementer. +type InitiationAwareEndpoint interface { + // InitiationMessagePublicKey is called when a handshake initiation message + // has been received, and the sender's public key has been identified, but + // BEFORE an attempt has been made to verify it. + InitiationMessagePublicKey(peerPublicKey [32]byte) +} + +// PeerAwareEndpoint is an optional Endpoint specialization for +// integrations that want to know about the outcome of Cryptokey Routing +// identification. +// +// If they receive a packet from a source they had not pre-identified, +// to learn the identification WireGuard can derive from the session +// or handshake. +// +// A [PeerAwareEndpoint] may be installed as the [conn.Endpoint] following +// successful decryption unless endpoint roaming has been disabled for +// the peer. +type PeerAwareEndpoint interface { + // FromPeer is called at least once per successfully Cryptokey Routing ID'd + // [ReceiveFunc] packets batch for a given node key. wireguard-go will + // always call it for the latest/tail packet in the batch, only ever + // suppressing calls for older packets. + FromPeer(peerPublicKey [32]byte) +} + var ( ErrBindAlreadyOpen = errors.New("bind is already open") ErrWrongEndpointType = errors.New("endpoint type does not correspond with bind type") diff --git a/device/noise-protocol.go b/device/noise-protocol.go index 98a03bc..ed4d82a 100644 --- a/device/noise-protocol.go +++ b/device/noise-protocol.go @@ -12,6 +12,7 @@ import ( "sync" "time" + "github.com/sagernet/wireguard-go/conn" "github.com/sagernet/wireguard-go/tai64n" "golang.org/x/crypto/blake2s" "golang.org/x/crypto/chacha20poly1305" @@ -337,7 +338,7 @@ func (device *Device) CreateMessageInitiation(peer *Peer) (*MessageInitiation, e return &msg, nil } -func (device *Device) ConsumeMessageInitiation(msg *MessageInitiation) *Peer { +func (device *Device) ConsumeMessageInitiation(msg *MessageInitiation, endpoint conn.Endpoint) *Peer { var ( hash [blake2s.Size]byte chainKey [blake2s.Size]byte @@ -371,6 +372,11 @@ func (device *Device) ConsumeMessageInitiation(msg *MessageInitiation) *Peer { // lookup peer + initEP, ok := endpoint.(conn.InitiationAwareEndpoint) + if ok { + initEP.InitiationMessagePublicKey(peerPK) + } + peer := device.LookupPeer(peerPK) if peer == nil || !peer.isRunning.Load() { return nil diff --git a/device/receive.go b/device/receive.go index 2c40730..e13c987 100644 --- a/device/receive.go +++ b/device/receive.go @@ -359,7 +359,7 @@ func (device *Device) RoutineHandshake(id int) { // consume initiation - peer := device.ConsumeMessageInitiation(&msg) + peer := device.ConsumeMessageInitiation(&msg, elem.endpoint) if peer == nil { device.log.Verbosef("Received invalid initiation message from %s", elem.endpoint.DstToString()) goto skip @@ -459,6 +459,9 @@ func (peer *Peer) RoutineSequentialReceiver(maxBatchSize int) { peer.timersHandshakeComplete() peer.SendStagedPackets() } + if ep, ok := elem.endpoint.(conn.PeerAwareEndpoint); ok { + ep.FromPeer(peer.handshake.remoteStatic) + } rxBytesLen += uint64(len(elem.packet) + MinMessageSize) if len(elem.packet) == 0 {