flamenet-messenger-android/engine/src/main/kotlin/io/flamenet/e2e/DoubleRatchet.kt

package io.flamenet.e2e

import org.json.JSONObject
import javax.crypto.Mac
import javax.crypto.spec.SecretKeySpec

/**
 * The Double Ratchet (spec §6), with unencrypted headers.
 *
 * Two ratchets turning at different rates. The **symmetric** one advances once
 * per message, so a compromised message key reveals nothing about earlier ones.
 * The **Diffie-Hellman** one advances whenever the other side sends a new
 * ratchet public key, which is what lets a conversation recover after a
 * compromise rather than staying broken forever.
 *
 * ```
 * MK  = HMAC-SHA256(key = CK, msg = 0x01)
 * CK' = HMAC-SHA256(key = CK, msg = 0x02)
 * (RK', CK') = HKDF-SHA256(IKM = dh_out, salt = RK,
 *                          info = "FlamenetE2E_Ratchet_v1", L = 64)
 * ```
 *
 * ⚠️ **Everything here is state, and state is where this goes wrong.** X3DH is
 * a pure function and either agrees or does not. A ratchet can agree for six
 * messages and then diverge on the seventh because one side counted a skipped
 * message differently — and the symptom, again, is only that something will not
 * decrypt. The tests below therefore run conversations, out of order and with
 * gaps, rather than checking any single step.
 *
 * ⚠️ **Never encrypt twice under one message key.** Two AES-GCM encryptions
 * under the same key and nonce leak the XOR of both plaintexts and destroy the
 * authentication guarantee — it is a catastrophic failure, not a degraded one.
 * The sending chain advances before the key is used, so a key cannot be reused
 * even if a caller retries.
 *
 * A port of `engine/src/e2e.js`, which is normative.
 */
