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 121122123124125126127128 private const val MAX_CHUNK_BYTES = 4096129 private val MAX_ENCODED_CHUNK_CHARS = base64Url(ByteArray(MAX_CHUNK_BYTES)).length130131 132133134135136 private val MAX_RECONSTRUCTED_BYTES = MAX_INFLATED_BYTES + 1024 * 1024137138 private val VERIFY_CODE = Regex("PX-[0-9A-Z]{4}-[0-9A-Z]{4}")139140 141142143144145146 private val MAX_STRIPES: Int =147 ((MAX_INFLATED_BYTES + MAX_DATA_SHARDS_PER_STRIPE * CHUNK_SIZE - 1) /148 (MAX_DATA_SHARDS_PER_STRIPE * CHUNK_SIZE)).toInt() + 1149150 151152153154155156157158159160161162163164 internal fun isValidEnvelope(env: ChunkEnvelope): Boolean =165 env.v == CODEX_VERSION &&166 VERIFY_CODE.matches(env.c) &&167 env.capS in 1..MAX_STRIPES &&168 env.s in 0 until env.capS &&169 env.t in 1..MAX_TOTAL_SHARDS_PER_STRIPE &&170 env.i in 0 until env.t &&171 env.d.length in 1..MAX_ENCODED_CHUNK_CHARS &&172 (env.h.length == 16 || env.h.length == 32)173174 fun verifyChunkIntegrity(env: ChunkEnvelope): Boolean {175 if (env.h.length != 16 && env.h.length != 32) return false176 return sha256Hex(fromBase64Url(env.d)).take(env.h.length).equals(env.h, ignoreCase = true)177 }178179 180181182183184 fun reconstructCompressed(envelopes: List<ChunkEnvelope>): ByteArray {185 val good = envelopes.filter { isValidEnvelope(it) && verifyChunkIntegrity(it) }186 if (good.isEmpty()) throw CodexException(CodexErrorCode.INSUFFICIENT_CHUNKS, "no valid chunks")187 val verifyCode = good.first().c188 val stripeCount = good.first().capS189 if (good.any { it.c != verifyCode || it.capS != stripeCount }) {190 throw CodexException(CodexErrorCode.INVALID_ENVELOPE, "mixed verify codes or stripe counts")191 }192 val byStripe = good.groupBy { it.s }193 var reconstructedBytes = 0L194 val out = java.io.ByteArrayOutputStream()195 for (sIdx in 0 until stripeCount) {196 val chunks = byStripe[sIdx]197 ?: throw CodexException(CodexErrorCode.INSUFFICIENT_CHUNKS, "stripe $sIdx missing entirely")198 val total = chunks.first().t199 if (chunks.any { it.t != total }) {200 throw CodexException(CodexErrorCode.INVALID_ENVELOPE, "mixed shard totals in stripe $sIdx")201 }202 val (dataShards, parityShards) = deriveShardCounts(total)203 val shards = Array(total) { ByteArray(0) }204 val present = BooleanArray(total)205 for (env in chunks) {206 if (env.i in 0 until total) {207 val decoded = fromBase64Url(env.d)208 reconstructedBytes += decoded.size209 if (reconstructedBytes > MAX_RECONSTRUCTED_BYTES) {210 throw CodexException(211 CodexErrorCode.PAYLOAD_TOO_LARGE,212 "Reconstructed payload exceeds the ${MAX_INFLATED_BYTES / (1024 * 1024)} MiB cap",213 )214 }215 shards[env.i] = decoded216 present[env.i] = true217 }218 }219 ReedSolomon.reconstruct(shards, dataShards, parityShards, present)220 for (d in 0 until dataShards) out.write(shards[d])221 }222 return out.toByteArray()223 }224}225226data class CodexLayout(val stripeCount: Int, val dataChunks: Int, val parityChunks: Int, val totalChunks: Int)