diff --git a/conn/conn.go b/conn/conn.go index 2a04e6d..b949641 100644 --- a/conn/conn.go +++ b/conn/conn.go @@ -86,6 +86,21 @@ 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. diff --git a/device/noise-protocol.go b/device/noise-protocol.go index 555ce91..ad5838e 100644 --- a/device/noise-protocol.go +++ b/device/noise-protocol.go @@ -16,6 +16,7 @@ import ( "golang.org/x/crypto/chacha20poly1305" "golang.org/x/crypto/poly1305" + "github.com/tailscale/wireguard-go/conn" "github.com/tailscale/wireguard-go/tai64n" ) @@ -338,7 +339,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 @@ -372,6 +373,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/noise_test.go b/device/noise_test.go index 7d6af1d..160bee5 100644 --- a/device/noise_test.go +++ b/device/noise_test.go @@ -8,6 +8,7 @@ package device import ( "bytes" "encoding/binary" + "net/netip" "testing" "github.com/tailscale/wireguard-go/conn" @@ -56,6 +57,26 @@ func assertEqual(t *testing.T, a, b []byte) { } } +type initAwareEP struct { + calledWith *[32]byte +} + +var _ conn.Endpoint = (*initAwareEP)(nil) +var _ conn.InitiationAwareEndpoint = (*initAwareEP)(nil) + +func (i *initAwareEP) ClearSrc() {} +func (i *initAwareEP) SrcToString() string { return "" } +func (i *initAwareEP) DstToString() string { return "" } +func (i *initAwareEP) DstToBytes() []byte { return nil } +func (i *initAwareEP) DstIP() netip.Addr { return netip.Addr{} } +func (i *initAwareEP) SrcIP() netip.Addr { return netip.Addr{} } + +func (i *initAwareEP) InitiationMessagePublicKey(peerPublicKey [32]byte) { + calledWith := [32]byte{} + copy(calledWith[:], peerPublicKey[:]) + i.calledWith = &calledWith +} + func TestNoiseHandshake(t *testing.T) { dev1 := randDevice(t) dev2 := randDevice(t) @@ -93,10 +114,17 @@ func TestNoiseHandshake(t *testing.T) { writer := bytes.NewBuffer(packet) err = binary.Write(writer, binary.LittleEndian, msg1) assertNil(t, err) - peer := dev2.ConsumeMessageInitiation(msg1) + initEP := &initAwareEP{} + peer := dev2.ConsumeMessageInitiation(msg1, initEP) if peer == nil { t.Fatal("handshake failed at initiation message") } + if initEP.calledWith == nil { + t.Fatal("initAwareEP never called") + } + if *initEP.calledWith != dev1.staticIdentity.publicKey { + t.Fatal("initAwareEP called with unexpected public key") + } assertEqual( t, diff --git a/device/receive.go b/device/receive.go index bc37f91..e74de1a 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