1package com.peachpasswords.codex23import kotlinx.serialization.json.Json4import java.security.MessageDigest5import java.util.Base646import java.util.zip.Deflater7import java.util.zip.Inflater89101112131415object CodexCodec {16 val json = Json { encodeDefaults = true; ignoreUnknownKeys = true }1718 private val urlEncoder = Base64.getUrlEncoder().withoutPadding()19 private val urlDecoder = Base64.getUrlDecoder()2021 2223242526 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 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 while (!inflater.finished()) {53 val n = inflater.inflate(buf)54 if (n == 0 && inflater.needsInput()) break 55 out.write(buf, 0, n)56 if (out.size() > MAX_INFLATED_BYTES) {57 inflater.end()58 throw CodexException(59 CodexErrorCode.PAYLOAD_TOO_LARGE,60 "Inflated payload exceeds the ${MAX_INFLATED_BYTES / (1024 * 1024)} MiB cap",61 )62 }63 }64 inflater.end()65 return out.toByteArray()66 }6768 6970 fun planStripeShardCounts(totalDataShards: Int): IntArray {71 if (totalDataShards <= 0) return IntArray(0)72 val numStripes = maxOf(1, Math.ceil(totalDataShards.toDouble() / MAX_DATA_SHARDS_PER_STRIPE).toInt())73 val base = totalDataShards / numStripes74 val remainder = totalDataShards % numStripes75 return IntArray(numStripes) { base + if (it < remainder) 1 else 0 }76 }7778 data class Stripe(val dataShards: Int, val parityShards: Int, val shards: Array<ByteArray>) {79 val totalChunks: Int get() = dataShards + parityShards80 }8182 83 fun encodeToEnvelopes(compressed: ByteArray, verifyCode: String): Pair<List<ChunkEnvelope>, CodexLayout> {84 val totalDataShards = maxOf(1, Math.ceil(compressed.size.toDouble() / CHUNK_SIZE).toInt())85 val counts = planStripeShardCounts(totalDataShards)86 val envelopes = ArrayList<ChunkEnvelope>()87 var offset = 088 var totalData = 089 var totalParity = 090 for ((sIdx, dataShards) in counts.withIndex()) {91 val region = dataShards * CHUNK_SIZE92 val parityShards = computeParityShards(dataShards)93 val total = dataShards + parityShards94 val shards = Array(total) { ByteArray(CHUNK_SIZE) }95 for (d in 0 until dataShards) {96 val start = offset + d * CHUNK_SIZE97 val len = (compressed.size - start).coerceIn(0, CHUNK_SIZE)98 if (len > 0) System.arraycopy(compressed, start, shards[d], 0, len)99 }100 ReedSolomon.encode(shards, dataShards, parityShards)101 for (i in 0 until total) {102 envelopes.add(buildEnvelope(shards[i], i, total, i >= dataShards, verifyCode, sIdx, counts.size))103 }104 offset += region105 totalData += dataShards106 totalParity += parityShards107 }108 return envelopes to CodexLayout(counts.size, totalData, totalParity, totalData + totalParity)109 }110111 private fun buildEnvelope(112 chunk: ByteArray, index: Int, total: Int, isParity: Boolean,113 verifyCode: String, stripe: Int, stripeCount: Int,114 ): ChunkEnvelope = ChunkEnvelope(115 v = 3, c = verifyCode, s = stripe, capS = stripeCount,116 i = index, t = total, p = isParity,117 d = base64Url(chunk), h = sha256Hex(chunk).take(32),118 )119120 121 private val ENCODED_CHUNK_CHARS = base64Url(ByteArray(CHUNK_SIZE)).length122123 private val VERIFY_CODE = Regex("PX-[0-9A-Z]{4}-[0-9A-Z]{4}")124125 126127128129130131 private val MAX_STRIPES: Int =132 ((MAX_INFLATED_BYTES + MAX_DATA_SHARDS_PER_STRIPE * CHUNK_SIZE - 1) /133 (MAX_DATA_SHARDS_PER_STRIPE * CHUNK_SIZE)).toInt() + 1134135 136137138139140141142143144145146147 internal fun isValidEnvelope(env: ChunkEnvelope): Boolean =148 env.v == CODEX_VERSION &&149 VERIFY_CODE.matches(env.c) &&150 env.capS in 1..MAX_STRIPES &&151 env.s in 0 until env.capS &&152 env.t in 1..MAX_TOTAL_SHARDS_PER_STRIPE &&153 env.i in 0 until env.t &&154 env.d.length == ENCODED_CHUNK_CHARS &&155 (env.h.length == 16 || env.h.length == 32)156157 fun verifyChunkIntegrity(env: ChunkEnvelope): Boolean {158 if (env.h.length != 16 && env.h.length != 32) return false159 return sha256Hex(fromBase64Url(env.d)).take(env.h.length).equals(env.h, ignoreCase = true)160 }161162 163164165166167 fun reconstructCompressed(envelopes: List<ChunkEnvelope>): ByteArray {168 val good = envelopes.filter { isValidEnvelope(it) && verifyChunkIntegrity(it) }169 if (good.isEmpty()) throw CodexException(CodexErrorCode.INSUFFICIENT_CHUNKS, "no valid chunks")170 val verifyCode = good.first().c171 val stripeCount = good.first().capS172 if (good.any { it.c != verifyCode || it.capS != stripeCount }) {173 throw CodexException(CodexErrorCode.INVALID_ENVELOPE, "mixed verify codes or stripe counts")174 }175 val byStripe = good.groupBy { it.s }176 val out = java.io.ByteArrayOutputStream()177 for (sIdx in 0 until stripeCount) {178 val chunks = byStripe[sIdx]179 ?: throw CodexException(CodexErrorCode.INSUFFICIENT_CHUNKS, "stripe $sIdx missing entirely")180 val total = chunks.first().t181 if (chunks.any { it.t != total }) {182 throw CodexException(CodexErrorCode.INVALID_ENVELOPE, "mixed shard totals in stripe $sIdx")183 }184 val (dataShards, parityShards) = deriveShardCounts(total)185 if (dataShards > MAX_DATA_SHARDS_PER_STRIPE) {186 throw CodexException(187 CodexErrorCode.INVALID_ENVELOPE,188 "stripe $sIdx claims $dataShards data shards",189 )190 }191 val shards = Array(total) { ByteArray(0) }192 val present = BooleanArray(total)193 for (env in chunks) {194 if (env.i in 0 until total) {195 shards[env.i] = fromBase64Url(env.d)196 present[env.i] = true197 }198 }199 ReedSolomon.reconstruct(shards, dataShards, parityShards, present)200 for (d in 0 until dataShards) out.write(shards[d])201 }202 return out.toByteArray()203 }204}205206data class CodexLayout(val stripeCount: Int, val dataChunks: Int, val parityChunks: Int, val totalChunks: Int)