handshake: Add ix_psk2_mlkem1024, IX with an ephemeral ML-KEM-1024 exchange in the psk2 slot

handshake/credential.go

@@ -35,7 +35,9 @@ func NewCredential(
 }
 
 // buildHandshakeState creates a noise.HandshakeState from this credential.
-func (hc *Credential) buildHandshakeState(initiator bool, pattern noise.HandshakePattern) (*noise.HandshakeState, error) {
+// pskPlacement selects the noise psk modifier; 0 means none. A psk for
+// placements 2 and up is supplied later with SetPresharedKey.
+func (hc *Credential) buildHandshakeState(initiator bool, pattern noise.HandshakePattern, pskPlacement int) (*noise.HandshakeState, error) {
 	return noise.NewHandshakeState(noise.Config{
 		CipherSuite:           hc.cipherSuite,
 		Random:                rand.Reader,
@@ -43,7 +45,7 @@ func (hc *Credential) buildHandshakeState(initiator bool, pattern noise.Handshak
 		Initiator:             initiator,
 		StaticKeypair:         noise.DHKey{Private: hc.privateKey, Public: hc.Cert.PublicKey()},
 		PresharedKey:          []byte{},
-		PresharedKeyPlacement: 0,
+		PresharedKeyPlacement: pskPlacement,
 	})
 }
 

handshake/handshake.proto

@@ -24,6 +24,9 @@ message NebulaHandshakeDetails {
   uint64 Cookie = 4 [deprecated = true];
   uint64 Time = 5;
   uint32 CertVersion = 8;
+  // KEMKey is the initiator's ephemeral ML-KEM-1024 encapsulation key. It is
+  // only present in msg1 of the ix_psk2_mlkem1024 subtype.
+  bytes KEMKey = 9;
   // reserved for WIP multiport
   reserved 6, 7;
 }

handshake/kem.go

@@ -0,0 +1,68 @@
+package handshake
+
+import (
+	"crypto/hkdf"
+	"crypto/mlkem"
+	"crypto/sha256"
+	"errors"
+	"fmt"
+)
+
+// kemPSKSalt separates the psk derivation from any other use of the secret.
+const kemPSKSalt = "nebula ix_psk2_mlkem1024 v1"
+
+var ErrInvalidKEMKey = errors.New("peer sent an invalid ML-KEM encapsulation key")
+
+// kemPSK derives the noise psk from an ML-KEM shared secret. The ciphertext is
+// bound in because, unlike the encapsulation key, it travels outside the noise
+// transcript.
+func kemPSK(sharedKey, ciphertext []byte) ([]byte, error) {
+	secret := make([]byte, 0, len(sharedKey)+len(ciphertext))
+	secret = append(secret, sharedKey...)
+	secret = append(secret, ciphertext...)
+	return hkdf.Extract(sha256.New, secret, []byte(kemPSKSalt))
+}
+
+// kemEncapsulate is the responder half: encapsulate to the initiator's key,
+// install the psk, and remember the ciphertext to prefix the response with.
+func (m *Machine) kemEncapsulate(encapsulationKey []byte) error {
+	ek, err := mlkem.NewEncapsulationKey1024(encapsulationKey)
+	if err != nil {
+		return fmt.Errorf("%w: %w", ErrInvalidKEMKey, err)
+	}
+	sharedKey, ct := ek.Encapsulate()
+	psk, err := kemPSK(sharedKey, ct)
+	if err != nil {
+		return err
+	}
+	if err := m.hs.SetPresharedKey(psk); err != nil {
+		return err
+	}
+	m.kemCiphertext = ct
+	return nil
+}
+
+// kemDecapsulate is the initiator half: recover the shared secret from the
+// ciphertext prefixing the response and install the psk. It returns the rest of
+// the message, which is the noise message.
+func (m *Machine) kemDecapsulate(msg []byte) ([]byte, error) {
+	if len(msg) < mlkem.CiphertextSize1024 {
+		return nil, ErrPacketTooShort
+	}
+	if m.kemKey == nil {
+		return nil, errors.New("no ML-KEM decapsulation key for this handshake")
+	}
+	ct := msg[:mlkem.CiphertextSize1024]
+	sharedKey, err := m.kemKey.Decapsulate(ct)
+	if err != nil {
+		return nil, err
+	}
+	psk, err := kemPSK(sharedKey, ct)
+	if err != nil {
+		return nil, err
+	}
+	if err := m.hs.SetPresharedKey(psk); err != nil {
+		return nil, err
+	}
+	return msg[mlkem.CiphertextSize1024:], nil
+}

handshake/kem_test.go

@@ -0,0 +1,208 @@
+package handshake
+
+import (
+	"bytes"
+	"crypto/mlkem"
+	"net/netip"
+	"testing"
+	"time"
+
+	"github.com/flynn/noise"
+	"github.com/slackhq/nebula/cert"
+	ct "github.com/slackhq/nebula/cert_test"
+	"github.com/slackhq/nebula/header"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func newKEMMachine(t *testing.T, cs *testCertState, v CertVerifier, initiator bool, index uint32) *Machine {
+	t.Helper()
+	m, err := NewMachine(
+		cs.version, cs.getCredential, v,
+		func() (uint32, error) { return index, nil },
+		initiator, header.HandshakeIXPSK2MLKEM1024,
+	)
+	require.NoError(t, err)
+	return m
+}
+
+type kemFixture struct {
+	initCS, respCS *testCertState
+	v              CertVerifier
+}
+
+func newKEMFixture(t *testing.T, cipher noise.CipherFunc) kemFixture {
+	ca, _, caKey, _ := ct.NewTestCaCert(
+		cert.Version2, cert.Curve_CURVE25519, time.Time{}, time.Time{}, nil, nil, nil,
+	)
+	return kemFixture{
+		initCS: newTestCertStateWithCipher(t, ca, caKey, "initiator", []netip.Prefix{netip.MustParsePrefix("10.0.0.1/24")}, cipher),
+		respCS: newTestCertStateWithCipher(t, ca, caKey, "responder", []netip.Prefix{netip.MustParsePrefix("10.0.0.2/24")}, cipher),
+		v:      testVerifier(ct.NewTestCAPool(ca)),
+	}
+}
+
+func requireKeysMatch(t *testing.T, a, b *Result) {
+	t.Helper()
+	c, err := a.EKey.Encrypt(nil, nil, []byte("hello"))
+	require.NoError(t, err)
+	p, err := b.DKey.Decrypt(nil, nil, c)
+	require.NoError(t, err)
+	assert.Equal(t, []byte("hello"), p)
+
+	c, err = b.EKey.Encrypt(nil, nil, []byte("world"))
+	require.NoError(t, err)
+	p, err = a.DKey.Decrypt(nil, nil, c)
+	require.NoError(t, err)
+	assert.Equal(t, []byte("world"), p)
+}
+
+func TestKEMHandshake(t *testing.T) {
+	for _, cipher := range []noise.CipherFunc{noise.CipherChaChaPoly, noise.CipherAESGCM} {
+		t.Run(cipher.CipherName(), func(t *testing.T) {
+			f := newKEMFixture(t, cipher)
+			initM := newKEMMachine(t, f.initCS, f.v, true, 1000)
+			respM := newKEMMachine(t, f.respCS, f.v, false, 2000)
+
+			msg1, err := initM.Initiate(nil)
+			require.NoError(t, err)
+			assert.Equal(t, header.HandshakeIXPSK2MLKEM1024, header.MessageSubType(msg1[1]))
+			assert.Greater(t, len(msg1), header.Len+mlkem.EncapsulationKeySize1024)
+
+			msg2, respR, err := respM.ProcessPacket(nil, msg1)
+			require.NoError(t, err)
+			require.NotNil(t, respR)
+			assert.Greater(t, len(msg2), header.Len+mlkem.CiphertextSize1024)
+
+			_, initR, err := initM.ProcessPacket(nil, msg2)
+			require.NoError(t, err)
+			require.NotNil(t, initR)
+
+			assert.Equal(t, "responder", initR.RemoteCert.Certificate.Name())
+			assert.Equal(t, "initiator", respR.RemoteCert.Certificate.Name())
+			assert.Equal(t, uint32(2000), initR.RemoteIndex)
+			assert.Equal(t, uint32(1000), respR.RemoteIndex)
+			requireKeysMatch(t, initR, respR)
+		})
+	}
+}
+
+// TestKEMSecretMatters checks the ML-KEM secret really reaches the transport
+// keys: an initiator holding a different decapsulation key cannot complete even
+// though every DH input is genuine.
+func TestKEMSecretMatters(t *testing.T) {
+	f := newKEMFixture(t, noise.CipherChaChaPoly)
+	initM := newKEMMachine(t, f.initCS, f.v, true, 1000)
+	respM := newKEMMachine(t, f.respCS, f.v, false, 2000)
+
+	msg1, err := initM.Initiate(nil)
+	require.NoError(t, err)
+	msg2, _, err := respM.ProcessPacket(nil, msg1)
+	require.NoError(t, err)
+
+	other, err := mlkem.GenerateKey1024()
+	require.NoError(t, err)
+	initM.kemKey = other
+
+	_, res, err := initM.ProcessPacket(nil, msg2)
+	require.Error(t, err)
+	assert.Nil(t, res)
+}
+
+func TestKEMTamperedCiphertext(t *testing.T) {
+	f := newKEMFixture(t, noise.CipherChaChaPoly)
+	initM := newKEMMachine(t, f.initCS, f.v, true, 1000)
+	respM := newKEMMachine(t, f.respCS, f.v, false, 2000)
+
+	msg1, err := initM.Initiate(nil)
+	require.NoError(t, err)
+	msg2, respR, err := respM.ProcessPacket(nil, msg1)
+	require.NoError(t, err)
+
+	t.Run("truncated", func(t *testing.T) {
+		_, res, err := initM.ProcessPacket(nil, msg2[:header.Len+mlkem.CiphertextSize1024-1])
+		require.ErrorIs(t, err, ErrPacketTooShort)
+		assert.Nil(t, res)
+		assert.False(t, initM.Failed())
+	})
+
+	t.Run("flipped ciphertext bit", func(t *testing.T) {
+		bad := bytes.Clone(msg2)
+		bad[header.Len+10] ^= 1
+		_, res, err := initM.ProcessPacket(nil, bad)
+		require.Error(t, err)
+		assert.Nil(t, res)
+		assert.False(t, initM.Failed(), "a forged ciphertext must not kill the handshake")
+	})
+
+	t.Run("genuine response still completes", func(t *testing.T) {
+		_, initR, err := initM.ProcessPacket(nil, msg2)
+		require.NoError(t, err)
+		require.NotNil(t, initR)
+		requireKeysMatch(t, initR, respR)
+	})
+}
+
+// TestKEMBadMsg1 drives processPayload by hand. With the psk modifier msg1's
+// payload is encrypted under a key derived from the public ephemeral, so a
+// forged payload would need re-encrypting; the parsing rules are what matter.
+func TestKEMBadMsg1(t *testing.T) {
+	f := newKEMFixture(t, noise.CipherChaChaPoly)
+	msg1Flags := subtypeInfos[header.HandshakeIXPSK2MLKEM1024].msgs[0]
+	cert := f.initCS.creds[cert.Version2]
+	payload := func(kemKey []byte) []byte {
+		return MarshalPayload(nil, Payload{
+			Cert:           cert.Bytes,
+			CertVersion:    uint32(cert.Cert.Version()),
+			InitiatorIndex: 1000,
+			Time:           1,
+			KEMKey:         kemKey,
+		})
+	}
+	dk, err := mlkem.GenerateKey1024()
+	require.NoError(t, err)
+
+	t.Run("missing KEM key", func(t *testing.T) {
+		m := newKEMMachine(t, f.respCS, f.v, false, 2000)
+		err := m.processPayload(payload(nil), msg1Flags)
+		require.ErrorIs(t, err, ErrUnexpectedContent)
+		assert.True(t, m.Failed())
+	})
+
+	t.Run("invalid KEM key", func(t *testing.T) {
+		m := newKEMMachine(t, f.respCS, f.v, false, 2000)
+		// Coefficients of 0xfff exceed the modulus, which ML-KEM must reject.
+		bad := bytes.Repeat([]byte{0xff}, mlkem.EncapsulationKeySize1024)
+		err := m.kemEncapsulate(bad)
+		require.ErrorIs(t, err, ErrInvalidKEMKey)
+		err = m.kemEncapsulate(bad[:100])
+		require.ErrorIs(t, err, ErrInvalidKEMKey)
+	})
+
+	t.Run("valid KEM key", func(t *testing.T) {
+		m := newKEMMachine(t, f.respCS, f.v, false, 2000)
+		require.NoError(t, m.kemEncapsulate(dk.EncapsulationKey().Bytes()))
+		assert.Len(t, m.kemCiphertext, mlkem.CiphertextSize1024)
+	})
+
+	t.Run("KEM key on the classical subtype", func(t *testing.T) {
+		m := newTestMachine(t, f.respCS, f.v, false, 2000)
+		classical := subtypeInfos[header.HandshakeIXPSK0].msgs[0]
+		err := m.processPayload(payload(dk.EncapsulationKey().Bytes()), classical)
+		require.ErrorIs(t, err, ErrUnexpectedContent)
+		assert.True(t, m.Failed())
+	})
+}
+
+// TestKEMDistinctFromClassical makes sure the two subtypes cannot be mixed: a
+// classical responder cannot answer a KEM initiator.
+func TestKEMDistinctFromClassical(t *testing.T) {
+	f := newKEMFixture(t, noise.CipherChaChaPoly)
+	initM := newKEMMachine(t, f.initCS, f.v, true, 1000)
+	respM := newTestMachine(t, f.respCS, f.v, false, 2000)
+
+	msg1, err := initM.Initiate(nil)
+	require.NoError(t, err)
+	_, _, err = respM.ProcessPacket(nil, msg1)
+	require.ErrorIs(t, err, ErrSubtypeMismatch)
+}

handshake/machine.go

@@ -2,6 +2,8 @@ package handshake
 
 import (
 	"bytes"
+	"crypto/mlkem"
+	"errors"
 	"fmt"
 	"slices"
 	"time"
@@ -67,6 +69,9 @@ type Machine struct {
 	remoteCertSet  bool
 	payloadSet     bool
 	failed         bool
+
+	kemKey        *mlkem.DecapsulationKey1024 // initiator's ephemeral ML-KEM key, until completion
+	kemCiphertext []byte                      // responder's ciphertext to prefix its response with
 }
 
 // NewMachine creates a handshake state machine. The subtype determines both
@@ -91,7 +96,7 @@ func NewMachine(
 		return nil, fmt.Errorf("%w: %v", ErrNoCredential, version)
 	}
 
-	hs, err := cred.buildHandshakeState(initiator, info.pattern)
+	hs, err := cred.buildHandshakeState(initiator, info.pattern, info.pskPlacement)
 	if err != nil {
 		return nil, fmt.Errorf("build noise state: %w", err)
 	}
@@ -137,7 +142,8 @@ func (m *Machine) requireComplete() error {
 	return nil
 }
 
-// myMsgFlags returns the flags for the current outgoing message.
+// myMsgFlags returns the flags for the current outgoing message, which are
+// also the flags of the next incoming message before it is read.
 func (m *Machine) myMsgFlags() msgFlags {
 	idx := m.hs.MessageIndex()
 	if idx < len(m.msgs) {
@@ -225,7 +231,18 @@ func (m *Machine) ProcessPacket(out, packet []byte) ([]byte, *Result, error) {
 	// noise returns (cs1, cs2) where cs1 is the initiator->responder cipher.
 	// For 3-message patterns where a responder finishes by reading the final
 	// message, this ordering would be wrong; revisit when XX/pqIX lands.
-	msg, eKey, dKey, err := m.hs.ReadMessage(nil, packet[header.Len:])
+	noiseMsg := packet[header.Len:]
+	if m.myMsgFlags().kemCiphertext {
+		var err error
+		// Installing a psk derived from a forged ciphertext is harmless: ReadMessage
+		// then fails without advancing, and the next packet installs its own.
+		noiseMsg, err = m.kemDecapsulate(noiseMsg)
+		if err != nil {
+			return nil, nil, err
+		}
+	}
+
+	msg, eKey, dKey, err := m.hs.ReadMessage(nil, noiseMsg)
 	if err != nil {
 		// Noise ReadMessage failed. The noise library checkpoints and rolls back
 		// on failure, so the Machine is still alive. The caller can retry with
@@ -276,6 +293,8 @@ func (m *Machine) ProcessPacket(out, packet []byte) ([]byte, *Result, error) {
 }
 
 func (m *Machine) completed(eKey, dKey *noise.CipherState) *Result {
+	m.kemKey = nil
+	m.kemCiphertext = nil
 	m.result.EKey = eKey
 	m.result.DKey = dKey
 	m.result.MessageIndex = uint64(m.hs.MessageIndex())
@@ -284,7 +303,7 @@ func (m *Machine) completed(eKey, dKey *noise.CipherState) *Result {
 
 func (m *Machine) processPayload(msg []byte, flags msgFlags) error {
 	if len(msg) == 0 {
-		if flags.expectsPayload || flags.expectsCert {
+		if flags.expectsPayload || flags.expectsCert || flags.kemKey {
 			m.failed = true
 			return ErrMissingContent
 		}
@@ -310,6 +329,12 @@ func (m *Machine) processPayload(msg []byte, flags msgFlags) error {
 		return ErrUnexpectedContent
 	}
 
+	hasKEMKey := len(payload.KEMKey) > 0
+	if hasKEMKey != flags.kemKey {
+		m.failed = true
+		return ErrUnexpectedContent
+	}
+
 	// Process payload
 	if flags.expectsPayload {
 		var remoteIndex uint32
@@ -336,6 +361,13 @@ func (m *Machine) processPayload(msg []byte, flags msgFlags) error {
 		}
 	}
 
+	if flags.kemKey {
+		if err := m.kemEncapsulate(payload.KEMKey); err != nil {
+			m.failed = true
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -380,11 +412,21 @@ func (m *Machine) validateCert(payload Payload) error {
 }
 
 func (m *Machine) marshalOutgoing(flags msgFlags) ([]byte, error) {
-	if !flags.expectsPayload && !flags.expectsCert {
+	if !flags.expectsPayload && !flags.expectsCert && !flags.kemKey {
 		return nil, nil
 	}
 
 	var p Payload
+	if flags.kemKey {
+		if m.kemKey == nil {
+			dk, err := mlkem.GenerateKey1024()
+			if err != nil {
+				return nil, fmt.Errorf("generate ML-KEM key: %w", err)
+			}
+			m.kemKey = dk
+		}
+		p.KEMKey = m.kemKey.EncapsulationKey().Bytes()
+	}
 	if flags.expectsPayload {
 		if !m.indexAllocated {
 			index, err := m.allocIndex()
@@ -436,6 +478,13 @@ func (m *Machine) buildResponse(out []byte) ([]byte, *noise.CipherState, *noise.
 		uint64(m.hs.MessageIndex()+1),
 	)
 
+	if flags.kemCiphertext {
+		if m.kemCiphertext == nil {
+			return nil, nil, nil, errors.New("no ML-KEM ciphertext to send")
+		}
+		out = append(out, m.kemCiphertext...)
+	}
+
 	// noise.WriteMessage appends the encrypted handshake message to out,
 	// reusing capacity when present.
 	//

handshake/patterns.go

@@ -11,13 +11,16 @@ import (
 type msgFlags struct {
 	expectsPayload bool // message carries indexes and time
 	expectsCert    bool // message carries the certificate
+	kemKey         bool // payload carries an ML-KEM encapsulation key
+	kemCiphertext  bool // message is prefixed with the ML-KEM ciphertext for the peer's key
 }
 
 // subtypeInfo bundles the noise pattern with the per-message flags for a
 // given handshake subtype.
 type subtypeInfo struct {
-	pattern noise.HandshakePattern
-	msgs    []msgFlags
+	pattern      noise.HandshakePattern
+	pskPlacement int // noise psk modifier position, 0 means no psk token
+	msgs         []msgFlags
 }
 
 // subtypeInfos defines the noise pattern and message content layout for each
@@ -32,6 +35,22 @@ var subtypeInfos = map[header.MessageSubType]subtypeInfo{
 		},
 	},
 
+	// IX with an ephemeral ML-KEM-1024 exchange:
+	//   msg1 (I->R): payload + cert + encapsulation key
+	//   msg2 (R->I): ciphertext || noise(payload + cert)
+	// The responder encapsulates after reading msg1 and mixes the shared secret in
+	// through the psk2 token at the end of msg2. The ciphertext travels outside the
+	// noise message because the initiator needs the secret before it can process
+	// that token.
+	header.HandshakeIXPSK2MLKEM1024: {
+		pattern:      noise.HandshakeIX,
+		pskPlacement: 2,
+		msgs: []msgFlags{
+			{expectsPayload: true, expectsCert: true, kemKey: true},
+			{expectsPayload: true, expectsCert: true, kemCiphertext: true},
+		},
+	},
+
 	// XX: 3 messages
 	//   msg1 (I->R): payload only
 	//   msg2 (R->I): payload + cert

handshake/payload.go

@@ -20,6 +20,7 @@ type Payload struct {
 	ResponderIndex uint32
 	Time           uint64
 	CertVersion    uint32
+	KEMKey         []byte // ML-KEM encapsulation key, only in subtypes that carry one
 }
 
 // Proto field numbers for NebulaHandshakeDetails
@@ -29,6 +30,7 @@ const (
 	fieldResponderIndex = 3 // uint32
 	fieldTime           = 5 // uint64
 	fieldCertVersion    = 8 // uint32
+	fieldKEMKey         = 9 // bytes
 )
 
 // MarshalPayload encodes a handshake payload in protobuf wire format compatible
@@ -57,6 +59,10 @@ func MarshalPayload(out []byte, p Payload) []byte {
 		details = protowire.AppendTag(details, fieldCertVersion, protowire.VarintType)
 		details = protowire.AppendVarint(details, uint64(p.CertVersion))
 	}
+	if len(p.KEMKey) > 0 {
+		details = protowire.AppendTag(details, fieldKEMKey, protowire.BytesType)
+		details = protowire.AppendBytes(details, p.KEMKey)
+	}
 
 	out = protowire.AppendTag(out, 1, protowire.BytesType)
 	out = protowire.AppendBytes(out, details)
@@ -161,6 +167,16 @@ func unmarshalPayloadDetails(p *Payload, b []byte) error {
 			}
 			p.CertVersion = uint32(v)
 			b = b[n:]
+		case fieldKEMKey:
+			if typ != protowire.BytesType {
+				return errInvalidHandshakeDetails
+			}
+			v, n := protowire.ConsumeBytes(b)
+			if n < 0 {
+				return errInvalidHandshakeDetails
+			}
+			p.KEMKey = append([]byte(nil), v...)
+			b = b[n:]
 		default:
 			n := protowire.ConsumeFieldValue(num, typ, b)
 			if n < 0 {