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

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.
 	//