class DoubleRatchet private constructor(
    private var rootKey: ByteArray,
    private var sendingChain: ByteArray?,
    private var receivingChain: ByteArray?,
    private var selfRatchet: java.security.KeyPair?,
    private var remoteRatchet: ByteArray?,
    /** Messages sent in the current sending chain. */
    private var sent: Int,
    /** Messages received in the current receiving chain. */
    private var received: Int,
    /** Length of the PREVIOUS sending chain, so the far side knows how many to skip. */
    private var previousChainLength: Int,
) {

    /** Keys for messages that arrived out of order, or have not arrived yet. */
    private val skipped = LinkedHashMap<SkippedKey, ByteArray>()

    data class SkippedKey(val ratchetPublic: String, val index: Int)

    data class Header(val dh: ByteArray, val pn: Int, val n: Int) {
        /**
         * The exact bytes that go into the AAD.
         *
         * Hand-built rather than serialised by a JSON library: the AAD has to be
         * reproduced byte-for-byte by three implementations, and an encoder is
         * 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)
    }

    data class Message(val header: Header, val ciphertext: ByteArray)

    companion object {
        private val RATCHET_INFO = "FlamenetE2E_Ratchet_v1".toByteArray(Charsets.UTF_8)
        private val MSGKEY_INFO = "FlamenetE2E_MsgKey_v1".toByteArray(Charsets.UTF_8)
        private val ZERO_SALT = ByteArray(32)

        /**
         * A cap on how many keys a single header may make us derive.
         *
         * Without one, a header claiming `n = 2_000_000_000` is a denial of
         * service that costs the sender nothing: the receiver dutifully derives
         * two billion keys. The number is generous for real reordering and
         * ruinous for that.
         */
        const val MAX_SKIP = 1000

        /**
         * How many skipped keys are kept across the whole session.
         *
         * `MAX_SKIP` bounds one header; without a second bound, a peer sending
         * a thousand-message gap over and over grows this map forever. Oldest
         * out first — a key that has waited through two thousand later messages
         * is not arriving.
         */
        const val MAX_STORE = 2000

        /** Matches `Aead`'s GCM nonce length; the ratchet strips it on the wire. */
        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 {
            val self = X25519.generate()
            val (rk, ck) = rootStep(sharedSecret, X25519.agree(self.private, remoteSignedPreKey))
            return DoubleRatchet(
                rootKey = rk, sendingChain = ck, receivingChain = null,
                selfRatchet = self, remoteRatchet = remoteSignedPreKey,
                sent = 0, received = 0, previousChainLength = 0,
            )
        }

        /**
         * Bob, whose signed prekey IS his first ratchet key.
         *
         * 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 =
            DoubleRatchet(
                rootKey = sharedSecret, sendingChain = null, receivingChain = null,
                selfRatchet = signedPreKey, remoteRatchet = null,
                sent = 0, received = 0, previousChainLength = 0,
            )

        /** The other half of [archive]. */
        fun fromArchive(archive: JSONObject): DoubleRatchet {
            fun opt(key: String) = archive.optString(key, "").takeIf { it.isNotEmpty() }?.let(Bytes::unb64)
            val selfPrivate = opt("dhs")
            val selfPublic = opt("dhsPub")
            val ratchet = DoubleRatchet(
                rootKey = Bytes.unb64(archive.getString("rk")),
                sendingChain = opt("cks"),
                receivingChain = opt("ckr"),
                selfRatchet = if (selfPrivate != null && selfPublic != null) {
                    java.security.KeyPair(
                        X25519.publicKeyFromRaw(selfPublic),
                        X25519.privateKeyFromPkcs8(selfPrivate),
                    )
                } else null,
                remoteRatchet = opt("dhr"),
                sent = archive.getInt("ns"),
                received = archive.getInt("nr"),
                previousChainLength = archive.getInt("pn"),
            )
            val skipped = archive.optJSONArray("skipped") ?: org.json.JSONArray()
            for (i in 0 until skipped.length()) {
                val entry = skipped.getJSONObject(i)
                ratchet.skipped[SkippedKey(entry.getString("dh"), entry.getInt("n"))] =
                    Bytes.unb64(entry.getString("mk"))
            }
            return ratchet
        }

        private fun rootStep(rootKey: ByteArray, dhOutput: ByteArray): Pair<ByteArray, ByteArray> {
            val out = Hkdf.derive(dhOutput, rootKey, RATCHET_INFO, 64)
            return out.copyOfRange(0, 32) to out.copyOfRange(32, 64)
        }

        private fun hmac(key: ByteArray, msg: ByteArray): ByteArray {
            val mac = Mac.getInstance("HmacSHA256")
            mac.init(SecretKeySpec(key, "HmacSHA256"))
            return mac.doFinal(msg)
        }

        /** One symmetric step: a message key, and the chain key that replaces it. */
        private fun chainStep(chainKey: ByteArray): Pair<ByteArray, ByteArray> =
            hmac(chainKey, byteArrayOf(0x01)) to hmac(chainKey, byteArrayOf(0x02))

        /** A message key expands into an AES key and the nonce it is used with. */
        private fun aeadParams(messageKey: ByteArray): Pair<ByteArray, ByteArray> {
            val okm = Hkdf.derive(messageKey, ZERO_SALT, MSGKEY_INFO, 44)
            return okm.copyOfRange(0, 32) to okm.copyOfRange(32, 44)
        }
    }

    fun encrypt(plaintext: ByteArray, associatedData: ByteArray): Message {
        val chain = sendingChain ?: error("cannot send before the first ratchet step")
        val (messageKey, nextChain) = chainStep(chain)
        // Advanced BEFORE use, so a caller that retries cannot encrypt twice
        // under one key.
        sendingChain = nextChain

        val header = Header(
            dh = X25519.rawFromPublicKey(selfRatchet!!.public),
            pn = previousChainLength,
            n = sent,
        )
        sent += 1

        val (aesKey, nonce) = aeadParams(messageKey)
        // `Aead.seal` returns `nonce || ct || tag`, and the nonce is dropped
        // here: the ratchet DERIVES it from the message key, so transmitting it
        // would put twelve redundant bytes on every message and invite a
        // receiver to trust a nonce off the wire instead of the one the key
        // says it must be.
        val sealed = Aead.seal(aesKey, nonce, plaintext, associatedData + header.canonicalJson())
        return Message(header, sealed.copyOfRange(NONCE_LEN, sealed.size))
    }

    fun decrypt(message: Message, associatedData: ByteArray): ByteArray {
        val aad = associatedData + message.header.canonicalJson()

        // A key held for a message that never arrived, or arrived late.
        //
        // 🔴 **Removed only AFTER the tag verifies (audit LOW).** Consuming it
        // first meant a forged ciphertext for a stashed (dh, n) burned the key
        // the real delayed message needed — the genuine message could then
        // never be opened. Today the caller discards the whole session object
        // on a failed decrypt, so the loss never reaches the vault; that is a
        // property of the CALLER, and the moment anything caches a ratchet in
        // memory it would become a real message-loss bug. iOS has always
        // removed after the open for exactly this reason.
        val skippedKey = SkippedKey(Bytes.b64(message.header.dh), message.header.n)
        skipped[skippedKey]?.let { mk ->
            val (aesKey, nonce) = aeadParams(mk)
            val plaintext = Aead.open(aesKey, nonce + message.ciphertext, aad)
            // LinkedHashMap keeps insertion order, so removing from it is the
            // whole eviction bookkeeping — there is no separate order list.
            skipped.remove(skippedKey)
            return plaintext
        }

        // A new ratchet key from the far side: catch up on their old chain,
        // then turn the DH ratchet.
        if (remoteRatchet == null || !message.header.dh.contentEquals(remoteRatchet)) {
            skipTo(message.header.pn)
            dhRatchet(message.header.dh)
        }
        skipTo(message.header.n)

        val chain = receivingChain ?: error("no receiving chain")
        val (messageKey, nextChain) = chainStep(chain)
        receivingChain = nextChain
        received += 1

        val (aesKey, nonce) = aeadParams(messageKey)
        return Aead.open(aesKey, nonce + message.ciphertext, aad)
    }

    /** Derive and store keys for messages we have not seen yet. */
    private fun skipTo(target: Int) {
        val chain = receivingChain ?: return
        if (target < received) return
        require(target - received <= MAX_SKIP) {
            "refusing to derive ${target - received} skipped keys"
        }
        var ck = chain
        var n = received
        val remote = Bytes.b64(remoteRatchet ?: return)
        while (n < target) {
            val (mk, next) = chainStep(ck)
            skipped[SkippedKey(remote, n)] = mk
            while (skipped.size > MAX_STORE) {
                skipped.remove(skipped.keys.first())
            }
            ck = next
            n += 1
        }
        receivingChain = ck
        received = n
    }

    /**
     * Turn the DH ratchet: one root step to open a receiving chain against
     * 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) {
        previousChainLength = sent
        sent = 0
        received = 0
        remoteRatchet = theirRatchet

        var (rk, ck) = rootStep(rootKey, X25519.agree(selfRatchet!!.private, theirRatchet))
        rootKey = rk
        receivingChain = ck

        selfRatchet = X25519.generate()
        val stepped = rootStep(rootKey, X25519.agree(selfRatchet!!.private, theirRatchet))
        rootKey = stepped.first
        sendingChain = stepped.second
    }

    /**
     * Everything the conversation needs to continue after a restart.
     *
     * ⚠️ A ratchet that is not persisted is not a bug you notice — it is a bug
     * your *contacts* notice. A client that loses session state falls back to
     * minting a new device identity, and everyone who has ever talked to that
     * account sees "security keys changed", which is the warning that is
     * supposed to mean an attack.
     *
     * Skipped message keys travel too. They are the keys for messages that have
     * not arrived yet, and dropping them at a restart loses exactly the
     * messages that were already in flight.
     */
    fun archive(): JSONObject {
        val skippedJson = org.json.JSONArray()
        skipped.forEach { (key, messageKey) ->
            skippedJson.put(JSONObject()
                .put("dh", key.ratchetPublic)
                .put("n", key.index)
                .put("mk", Bytes.b64(messageKey)))
        }
        return JSONObject()
            .put("rk", Bytes.b64(rootKey))
            .put("cks", sendingChain?.let { Bytes.b64(it) })
            .put("ckr", receivingChain?.let { Bytes.b64(it) })
            .put("dhs", selfRatchet?.let { Bytes.b64(it.private.encoded) })
            .put("dhsPub", selfRatchet?.let { Bytes.b64(X25519.rawFromPublicKey(it.public)) })
            .put("dhr", remoteRatchet?.let { Bytes.b64(it) })
            .put("ns", sent)
            .put("nr", received)
            .put("pn", previousChainLength)
            .put("skipped", skippedJson)
    }
}