flamenet-messenger-android/engine/src/test/kotlin/io/flamenet/e2e/VaultTest.kt

package io.flamenet.e2e

import org.json.JSONObject
import javax.crypto.KeyGenerator
import javax.crypto.SecretKey
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertFailsWith
import kotlin.test.assertFalse
import kotlin.test.assertNotEquals
import kotlin.test.assertNull
import kotlin.test.assertTrue

class VaultTest {

    /** Stands in for a directory on a phone, and lets a test see the raw bytes. */
    private class MapStorage : Vault.Storage {
        val records = LinkedHashMap<String, ByteArray>()
        override fun read(name: String) = records[name]
        override fun write(name: String, bytes: ByteArray) { records[name] = bytes }
        override fun delete(name: String) { records.remove(name) }
        override fun names() = records.keys.toList()
    }

    private fun key(): SecretKey =
        KeyGenerator.getInstance("AES").apply { init(256) }.generateKey()

    private fun bundleFor(identity: DeviceIdentity): PeerBundle {
        val oneTime = identity.generateOneTimePreKeys(1).first()
        return PeerBundle.fromWire(identity.registrationWire().put("opk",
            JSONObject().put("id", oneTime.id).put("pub", Bytes.b64(oneTime.publicKey))))
    }

    @Test
    fun `an identity survives being stored and read back`() {
        val storage = MapStorage()
        val secret = key()
        val identity = DeviceIdentity.create()
        identity.generateOneTimePreKeys(2)

        Vault(storage, secret).saveIdentity(identity)
        val restored = Vault(storage, secret).loadIdentity()!!

        assertEquals(identity.deviceId, restored.deviceId)
        assertEquals(2, restored.oneTimePreKeyCount)

        // The private halves came back too, which only a real handshake proves.
        val peer = DeviceIdentity.create()
        val session = E2ESession.initiate(peer, bundleFor(restored))
        val envelope = session.encrypt("after a cold start")
        assertEquals("after a cold start", E2ESession.respond(restored, envelope).decrypt(envelope))
    }

    @Test
    fun `a conversation survives being stored and read back`() {
        val storage = MapStorage()
        val secret = key()
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()

        val a = E2ESession.initiate(alice, bundleFor(bob))
        val first = a.encrypt("hello")
        val b = E2ESession.respond(bob, first)
        b.decrypt(first)

        Vault(storage, secret).saveSession(alice.deviceId, b)
        val restored = Vault(storage, secret).loadSession(alice.deviceId)!!
        assertEquals("still talking", restored.decrypt(a.encrypt("still talking")))
        assertEquals(listOf(alice.deviceId), Vault(storage, secret).knownPeers())
    }

    /**
     * 🔴 **The chain a re-establishment replaces must survive it.**
     *
     * Both sides open the chat at once, so each runs X3DH before seeing the
     * other's prekey and each adopts the other's chain to read what arrived.
     * From then on the peer sends on the chain being replaced. A client that
     * stored only the newest session could not open any of it — silently,
     * because an envelope that will not decrypt is skipped, not reported.
     *
     * This is the storage half of the fix. `E2ESessionTest` covers the engine
     * primitives, and passes with or without it, because a test holding both
     * sessions in local variables never exercises what a client keeps.
     */
    @Test
    fun `a retired session is kept and still opens its chain`() {
        val storage = MapStorage()
        val secret = key()
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()

        // Bob starts his own chain before hearing from Alice.
        val bobsOwn = E2ESession.initiate(bob, bundleFor(alice))
        val bobsFirst = bobsOwn.encrypt("from bob")
        Vault(storage, secret).saveSession(bob.deviceId, bobsOwn)

        // Alice's message arrives on a different chain, so Alice's is adopted
        // and Bob's own is retired rather than dropped.
        val alices = E2ESession.initiate(alice, bundleFor(bob))
        val alicesFirst = alices.encrypt("from alice")
        val adopted = E2ESession.respond(bob, alicesFirst)
        assertEquals("from alice", adopted.decrypt(alicesFirst))

        val vault = Vault(storage, secret)
        vault.saveRetiredSession(bob.deviceId, vault.loadSession(bob.deviceId)!!)
        vault.saveSession(bob.deviceId, adopted)

        // Alice is still on the chain Bob started — she answers on it.
        val onTheRetiredChain = E2ESession.respond(alice, bobsFirst)
        assertEquals("from bob", onTheRetiredChain.decrypt(bobsFirst))
        val reply = onTheRetiredChain.encrypt("bob, this is the old chain")

        val fresh = Vault(storage, secret)
        assertNull(runCatching { fresh.loadSession(bob.deviceId)!!.decrypt(reply) }.getOrNull(),
            "the adopted session must NOT open the retired chain — otherwise this " +
                "test would pass without the retired copy existing at all")
        assertEquals("bob, this is the old chain",
            fresh.loadRetiredSession(bob.deviceId)!!.decrypt(reply))
    }

