engine/src/main/kotlin/io/flamenet/e2e/DoubleRatchet.kt

@@ -47,6 +47,24 @@ class DoubleRatchet private constructor(
     private var received: Int,
     /** Length of the PREVIOUS sending chain, so the far side knows how many to skip. */
     private var previousChainLength: Int,
+    /**
+     * Our current ML-KEM keypair (spec §13.2 `PQs`). Null for a v1/v2 session,
+     * which is what a peer that does not advertise v3 gets.
+     */
+    private var pqSelf: PQKEM.KeyPair? = null,
+    /** The peer's current ML-KEM public key (`PQr`). */
+    private var pqRemote: ByteArray? = null,
+    /**
+     * `{pub, ct}` to attach to outgoing messages (`PQpending`).
+     *
+     * ⚠️ **Repeated on every message of the sending chain, not just the first
+     * (§13.3).** Sending it once would mean losing that single message strands
+     * the whole chain: the receiver never gets the material to advance its root
+     * key, and nothing later in the chain can be decrypted. It is replaced by
+     * the next ratchet step, and is NOT cleared when the peer answers — it
+     * belongs to the sending chain.
+     */
+    private var pqPending: Pair<ByteArray, ByteArray>? = null,
 ) {
 
     /** Keys for messages that arrived out of order, or have not arrived yet. */
@@ -54,7 +72,10 @@ class DoubleRatchet private constructor(
 
     data class SkippedKey(val ratchetPublic: String, val index: Int)
 
-    data class Header(val dh: ByteArray, val pn: Int, val n: Int) {
+    data class Header(
+        val dh: ByteArray, val pn: Int, val n: Int,
+        val pqdh: ByteArray? = null, val pqct: ByteArray? = null,
+    ) {
         /**
          * The exact bytes that go into the AAD.
          *
@@ -63,15 +84,32 @@ class DoubleRatchet private constructor(
          * free to reorder keys or add spaces between releases. Key order is
          * `dh, pn, n` and there are no spaces — both are part of the protocol.
          */
-        fun canonicalJson(): ByteArray =
-            ("{\"dh\":\"" + Bytes.b64(dh) + "\",\"pn\":" + pn + ",\"n\":" + n + "}")
-                .toByteArray(Charsets.UTF_8)
+        fun canonicalJson(): ByteArray {
+            val sb = StringBuilder("{\"dh\":\"").append(Bytes.b64(dh))
+                .append("\",\"pn\":").append(pn).append(",\"n\":").append(n)
+            // 🔴 Omitted ENTIRELY when absent, never emitted as null (§13.3).
+            // The AAD is reproduced byte-for-byte by three implementations; a
+            // header that spelled absence differently would fail every open.
+            if (pqdh != null && pqct != null) {
+                sb.append(",\"pqdh\":\"").append(Bytes.b64(pqdh))
+                  .append("\",\"pqct\":\"").append(Bytes.b64(pqct)).append("\"")
+            }
+            return sb.append("}").toString().toByteArray(Charsets.UTF_8)
+        }
     }
 
     data class Message(val header: Header, val ciphertext: ByteArray)
 
     companion object {
         private val RATCHET_INFO = "FlamenetE2E_Ratchet_v1".toByteArray(Charsets.UTF_8)
+
+        /**
+         * 🔴 **The info string MUST change for v3 (spec §13.2).** A v3 IKM is
+         * `dh_out || ss` — 64 bytes where v1 is 32 — so the two cannot collide
+         * by length, but the version string is what makes the separation
+         * explicit rather than incidental.
+         */
+        private val RATCHET_INFO_V3 = "FlamenetE2E_Ratchet_v3".toByteArray(Charsets.UTF_8)
         private val MSGKEY_INFO = "FlamenetE2E_MsgKey_v1".toByteArray(Charsets.UTF_8)
         private val ZERO_SALT = ByteArray(32)
 
@@ -99,13 +137,32 @@ class DoubleRatchet private constructor(
         private const val NONCE_LEN = 12
 
         /** Alice, who has just completed X3DH and knows Bob's signed prekey. */
-        fun initiator(sharedSecret: ByteArray, remoteSignedPreKey: ByteArray): DoubleRatchet {
+        fun initiator(
+            sharedSecret: ByteArray,
+            remoteSignedPreKey: ByteArray,
+            /**
+             * The PQ prekey X3DH encapsulated to (§12.5) — the SAME key, which
+             * is what makes it a valid `PQr` to start from. Null runs v2.
+             */
+            remotePQPreKey: ByteArray? = null,
+        ): DoubleRatchet {
             val self = X25519.generate()
-            val (rk, ck) = rootStep(sharedSecret, X25519.agree(self.private, remoteSignedPreKey))
+            var pqSelf: PQKEM.KeyPair? = null
+            var pqPending: Pair<ByteArray, ByteArray>? = null
+            var ss: ByteArray? = null
+            if (remotePQPreKey != null) {
+                pqSelf = PQKEM.generate()
+                val enc = PQKEM.encapsulate(remotePQPreKey)
+                ss = enc.sharedSecret
+                pqPending = pqSelf.publicKey to enc.ciphertext
+            }
+            val (rk, ck) = rootStep(
+                sharedSecret, X25519.agree(self.private, remoteSignedPreKey), ss)
             return DoubleRatchet(
                 rootKey = rk, sendingChain = ck, receivingChain = null,
                 selfRatchet = self, remoteRatchet = remoteSignedPreKey,
                 sent = 0, received = 0, previousChainLength = 0,
+                pqSelf = pqSelf, pqRemote = remotePQPreKey, pqPending = pqPending,
             )
         }
 
@@ -115,11 +172,21 @@ class DoubleRatchet private constructor(
          * He does not step the root yet: he has nothing to step it with until a
          * message from Alice tells him her ratchet public key.
          */
-        fun responder(sharedSecret: ByteArray, signedPreKey: java.security.KeyPair): DoubleRatchet =
+        fun responder(
+            sharedSecret: ByteArray,
+            signedPreKey: java.security.KeyPair,
+            /**
+             * The seed of the PQ prekey the initiator encapsulated to, so the
+             * first inbound ratchet step can decapsulate. `PQr` is adopted from
+             * that first header (§13.2).
+             */
+            pqPreKeySeed: ByteArray? = null,
+        ): DoubleRatchet =
             DoubleRatchet(
                 rootKey = sharedSecret, sendingChain = null, receivingChain = null,
                 selfRatchet = signedPreKey, remoteRatchet = null,
                 sent = 0, received = 0, previousChainLength = 0,
+                pqSelf = pqPreKeySeed?.let(PQKEM::fromSeed),
             )
 
         /** The other half of [archive]. */
@@ -141,6 +208,13 @@ class DoubleRatchet private constructor(
                 sent = archive.getInt("ns"),
                 received = archive.getInt("nr"),
                 previousChainLength = archive.getInt("pn"),
+                // Rebuilt from the seed, so an archive written by an older
+                // build simply has none and the session runs v2.
+                pqSelf = opt("pqs")?.let(PQKEM::fromSeed),
+                pqRemote = opt("pqr"),
+                pqPending = opt("pqPendingPub")?.let { pub ->
+                    opt("pqPendingCt")?.let { ct -> pub to ct }
+                },
             )
             val skipped = archive.optJSONArray("skipped") ?: org.json.JSONArray()
             for (i in 0 until skipped.length()) {
@@ -151,8 +225,12 @@ class DoubleRatchet private constructor(
             return ratchet
         }
 
-        private fun rootStep(rootKey: ByteArray, dhOutput: ByteArray): Pair<ByteArray, ByteArray> {
-            val out = Hkdf.derive(dhOutput, rootKey, RATCHET_INFO, 64)
+        private fun rootStep(
+            rootKey: ByteArray, dhOutput: ByteArray, ss: ByteArray? = null,
+        ): Pair<ByteArray, ByteArray> {
+            val ikm = if (ss == null) dhOutput else dhOutput + ss
+            val info = if (ss == null) RATCHET_INFO else RATCHET_INFO_V3
+            val out = Hkdf.derive(ikm, rootKey, info, 64)
             return out.copyOfRange(0, 32) to out.copyOfRange(32, 64)
         }
 
@@ -184,6 +262,8 @@ class DoubleRatchet private constructor(
             dh = X25519.rawFromPublicKey(selfRatchet!!.public),
             pn = previousChainLength,
             n = sent,
+            pqdh = pqPending?.first,
+            pqct = pqPending?.second,
         )
         sent += 1
 
@@ -224,7 +304,7 @@ class DoubleRatchet private constructor(
         // then turn the DH ratchet.
         if (remoteRatchet == null || !message.header.dh.contentEquals(remoteRatchet)) {
             skipTo(message.header.pn)
-            dhRatchet(message.header.dh)
+            dhRatchet(message.header.dh, message.header.pqdh, message.header.pqct)
         }
         skipTo(message.header.n)
 
@@ -265,18 +345,40 @@ class DoubleRatchet private constructor(
      * their new key, a fresh keypair of our own, and a second root step to open
      * the sending chain they will decrypt with.
      */
-    private fun dhRatchet(theirRatchet: ByteArray) {
+    private fun dhRatchet(
+        theirRatchet: ByteArray, pqdh: ByteArray? = null, pqct: ByteArray? = null,
+    ) {
         previousChainLength = sent
         sent = 0
         received = 0
+
+        // 🔴 **Decapsulate BEFORE adopting their new PQ key, and getting this
+        // order wrong is SILENT.** The peer encapsulated to the key we held
+        // going into this step, so `pqSelf` here must still be the old one
+        // (§13.2). Reversing these derives a different root key with no error
+        // until an AEAD open fails somewhere later.
+        val ssRecv = if (pqct != null) pqSelf?.let { PQKEM.decapsulate(pqct, it.seed) } else null
         remoteRatchet = theirRatchet
 
-        var (rk, ck) = rootStep(rootKey, X25519.agree(selfRatchet!!.private, theirRatchet))
+        val (rk, ck) = rootStep(
+            rootKey, X25519.agree(selfRatchet!!.private, theirRatchet), ssRecv)
         rootKey = rk
         receivingChain = ck
 
+        // Now adopt theirs and encapsulate to it for the sending chain.
+        var ssSend: ByteArray? = null
+        if (pqdh != null) {
+            pqRemote = pqdh
+            val mine = PQKEM.generate()
+            pqSelf = mine
+            val enc = PQKEM.encapsulate(pqdh)
+            ssSend = enc.sharedSecret
+            pqPending = mine.publicKey to enc.ciphertext
+        }
+
         selfRatchet = X25519.generate()
-        val stepped = rootStep(rootKey, X25519.agree(selfRatchet!!.private, theirRatchet))
+        val stepped = rootStep(
+            rootKey, X25519.agree(selfRatchet!!.private, theirRatchet), ssSend)
         rootKey = stepped.first
         sendingChain = stepped.second
     }
@@ -312,6 +414,14 @@ class DoubleRatchet private constructor(
             .put("ns", sent)
             .put("nr", received)
             .put("pn", previousChainLength)
+            // ⚠️ A v3 session that lost these on a restart would silently drop
+            // to X25519-only healing: the ratchet keeps working, so nothing
+            // fails, and the post-quantum property is just gone.
+            .put("pqs", pqSelf?.let { Bytes.b64(it.seed) })
+            .put("pqsPub", pqSelf?.let { Bytes.b64(it.publicKey) })
+            .put("pqr", pqRemote?.let { Bytes.b64(it) })
+            .put("pqPendingPub", pqPending?.let { Bytes.b64(it.first) })
+            .put("pqPendingCt", pqPending?.let { Bytes.b64(it.second) })
             .put("skipped", skippedJson)
     }
 }

engine/src/main/kotlin/io/flamenet/e2e/E2ESession.kt

@@ -91,7 +91,11 @@ class E2ESession private constructor(
                 .put("opk_id", bundle.oneTimePreKeyId)
 
             return E2ESession(
-                ratchet = DoubleRatchet.initiator(initiated.sharedSecret, bundle.signedPreKey),
+                // 🔴 `chosen?.publicKey` is the PQ prekey X3DH just encapsulated
+                // to, and the ratchet has to start from that exact key or the
+                // responder cannot decapsulate the first step (§13.2).
+                ratchet = DoubleRatchet.initiator(
+                    initiated.sharedSecret, bundle.signedPreKey, chosen?.publicKey),
                 associatedData = X3DH.associatedData(
                     identity.identitySigningPublic, bundle.identitySigning,
                     initiated.pqCiphertext),
@@ -150,7 +154,9 @@ class E2ESession private constructor(
             )
 
             val session = E2ESession(
-                ratchet = DoubleRatchet.responder(sharedSecret, signedPreKey),
+                // The seed of the PQ prekey the initiator used, so the first
+                // inbound ratchet step can decapsulate to it.
+                ratchet = DoubleRatchet.responder(sharedSecret, signedPreKey, pq?.first),
                 associatedData = X3DH.associatedData(
                     theirIdentitySigning, identity.identitySigningPublic, pq?.second),
                 pendingHandshake = null,
@@ -373,7 +379,11 @@ data class Envelope(
         .put("header", JSONObject()
             .put("dh", Bytes.b64(header.dh))
             .put("pn", header.pn)
-            .put("n", header.n))
+            .put("n", header.n)
+            // §13.3 — present only on a PQ step. `JSONObject.put(String, null)`
+            // omits the key, which is exactly what the canonical form needs.
+            .put("pqdh", header.pqdh?.let(Bytes::b64))
+            .put("pqct", header.pqct?.let(Bytes::b64)))
         .put("x3dh", handshake)
         .put("pq", pq?.let {
             JSONObject().put("kind", it.kind).put("id", it.id)
@@ -397,6 +407,8 @@ data class Envelope(
                     dh = Bytes.unb64(header.getString("dh")),
                     pn = header.getInt("pn"),
                     n = header.getInt("n"),
+                    pqdh = header.optString("pqdh").ifEmpty { null }?.let(Bytes::unb64),
+                    pqct = header.optString("pqct").ifEmpty { null }?.let(Bytes::unb64),
                 ),
                 handshake = json.optJSONObject("x3dh"),
                 pq = json.optJSONObject("pq")?.let {

engine/src/test/kotlin/io/flamenet/e2e/DoubleRatchetTest.kt

@@ -47,6 +47,202 @@ class DoubleRatchetTest {
         return alice to bob
     }
 
+    /**
+     * The same pair, running protocol v3 — a fresh ML-KEM secret mixed into
+     * every ratchet step (§13).
+     *
+     * Bob's PQ prekey is the shared starting point: Alice encapsulates to it,
+     * and Bob holds its seed so he can decapsulate the first step.
+     */
+    private fun pqPair(): Triple<DoubleRatchet, DoubleRatchet, PQKEM.KeyPair> {
+        val bobSigned = X25519.generate()
+        val aliceIdentity = X25519.generate()
+        val bobIdentity = X25519.generate()
+        val bobPQ = PQKEM.generate()
+
+        val initiated = X3DH.initiate(
+            aliceIdentity.private,
+            X3DH.Bundle(
+                identityKeyDH = X25519.rawFromPublicKey(bobIdentity.public),
+                signedPreKey = X25519.rawFromPublicKey(bobSigned.public),
+                oneTimePreKey = null,
+            ),
+            bobPQ = bobPQ.publicKey,
+        )
+        val bobSecret = X3DH.respond(
+            identityPrivateDH = bobIdentity.private,
+            signedPreKeyPrivate = bobSigned.private,
+            oneTimePrivate = null,
+            initiatorIdentityDH = X25519.rawFromPublicKey(aliceIdentity.public),
+            initiatorEphemeral = initiated.ephemeralPublic,
+            pq = bobPQ.seed to initiated.pqCiphertext!!,
+        )
+        assertContentEquals(initiated.sharedSecret, bobSecret, "PQ X3DH disagreed")
+
+        val alice = DoubleRatchet.initiator(
+            initiated.sharedSecret, X25519.rawFromPublicKey(bobSigned.public), bobPQ.publicKey)
+        val bob = DoubleRatchet.responder(bobSecret, bobSigned, bobPQ.seed)
+        return Triple(alice, bob, bobPQ)
+    }
+
+    // MARK: - Protocol v3, the post-quantum ratchet (§13)
+
+    /**
+     * 🔴 **The cross-engine contract, pinned as bytes.**
+     *
+     * The header goes into the AAD, so three implementations must produce
+     * byte-identical JSON or nothing opens across engines. Key order is
+     * `dh, pn, n, pqdh, pqct`, there are no spaces, and absent PQ fields are
+     * omitted ENTIRELY rather than written as null (§13.3).
+     *
+     * This is the assertion that would catch a JSON library being introduced
+     * here, or a field being reordered — neither of which any conversation
+     * test can see, because both sides would drift together.
+     */
+    @Test
+    fun `the canonical header is exactly the bytes the spec names`() {
+        val dh = ByteArray(32) { it.toByte() }
+        val pqdh = ByteArray(8) { (it + 100).toByte() }
+        val pqct = ByteArray(8) { (it + 200).toByte() }
+
+        assertEquals(
+            "{\"dh\":\"" + Bytes.b64(dh) + "\",\"pn\":3,\"n\":7}",
+            String(DoubleRatchet.Header(dh, 3, 7).canonicalJson(), Charsets.UTF_8),
+            "the v2 header changed shape",
+        )
+        assertEquals(
+            "{\"dh\":\"" + Bytes.b64(dh) + "\",\"pn\":3,\"n\":7" +
+                ",\"pqdh\":\"" + Bytes.b64(pqdh) + "\",\"pqct\":\"" + Bytes.b64(pqct) + "\"}",
+            String(DoubleRatchet.Header(dh, 3, 7, pqdh, pqct).canonicalJson(), Charsets.UTF_8),
+            "the v3 header changed shape",
+        )
+        // ⚠️ Half a PQ step is not a PQ step. One field without the other must
+        // fall back to the v2 shape rather than emit a malformed header.
+        assertEquals(
+            "{\"dh\":\"" + Bytes.b64(dh) + "\",\"pn\":0,\"n\":0}",
+            String(DoubleRatchet.Header(dh, 0, 0, pqdh, null).canonicalJson(), Charsets.UTF_8),
+            "a half-populated PQ step leaked into the header",
+        )
+    }
+
+    @Test
+    fun `a v3 conversation going both ways stays in step`() {
+        val (alice, bob, _) = pqPair()
+        // Four direction changes: each one turns the DH ratchet AND the PQ
+        // ratchet in lockstep, so an ordering mistake in either compounds.
+        assertContentEquals("a1".toByteArray(), bob.decrypt(alice.encrypt("a1".toByteArray(), AD), AD))
+        assertContentEquals("b1".toByteArray(), alice.decrypt(bob.encrypt("b1".toByteArray(), AD), AD))
+        assertContentEquals("a2".toByteArray(), bob.decrypt(alice.encrypt("a2".toByteArray(), AD), AD))
+        assertContentEquals("b2".toByteArray(), alice.decrypt(bob.encrypt("b2".toByteArray(), AD), AD))
+    }
+
+    /**
+     * 🔴 **The property v3 exists for.** v2 already protects a whole session
+     * against harvest-now-decrypt-later, because the root chain never recovers
+     * once the adversary has missed the start. What v2 lacks is post-quantum
+     * *healing*: after a compromise, recovery rests entirely on X25519 steps,
+     * which a quantum adversary breaks. So every step must carry NEW ML-KEM
+     * material, not merely the handshake's.
+     */
+    @Test
+    fun `every ratchet step carries a fresh post-quantum key`() {
+        val (alice, bob, bobPQ) = pqPair()
+
+        val first = alice.encrypt("a1".toByteArray(), AD)
+        assertTrue(first.header.pqdh != null && first.header.pqct != null,
+            "the first sending chain carried no PQ step at all")
+        // Encapsulated to Bob's PREKEY, which is what he can decapsulate with.
+        val seen = mutableListOf(Bytes.b64(first.header.pqdh!!))
+        bob.decrypt(first, AD)
+
+        // Bob answers: a ratchet step, so a PQ key of his own.
+        val reply = bob.encrypt("b1".toByteArray(), AD)
+        assertTrue(reply.header.pqdh != null, "Bob's ratchet step carried no PQ key")
+        seen += Bytes.b64(reply.header.pqdh!!)
+        alice.decrypt(reply, AD)
+
+        val third = alice.encrypt("a2".toByteArray(), AD)
+        seen += Bytes.b64(third.header.pqdh!!)
+        bob.decrypt(third, AD)
+
+        assertEquals(seen.size, seen.toSet().size,
+            "a PQ key was reused across ratchet steps: $seen")
+        assertTrue(seen.none { it == Bytes.b64(bobPQ.publicKey) },
+            "a ratchet step reused the handshake prekey instead of a fresh key")
+    }
+
+    /**
+     * ⚠️ **Repeated on every message of a chain, not just the first (§13.3).**
+     * Sending it once would mean losing that one message strands the entire
+     * chain: the receiver never gets the material to advance its root key, and
+     * nothing later in the chain can ever be decrypted.
+     */
+    @Test
+    fun `the post-quantum step repeats until the peer answers`() {
+        val (alice, bob, _) = pqPair()
+        val first = alice.encrypt("a1".toByteArray(), AD)
+        val second = alice.encrypt("a2".toByteArray(), AD)
+        val third = alice.encrypt("a3".toByteArray(), AD)
+
+        for ((i, m) in listOf(first, second, third).withIndex()) {
+            assertTrue(m.header.pqct != null, "message $i of the chain dropped the PQ step")
+        }
+        assertContentEquals(first.header.pqdh, third.header.pqdh,
+            "the PQ key changed inside one sending chain")
+
+        // 🔴 And the whole chain still opens if the FIRST message is the one
+        // that goes missing — which is the loss the repetition exists for.
+        assertContentEquals("a2".toByteArray(), bob.decrypt(second, AD))
+        assertContentEquals("a3".toByteArray(), bob.decrypt(third, AD))
+    }
+
+    /**
+     * A stripped PQ step must not silently downgrade to a v1 root step. The
+     * header is bound into the AAD, so removing it fails the open — which is
+     * the difference between a downgrade attack and a broken message.
+     */
+    @Test
+    fun `a stripped post-quantum step does not open`() {
+        val (alice, bob, _) = pqPair()
+        val message = alice.encrypt("a1".toByteArray(), AD)
+        val stripped = DoubleRatchet.Message(
+            DoubleRatchet.Header(message.header.dh, message.header.pn, message.header.n),
+            message.ciphertext,
+        )
+        assertFailsWith<Exception> { bob.decrypt(stripped, AD) }
+    }
+
+    /**
+     * ⚠️ A v3 session that lost its PQ state on a restart would keep working
+     * and silently fall back to X25519-only healing — nothing fails, and the
+     * property is simply gone.
+     */
+    @Test
+    fun `a v3 session survives being archived`() {
+        val (alice, bob, _) = pqPair()
+        bob.decrypt(alice.encrypt("a1".toByteArray(), AD), AD)
+
+        val restoredBob = DoubleRatchet.fromArchive(bob.archive())
+        val restoredAlice = DoubleRatchet.fromArchive(alice.archive())
+
+        val reply = restoredBob.encrypt("b1".toByteArray(), AD)
+        assertTrue(reply.header.pqdh != null, "the restored session lost its PQ step")
+        assertContentEquals("b1".toByteArray(), restoredAlice.decrypt(reply, AD))
+    }
+
+    /**
+     * 🔴 A v3 device must still speak v2 to a peer that does not advertise it
+     * (§13.4), and the header then carries no PQ fields at all.
+     */
+    @Test
+    fun `a session without post-quantum material still runs v2`() {
+        val (alice, bob) = pair()
+        val message = alice.encrypt("hello".toByteArray(), AD)
+        assertTrue(message.header.pqdh == null && message.header.pqct == null,
+            "a v2 session emitted PQ header fields")
+        assertContentEquals("hello".toByteArray(), bob.decrypt(message, AD))
+    }
+
     @Test
     fun `a message survives the trip`() {
         val (alice, bob) = pair()