CodexRestoreAtomicityTest.kt

1package com.peachpasswords.vault23import kotlinx.coroutines.runBlocking4import kotlinx.serialization.json.JsonNull5import kotlinx.serialization.json.JsonPrimitive6import kotlinx.serialization.json.buildJsonObject7import org.junit.Assert.assertEquals8import org.junit.Assert.assertTrue9import org.junit.Test10import java.util.Base6411import java.util.concurrent.CountDownLatch12import java.util.concurrent.TimeUnit13import java.util.concurrent.atomic.AtomicReference1415/**16 * C-01: a Codex restore must not touch the device's vault until the candidate17 * has been decrypted and decoded in memory, and must roll back if any write in18 * the commit fails. An Android-shaped Codex parses structurally under ANY19 * password, so the old write-then-unlock order wiped the local vault on a typo.20 */21class CodexRestoreAtomicityTest {2223    private class FaultStorage : VaultStorage {24        var slotA: StoredVault? = null25        var slotB: String? = null26        var failSlotA = false27        var failSlotB = false28        var failNextSlotB = false29        var failedSlotBEntered: CountDownLatch? = null30        var releaseFailedSlotB: CountDownLatch? = null31        private var ver: Long = 032        private var settings = VaultSettings()33        override fun loadSlotA() = slotA34        override fun hasVault() = slotA != null35        override fun writeSlotA(slot: StoredVault): Boolean {36            if (failSlotA) return false37            slotA = slot38            return true39        }40        override fun loadSlotB() = slotB41        override fun writeSlotB(blobBase64: String): Boolean {42            if (failNextSlotB) {43                failNextSlotB = false44                failedSlotBEntered?.countDown()45                check(releaseFailedSlotB?.await(5, TimeUnit.SECONDS) != false) {46                    "test writer did not release failed slot-B write"47                }48                return false49            }50            if (failSlotB) return false51            slotB = blobBase6452            return true53        }54        var decoy: StoredVault? = null55        override fun loadDecoySlot() = decoy56        override fun writeDecoySlot(slot: StoredVault): Boolean { decoy = slot; return true }57        override fun hasDecoySlot() = decoy != null58        override fun clearDecoySlot() { decoy = null }59        override fun version() = ver60        override fun stampLocalMutation(): Long { ver = maxOf(System.currentTimeMillis(), ver + 1); return ver }61        override fun loadSettings() = settings62        override fun writeSettings(s: VaultSettings): Boolean { settings = s; return true }63        override fun clearVault(): Boolean { slotA = null; slotB = null; ver = 0; settings = VaultSettings(); return true }64    }6566    private fun sourceCodexBlob(): String = runBlocking {67        val src = VaultRepository(FaultStorage())68        src.createVault("source pass")69        src.addEntry(VaultEntry(id = "", title = "Incoming", type = "login", password = "p", createdAt = 0, modifiedAt = 0))70        src.codexVaultBlobBase64()!!71    }7273    /** Re-wrap an Android blob the way the browser writes it: slot as a JSON string. */74    private fun asBrowserShapedWrapper(androidBlob: String): String {75        val wrapper = String(Base64.getUrlDecoder().decode(androidBlob), Charsets.UTF_8)76        // H-04: codex now carries "a" and "b" as siblings. Match the first77        // balanced `{…}` after `"a":` — earlier regex matched through to the78        // wrapper's last brace and broke when the duress slot was added.79        val openChar = Regex("\"a\"\\s*:\\s*\\{").find(wrapper)!!.range.last80        val open = openChar + 181        var depth = 182        var end = open83        while (end < wrapper.length && depth > 0) {84            when (wrapper[end]) {85                '{' -> depth++86                '}' -> depth--87            }88            if (depth == 0) break89            end++90        }91        val slotA = wrapper.substring(openChar, end + 1)92        val browserWrapper = buildJsonObject {93            put("a", JsonPrimitive(slotA))94            put("b", JsonNull)95        }.toString()96        return Base64.getUrlEncoder().withoutPadding().encodeToString(browserWrapper.toByteArray(Charsets.UTF_8))97    }9899    /** An existing local vault holding one entry, plus its storage. */100    private fun existingVault(): Pair<VaultRepository, FaultStorage> = runBlocking {101        val storage = FaultStorage()102        val repo = VaultRepository(storage)103        repo.createVault("local pass")104        repo.addEntry(VaultEntry(id = "", title = "Local", type = "login", password = "p", createdAt = 0, modifiedAt = 0))105        repo.lock()106        repo to storage107    }108109    private fun assertLocalVaultIntact(repo: VaultRepository, storage: FaultStorage, before: StoredVault?) =110        runBlocking {111            assertEquals("slot A must be untouched", before, storage.slotA)112            assertEquals(UnlockResult.Success, repo.unlock("local pass"))113            assertTrue(repo.getEntries().any { it.title == "Local" })114        }115116    @Test fun wrong_password_on_android_shaped_codex_preserves_local_vault() {117        val blob = sourceCodexBlob()118        val (repo, storage) = existingVault()119        val before = storage.slotA120        runBlocking { assertEquals(UnlockResult.WrongPassword, repo.restoreFromCodex(blob, "not the pass")) }121        assertLocalVaultIntact(repo, storage, before)122    }123124    @Test fun wrong_password_on_browser_shaped_codex_preserves_local_vault() {125        val blob = asBrowserShapedWrapper(sourceCodexBlob())126        val (repo, storage) = existingVault()127        val before = storage.slotA128        runBlocking { assertEquals(UnlockResult.WrongPassword, repo.restoreFromCodex(blob, "not the pass")) }129        assertLocalVaultIntact(repo, storage, before)130    }131132    @Test fun slot_a_write_failure_preserves_local_vault() {133        val blob = sourceCodexBlob()134        val (repo, storage) = existingVault()135        val before = storage.slotA136        storage.failSlotA = true137        runBlocking {138            assertEquals(UnlockResult.Error("storage_write_failed"), repo.restoreFromCodex(blob, "source pass"))139        }140        storage.failSlotA = false141        assertLocalVaultIntact(repo, storage, before)142    }143144    @Test fun later_write_failure_rolls_back_to_local_vault() {145        val blob = sourceCodexBlob()146        val (repo, storage) = existingVault()147        val before = storage.slotA148        storage.failSlotB = true149        runBlocking {150            assertEquals(UnlockResult.Error("storage_write_failed"), repo.restoreFromCodex(blob, "source pass"))151        }152        storage.failSlotB = false153        assertLocalVaultIntact(repo, storage, before)154    }155156    @Test157    fun partial_codex_write_holds_epoch_fence_through_rollback() = runBlocking {158        val blob = sourceCodexBlob()159        val (repo, storage) = existingVault()160        val priorA = storage.slotA161        val priorB = storage.slotB162        val priorD = storage.decoy163        val failedSlotBEntered = CountDownLatch(1)164        val releaseFailedSlotB = CountDownLatch(1)165        val rollbackComplete = CountDownLatch(1)166        val releaseAfterRollback = CountDownLatch(1)167        storage.failNextSlotB = true168        storage.failedSlotBEntered = failedSlotBEntered169        storage.releaseFailedSlotB = releaseFailedSlotB170        repo.afterCodexRollback = {171            rollbackComplete.countDown()172            check(releaseAfterRollback.await(5, TimeUnit.SECONDS)) {173                "test did not release post-rollback fence"174            }175        }176177        val restoreResult = AtomicReference<UnlockResult>()178        val restoreDone = CountDownLatch(1)179        val restoreThread = Thread {180            restoreResult.set(runBlocking { repo.restoreFromCodex(blob, "source pass") })181            restoreDone.countDown()182        }183        restoreThread.start()184        assertTrue(185            "Codex transaction did not reach the partial slot-B write",186            failedSlotBEntered.await(5, TimeUnit.SECONDS),187        )188189        // Start lock while the epoch-held Codex transaction is still blocked in190        // its failed slot-B write. The lock must wait until A/B/D rollback has191        // completed, otherwise it can advance the epoch and strand incoming A.192        val lockStarted = CountDownLatch(1)193        val lockDone = CountDownLatch(1)194        val lockThread = Thread {195            lockStarted.countDown()196            repo.lock()197            lockDone.countDown()198        }199        lockThread.start()200        assertTrue("lock thread did not start", lockStarted.await(5, TimeUnit.SECONDS))201202        releaseFailedSlotB.countDown()203        assertTrue("Codex rollback did not finish", rollbackComplete.await(5, TimeUnit.SECONDS))204        assertEquals("rollback must restore prior slot A before lock completes", priorA, storage.slotA)205        assertEquals("rollback must restore prior noise slot B before lock completes", priorB, storage.slotB)206        assertEquals("rollback must restore prior duress slot before lock completes", priorD, storage.decoy)207        assertTrue(208            "lock must remain blocked until the epoch-held rollback returns",209            !lockDone.await(0, TimeUnit.NANOSECONDS),210        )211        releaseAfterRollback.countDown()212        assertTrue("Codex restore did not finish", restoreDone.await(5, TimeUnit.SECONDS))213        assertTrue("lock did not finish", lockDone.await(5, TimeUnit.SECONDS))214        restoreThread.join(5_000)215        lockThread.join(5_000)216217        assertEquals(UnlockResult.Error("storage_write_failed"), restoreResult.get())218        assertEquals("rollback must restore prior slot A", priorA, storage.slotA)219        assertEquals("rollback must restore prior noise slot B", priorB, storage.slotB)220        assertEquals("rollback must restore prior duress slot", priorD, storage.decoy)221    }222}