CodexCodec.kt

11.5 KB · raw

1package com.peachpasswords.codex23import kotlinx.serialization.json.Json4import java.security.MessageDigest5import java.util.Base646import java.util.zip.Deflater7import java.util.zip.Inflater89/**10 * Pure (Android-free) Codex byte pipeline: compress → stripe → Reed-Solomon →11 * envelopes, and the inverse for restore. Kept independent of Aztec/PDF so the12 * full round-trip is unit-testable on the JVM. Port of the browser13 * `generate.ts` / `restore.ts` / `erasure.ts` byte handling.14 */15object CodexCodec {16    val json = Json { encodeDefaults = true; ignoreUnknownKeys = true }1718    private val urlEncoder = Base64.getUrlEncoder().withoutPadding()19    private val urlDecoder = Base64.getUrlDecoder()2021    /**22     * F4: decompression bomb guard — mirrors PeachPayloadCodec's23     * MAX_SECTION_RAW_BYTES (16 MiB). A codex payload is the same vault JSON24     * the PeachScript sections carry, so the cap is identical by definition.25     */26    const val MAX_INFLATED_BYTES = 16L * 1024 * 10242728    fun base64Url(bytes: ByteArray): String = urlEncoder.encodeToString(bytes)29    fun fromBase64Url(s: String): ByteArray = urlDecoder.decode(s)3031    /** Uppercase hex of SHA-256(bytes) — per-chunk hash + vault fingerprint source. */32    fun sha256Hex(bytes: ByteArray): String {33        val digest = MessageDigest.getInstance("SHA-256").digest(bytes)34        val sb = StringBuilder(digest.size * 2)35        for (b in digest) sb.append("%02X".format(b.toInt() and 0xff))36        return sb.toString()37    }3839    fun deflate(data: ByteArray): ByteArray {40        val deflater = Deflater(Deflater.BEST_COMPRESSION).apply { setInput(data); finish() }41        val out = java.io.ByteArrayOutputStream(data.size / 2 + 32)42        val buf = ByteArray(8192)43        while (!deflater.finished()) out.write(buf, 0, deflater.deflate(buf))44        deflater.end()45        return out.toByteArray()46    }4748    fun inflate(data: ByteArray): ByteArray {49        val inflater = Inflater().apply { setInput(data) }50        val out = java.io.ByteArrayOutputStream(data.size * 3 + 32)51        val buf = ByteArray(8192)52        // PB-26-004: a malformed stream can return 0 bytes without finishing53        // or needing input, spinning the loop indefinitely. Treat every54        // zero-progress state as an explicit failure: needsInput means trailing55        // zero-pad; needsDictionary means a preset dictionary stream Codex56        // never produces; any other 0-return is malformed.57        while (!inflater.finished()) {58            val n = inflater.inflate(buf)59            if (n == 0) {60                if (inflater.finished()) break61                inflater.end()62                if (inflater.needsDictionary()) {63                    throw CodexException(64                        CodexErrorCode.CORRUPTED,65                        "Inflate requires a preset dictionary — Codex never produces one",66                    )67                }68                if (inflater.needsInput()) {69                    throw CodexException(70                        CodexErrorCode.CORRUPTED,71                        "Inflate input ended before the stream finished",72                    )73                }74                throw CodexException(75                    CodexErrorCode.CORRUPTED,76                    "Inflate made no progress without needing input",77                )78            }79            out.write(buf, 0, n)80            if (out.size() > MAX_INFLATED_BYTES) {81                inflater.end()82                throw CodexException(83                    CodexErrorCode.PAYLOAD_TOO_LARGE,84                    "Inflated payload exceeds the ${MAX_INFLATED_BYTES / (1024 * 1024)} MiB cap",85                )86            }87        }88        inflater.end()89        return out.toByteArray()90    }9192    // ── Stripe layout ─────────────────────────────────────────────────────9394    fun planStripeShardCounts(totalDataShards: Int): IntArray {95        if (totalDataShards <= 0) return IntArray(0)96        val numStripes = maxOf(1, Math.ceil(totalDataShards.toDouble() / MAX_DATA_SHARDS_PER_STRIPE).toInt())97        val base = totalDataShards / numStripes98        val remainder = totalDataShards % numStripes99        return IntArray(numStripes) { base + if (it < remainder) 1 else 0 }100    }101102    data class Stripe(val dataShards: Int, val parityShards: Int, val shards: Array<ByteArray>) {103        val totalChunks: Int get() = dataShards + parityShards104    }105106    /** Compress + stripe + RS-encode into envelopes for [verifyCode]. */107    fun encodeToEnvelopes(compressed: ByteArray, verifyCode: String): Pair<List<ChunkEnvelope>, CodexLayout> {108        val totalDataShards = maxOf(1, Math.ceil(compressed.size.toDouble() / CHUNK_SIZE).toInt())109        val counts = planStripeShardCounts(totalDataShards)110        val envelopes = ArrayList<ChunkEnvelope>()111        var offset = 0112        var totalData = 0113        var totalParity = 0114        for ((sIdx, dataShards) in counts.withIndex()) {115            val region = dataShards * CHUNK_SIZE116            val parityShards = computeParityShards(dataShards)117            val total = dataShards + parityShards118            val shards = Array(total) { ByteArray(CHUNK_SIZE) }119            for (d in 0 until dataShards) {120                val start = offset + d * CHUNK_SIZE121                val len = (compressed.size - start).coerceIn(0, CHUNK_SIZE)122                if (len > 0) System.arraycopy(compressed, start, shards[d], 0, len)123            }124            ReedSolomon.encode(shards, dataShards, parityShards)125            for (i in 0 until total) {126                envelopes.add(buildEnvelope(shards[i], i, total, i >= dataShards, verifyCode, sIdx, counts.size))127            }128            offset += region129            totalData += dataShards130            totalParity += parityShards131        }132        return envelopes to CodexLayout(counts.size, totalData, totalParity, totalData + totalParity)133    }134135    private fun buildEnvelope(136        chunk: ByteArray, index: Int, total: Int, isParity: Boolean,137        verifyCode: String, stripe: Int, stripeCount: Int,138    ): ChunkEnvelope = ChunkEnvelope(139        v = 3, c = verifyCode, s = stripe, capS = stripeCount,140        i = index, t = total, p = isParity,141        d = base64Url(chunk), h = sha256Hex(chunk).take(32),142    )143144    /**145     * Upper bound on one chunk's decoded bytes, not the generator's exact146     * [CHUNK_SIZE]: the browser's decoder derives the shard size from the147     * first scanned chunk precisely so Codexes printed before the current148     * chunk size (800 bytes) still restore, and it enforces no length of its149     * own. The bound sits above a full-range Aztec code's alphanumeric150     * capacity, so it rejects only what no barcode could have carried.151     */152    private const val MAX_CHUNK_BYTES = 4096153    private val MAX_ENCODED_CHUNK_CHARS = base64Url(ByteArray(MAX_CHUNK_BYTES)).length154155    /**156     * Budget for the whole reconstructed payload. [inflate] refuses anything157     * over [MAX_INFLATED_BYTES], so a payload larger than this can never158     * succeed; the headroom covers deflate framing on a payload at the cap.159     */160    private val MAX_RECONSTRUCTED_BYTES = MAX_INFLATED_BYTES + 1024 * 1024161162    private val VERIFY_CODE = Regex("PX-[0-9A-Z]{4}-[0-9A-Z]{4}")163164    /**165     * Stripe ceiling implied by the inflate cap: the largest legal payload166     * divided by the largest legal stripe, rounded up. The extra stripe covers167     * deflate's own framing overhead, which can push the compressed payload a168     * fraction above the uncompressed cap.169     */170    private val MAX_STRIPES: Int =171        ((MAX_INFLATED_BYTES + MAX_DATA_SHARDS_PER_STRIPE * CHUNK_SIZE - 1) /172            (MAX_DATA_SHARDS_PER_STRIPE * CHUNK_SIZE)).toInt() + 1173174    /**175     * Reject a scanned envelope whose declared geometry no real Codex can176     * produce, before anything is allocated from it (PAI-2026-002).177     *178     * The per-chunk hash covers `d` only — `t`, `S`, `s` and `i` are unauthenticated179     * integers, and reconstruction sized `Array(total)` straight from `t`. A180     * barcode with a valid hash over a tiny payload and t = Int.MAX_VALUE181     * therefore cost a multi-gigabyte allocation before it could be rejected.182     * Every bound here is the encoder's own, except the chunk size, which is a183     * ceiling rather than an exact match so older paper still restores (see184     * [MAX_CHUNK_BYTES]). [buildEnvelope] emits at most185     * [MAX_TOTAL_SHARDS_PER_STRIPE] shards per stripe and at most186     * [MAX_STRIPES] stripes.187     */188    internal fun isValidEnvelope(env: ChunkEnvelope): Boolean =189        env.v == CODEX_VERSION &&190            VERIFY_CODE.matches(env.c) &&191            env.capS in 1..MAX_STRIPES &&192            env.s in 0 until env.capS &&193            env.t in 1..MAX_TOTAL_SHARDS_PER_STRIPE &&194            env.i in 0 until env.t &&195            env.d.length in 1..MAX_ENCODED_CHUNK_CHARS &&196            (env.h.length == 16 || env.h.length == 32)197198    fun verifyChunkIntegrity(env: ChunkEnvelope): Boolean {199        if (env.h.length != 16 && env.h.length != 32) return false200        return sha256Hex(fromBase64Url(env.d)).take(env.h.length).equals(env.h, ignoreCase = true)201    }202203    /**204     * Reconstruct the compressed payload from scanned envelopes, tolerating205     * lost codes via Reed-Solomon. Groups by stripe, reconstructs each, then206     * concatenates the data regions in stripe order.207     */208    fun reconstructCompressed(envelopes: List<ChunkEnvelope>): ByteArray {209        val good = envelopes.filter { isValidEnvelope(it) && verifyChunkIntegrity(it) }210        if (good.isEmpty()) throw CodexException(CodexErrorCode.INSUFFICIENT_CHUNKS, "no valid chunks")211        val verifyCode = good.first().c212        val stripeCount = good.first().capS213        if (good.any { it.c != verifyCode || it.capS != stripeCount }) {214            throw CodexException(CodexErrorCode.INVALID_ENVELOPE, "mixed verify codes or stripe counts")215        }216        val byStripe = good.groupBy { it.s }217        var reconstructedBytes = 0L218        val out = java.io.ByteArrayOutputStream()219        for (sIdx in 0 until stripeCount) {220            val chunks = byStripe[sIdx]221                ?: throw CodexException(CodexErrorCode.INSUFFICIENT_CHUNKS, "stripe $sIdx missing entirely")222            val total = chunks.first().t223            if (chunks.any { it.t != total }) {224                throw CodexException(CodexErrorCode.INVALID_ENVELOPE, "mixed shard totals in stripe $sIdx")225            }226            val (dataShards, parityShards) = deriveShardCounts(total)227            val shards = Array(total) { ByteArray(0) }228            val present = BooleanArray(total)229            for (env in chunks) {230                if (env.i in 0 until total) {231                    val decoded = fromBase64Url(env.d)232                    reconstructedBytes += decoded.size233                    if (reconstructedBytes > MAX_RECONSTRUCTED_BYTES) {234                        throw CodexException(235                            CodexErrorCode.PAYLOAD_TOO_LARGE,236                            "Reconstructed payload exceeds the ${MAX_INFLATED_BYTES / (1024 * 1024)} MiB cap",237                        )238                    }239                    shards[env.i] = decoded240                    present[env.i] = true241                }242            }243            ReedSolomon.reconstruct(shards, dataShards, parityShards, present)244            for (d in 0 until dataShards) out.write(shards[d])245        }246        return out.toByteArray()247    }248}249250data class CodexLayout(val stripeCount: Int, val dataChunks: Int, val parityChunks: Int, val totalChunks: Int)