    /**
     * 🔴 **A failed decrypt damages the session object it was handed, so the
     * thing that gets retired must be re-read from the vault.**
     *
     * This is the bug the first version of the retirement fix shipped with, and
     * the reason `a retired session is kept and still opens its chain` did not
     * catch it: that test retires a PRISTINE session, which is not what the
     * receive path has in its hand. The real path has just proved the peer is on
     * another chain — by handing this session an envelope it could not open —
     * and `E2ESession` is a class, so that attempt consumed skipped keys and
     * stepped the chain. Retiring THAT instance stores something that can no
     * longer open anything, which looks exactly like not retiring at all.
     *
     * Asserts both halves: the mutated instance fails, the re-read one works.
     */
    @Test
    fun `the session retired after a failed decrypt must be re-read, not reused`() {
        val storage = MapStorage()
        val secret = key()
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()

        // Bob's own chain, stored, with a message Alice will later answer on it.
        val bobsOwn = E2ESession.initiate(bob, bundleFor(alice))
        val bobsFirst = bobsOwn.encrypt("from bob")
        Vault(storage, secret).saveSession(bob.deviceId, bobsOwn)

        val alicesReplyOnBobsChain = E2ESession.respond(alice, bobsFirst).let {
            it.decrypt(bobsFirst)
            it.encrypt("answering on the chain bob started")
        }

        // The receive path: Bob's stored session is handed a foreign envelope
        // and fails. THIS is the instance the buggy version retired.
        val foreign = E2ESession.initiate(alice, bundleFor(bob)).encrypt("different chain")
        val damaged = Vault(storage, secret).loadSession(bob.deviceId)!!
        assertNull(runCatching { damaged.decrypt(foreign) }.getOrNull(),
            "the foreign envelope must NOT open, or this test proves nothing")

        // Retiring the damaged instance loses the chain.
        val wrong = MapStorage().also { storage.records.forEach(it.records::put) }
        Vault(wrong, secret).saveRetiredSession(bob.deviceId, damaged)
        assertNull(
            runCatching {
                Vault(wrong, secret).loadRetiredSession(bob.deviceId)!!
                    .decrypt(alicesReplyOnBobsChain)
            }.getOrNull(),
            "a session damaged by a failed decrypt should NOT still open its chain — " +
                "if this ever starts passing, the hazard is gone and this test can go")

        // Re-reading the vault copy keeps it. This is the fix.
        val pristine = Vault(storage, secret).loadSession(bob.deviceId)!!
        Vault(storage, secret).saveRetiredSession(bob.deviceId, pristine)
        assertEquals("answering on the chain bob started",
            Vault(storage, secret).loadRetiredSession(bob.deviceId)!!
                .decrypt(alicesReplyOnBobsChain))
    }

    /** Signing out must take the retired chains too — they hold chain keys. */
    @Test
    fun `forgetting sessions forgets the retired ones as well`() {
        val storage = MapStorage()
        val secret = key()
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()

        val a = E2ESession.initiate(alice, bundleFor(bob))
        Vault(storage, secret).saveSession(bob.deviceId, a)
        Vault(storage, secret).saveRetiredSession(bob.deviceId, a)

        Vault(storage, secret).forgetAllSessions()

        assertNull(Vault(storage, secret).loadSession(bob.deviceId))
        assertNull(Vault(storage, secret).loadRetiredSession(bob.deviceId))
        assertTrue(storage.records.keys.none { it.contains("session") },
            "a record still names a session: ${storage.records.keys}")
    }

    /** What is on disk must not be readable, or none of this was worth doing. */
    @Test
    fun `nothing recognisable is written in the clear`() {
        val storage = MapStorage()
        val identity = DeviceIdentity.create()
        Vault(storage, key()).saveIdentity(identity)

        val onDisk = String(storage.records.values.first(), Charsets.ISO_8859_1)
        assertFalse(onDisk.contains(identity.deviceId))
        assertFalse(onDisk.contains("ikDH"))
        assertFalse(onDisk.contains(Bytes.b64(identity.identityDHPublic)))
    }

