conn,device: implement InitiationAwareEndpoint

To be implemented by [magicsock.lazyEndpoint], which is responsible for
triggering JIT peer configuration.

Updates tailscale/corp#20732
Updates tailscale/corp#30042

Signed-off-by: Jordan Whited <jordan@tailscale.com>

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.

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

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,

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