    /**
     * ⚠️ The failure this catches: a stored blob moved from one slot to
     * another. Without the name in the AAD, a peer's session pasted into a
     * different peer's slot decrypts perfectly and the app attributes their
     * messages to the wrong person.
     */
    @Test
    fun `a record moved to another slot does not open`() {
        val storage = MapStorage()
        val secret = key()
        val vault = Vault(storage, secret)
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()
        val session = E2ESession.initiate(alice, bundleFor(bob))
        vault.saveSession("peer-a", session)

        val blob = storage.records.getValue("session/peer-a")
        storage.records["session/peer-b"] = blob

        assertFailsWith<Vault.Corrupt> { vault.loadSession("peer-b") }
        // And the record in its own slot is untouched by the attempt.
        assertTrue(vault.loadSession("peer-a") != null)
    }

    @Test
    fun `a different key does not open the vault`() {
        val storage = MapStorage()
        Vault(storage, key()).saveIdentity(DeviceIdentity.create())
        assertFailsWith<Vault.Corrupt> { Vault(storage, key()).loadIdentity() }
    }

    @Test
    fun `a flipped bit is refused rather than half-read`() {
        val storage = MapStorage()
        val secret = key()
        Vault(storage, secret).saveIdentity(DeviceIdentity.create())
        val blob = storage.records.getValue("identity")
        blob[blob.size / 2] = (blob[blob.size / 2].toInt() xor 1).toByte()
        assertFailsWith<Vault.Corrupt> { Vault(storage, secret).loadIdentity() }
    }

    /**
     * ⚠️ A corrupt identity must NOT read as "no identity". The caller treats
     * null as a new install and mints a fresh device, which every contact sees
     * as "security keys changed" — the warning that is supposed to mean an
     * attack. Absent and unreadable are different answers.
     */
    @Test
    fun `an unreadable identity is not mistaken for a missing one`() {
        val storage = MapStorage()
        val secret = key()
        assertNull(Vault(storage, secret).loadIdentity())
        assertFalse(Vault(storage, secret).hasIdentity())

        Vault(storage, secret).saveIdentity(DeviceIdentity.create())
        storage.records["identity"] = ByteArray(64)
        assertTrue(Vault(storage, secret).hasIdentity())
        assertFailsWith<Vault.Corrupt> { Vault(storage, secret).loadIdentity() }
    }

    @Test
    fun `the inbox watermark round-trips and starts at zero`() {
        val storage = MapStorage()
        val secret = key()
        assertEquals(0, Vault(storage, secret).inboxWatermark)
        Vault(storage, secret).inboxWatermark = 4207
        assertEquals(4207, Vault(storage, secret).inboxWatermark)
    }

    @Test
    fun `forgetting removes rather than marks`() {
        val storage = MapStorage()
        val secret = key()
        val vault = Vault(storage, secret)
        val alice = DeviceIdentity.create()
        vault.saveIdentity(alice)
        vault.saveSession("peer-a", E2ESession.initiate(alice, bundleFor(DeviceIdentity.create())))
        vault.saveSession("peer-b", E2ESession.initiate(alice, bundleFor(DeviceIdentity.create())))

        vault.forget("peer-a")
        assertFalse(storage.records.keys.any { it.contains("peer-a") })
        assertEquals(listOf("peer-b"), vault.knownPeers())

        vault.forgetAllSessions()
        assertEquals(emptyList(), vault.knownPeers())
        assertTrue(vault.hasIdentity(), "signing out of a chat is not signing out of the app")

        // Sign-out takes the identity too: the next person to use this phone
        // must not inherit a device the relay still has messages for.
        vault.forgetEverything()
        assertFalse(vault.hasIdentity())
        assertEquals(emptyList(), storage.names())
    }

    @Test
    fun `two writes of the same value do not produce the same bytes`() {
        // A fresh nonce per write. Repeating one under the same key is the
        // catastrophic AES-GCM failure, not a degraded one.
        val storage = MapStorage()
        val secret = key()
        val identity = DeviceIdentity.create()
        Vault(storage, secret).saveIdentity(identity)
        val first = storage.records.getValue("identity").copyOf()
        Vault(storage, secret).saveIdentity(identity)
        assertNotEquals(Bytes.b64(first), Bytes.b64(storage.records.getValue("identity")))
    }
}