PeachScriptHardeningTest.kt

1package com.peachpasswords.vault23import com.github.luben.zstd.Zstd4import org.bouncycastle.crypto.digests.Blake3Digest5import org.bouncycastle.crypto.generators.Argon2BytesGenerator6import org.bouncycastle.crypto.params.Argon2Parameters7import org.junit.Assert.assertEquals8import org.junit.Assert.assertThrows9import org.junit.Assert.assertTrue10import org.junit.Test11import javax.crypto.Cipher12import javax.crypto.spec.GCMParameterSpec13import javax.crypto.spec.SecretKeySpec1415/**16 * M-02 regression tests: PeachScript parser must reject attacker-controlled17 * lengths/counts before any unbounded allocation. Each test crafts a payload18 * with a malicious header field and asserts that [PeachPayloadCodec.parsePayload]19 * refuses it with the expected error code.20 */21class PeachScriptHardeningTest {2223    private val passphrase = "test-passphrase"2425    @Test fun oversized_envelope_rejected() {26        // First guard fires before any allocation, so this just needs to be27        // above MAX_ENVELOPE_BYTES.28        val env = ByteArray(50 * 1024 * 1024 + 1)29        val ex = assertThrows(PeachPayloadCodec.PeachScriptException::class.java) {30            PeachPayloadCodec.parsePayload(env, passphrase)31        }32        assertEquals("PAYLOAD_TOO_LARGE", ex.message)33    }3435    @Test fun declared_length_exceeds_actual_remaining_rejected() {36        // v3 plaintext whose declared innerLen (1 MiB) is well under the cap37        // but vastly larger than the actual plaintext — must trip the38        // remaining-bytes check without overflowing Int arithmetic.39        val plaintext = ByteArray(64).apply {40            this[0] = 0x03 // FORMAT_V341            writeU32LE(this, 1, 1024L * 1024) // declared innerLen42        }43        val env = encryptPlaintext(plaintext, passphrase)44        val ex = assertThrows(PeachPayloadCodec.PeachScriptException::class.java) {45            PeachPayloadCodec.parsePayload(env, passphrase)46        }47        assertTrue(48            "expected PAYLOAD_TOO_SHORT or INNER_TOO_LARGE, got ${ex.message}",49            ex.message == "PAYLOAD_TOO_SHORT" || ex.message == "INNER_TOO_LARGE",50        )51    }5253    @Test fun decompression_bomb_rejected() {54        // Section header declares a 4 GiB uncompressed length but the compressed55        // payload is tiny — must be rejected before Zstd allocates the bomb size.56        val body = ByteArray(8 + 9 + 16).apply {57            writeU32BE(this, 0, 0x50434853L) // MAGIC58            this[4] = 0x01 // VERSION59            this[5] = 0x01 // sectionCount60            this[8] = 0x01 // SECTION_ENTRIES61            writeU32LE(this, 9, 0xFFFFFFFFL) // rawLen = 4 GiB (bomb)62            writeU32LE(this, 13, 16L) // compLen = 1663        }64        val env = encryptPlaintext(v3Plaintext(body), passphrase)65        val ex = assertThrows(PeachPayloadCodec.PeachScriptException::class.java) {66            PeachPayloadCodec.parsePayload(env, passphrase)67        }68        assertEquals("SECTION_TOO_LARGE", ex.message)69    }7071    @Test fun entry_count_overflow_rejected() {72        // Entries section declares a 4 GiB count — must be rejected before the73        // ArrayList pre-allocation runs.74        val entriesRaw = ByteArray(4).apply { writeU32LE(this, 0, 0xFFFFFFFFL) }75        val compressed = Zstd.compress(entriesRaw, 19)76        val body = ByteArray(8 + 9 + compressed.size).apply {77            writeU32BE(this, 0, 0x50434853L)78            this[4] = 0x0179            this[5] = 0x0180            this[8] = 0x01 // SECTION_ENTRIES81            writeU32LE(this, 9, entriesRaw.size.toLong())82            writeU32LE(this, 13, compressed.size.toLong())83            System.arraycopy(compressed, 0, this, 17, compressed.size)84        }85        val env = encryptPlaintext(v3Plaintext(body), passphrase)86        val ex = assertThrows(PeachPayloadCodec.PeachScriptException::class.java) {87            PeachPayloadCodec.parsePayload(env, passphrase)88        }89        assertEquals("ENTRY_COUNT_OVERFLOW", ex.message)90    }9192    @Test fun oversize_passkey_blob_rejected() {93        // One passkey entry whose blobLen is 4 GiB — must be rejected before94        // r.read() allocates the declared size.95        // Entry header layout: type(1) | id(16) | flags(1) | flags2(1) |96        // createdSec(4) | updatedDelta(4) | deletedSec(4) | titleIdx(2) | blobLen(4)97        val entry = ByteArray(33 + 4).apply {98            this[0] = 0x04 // ENTRY_TYPE_PASSKEY99            writeU32LE(this, 33, 0xFFFFFFFFL) // blobLen = 4 GiB100        }101        val entriesRaw = ByteArray(6 + entry.size).apply {102            writeU32LE(this, 0, 1L) // count = 1103            writeU16LE(this, 4, 0) // empty string table104            System.arraycopy(entry, 0, this, 6, entry.size)105        }106        val compressed = Zstd.compress(entriesRaw, 19)107        val body = ByteArray(8 + 9 + compressed.size).apply {108            writeU32BE(this, 0, 0x50434853L)109            this[4] = 0x01110            this[5] = 0x01111            this[8] = 0x01112            writeU32LE(this, 9, entriesRaw.size.toLong())113            writeU32LE(this, 13, compressed.size.toLong())114            System.arraycopy(compressed, 0, this, 17, compressed.size)115        }116        val env = encryptPlaintext(v3Plaintext(body), passphrase)117        val ex = assertThrows(PeachPayloadCodec.PeachScriptException::class.java) {118            PeachPayloadCodec.parsePayload(env, passphrase)119        }120        assertEquals("PASSKEY_BLOB_TOO_LARGE", ex.message)121    }122123    // ---- crypto / framing helpers (mirror production wire format) ---------124125    private fun v3Plaintext(body: ByteArray): ByteArray {126        val checksum = blake3Truncated(body)127        val inner = body + checksum128        val out = ByteArray(5 + inner.size)129        out[0] = 0x03 // FORMAT_V3130        writeU32LE(out, 1, inner.size.toLong())131        System.arraycopy(inner, 0, out, 5, inner.size)132        return out133    }134135    private fun encryptPlaintext(plaintext: ByteArray, passphrase: String): ByteArray {136        val salt = ByteArray(16).also { java.security.SecureRandom().nextBytes(it) }137        val nonce = ByteArray(12).also { java.security.SecureRandom().nextBytes(it) }138        val key = deriveKey(passphrase, salt)139        val ctTag = try {140            val cipher = Cipher.getInstance("AES/GCM/NoPadding")141            cipher.init(Cipher.ENCRYPT_MODE, SecretKeySpec(key, "AES"), GCMParameterSpec(128, nonce))142            cipher.doFinal(plaintext)143        } finally { key.fill(0) }144        return salt + nonce + ctTag145    }146147    private fun deriveKey(passphrase: String, salt: ByteArray): ByteArray {148        val params = Argon2Parameters.Builder(Argon2Parameters.ARGON2_id)149            .withSalt(salt).withParallelism(4).withMemoryAsKB(65536).withIterations(3).build()150        val gen = Argon2BytesGenerator().apply { init(params) }151        val out = ByteArray(32)152        gen.generateBytes(passphrase.toByteArray(Charsets.UTF_8), out)153        return out154    }155156    private fun blake3Truncated(data: ByteArray): ByteArray {157        val digest = Blake3Digest(256)158        digest.update(data, 0, data.size)159        val full = ByteArray(32)160        digest.doFinal(full, 0)161        return full.copyOfRange(0, 4)162    }163164    private fun writeU32LE(buf: ByteArray, off: Int, v: Long) {165        buf[off] = (v and 0xff).toByte()166        buf[off + 1] = ((v ushr 8) and 0xff).toByte()167        buf[off + 2] = ((v ushr 16) and 0xff).toByte()168        buf[off + 3] = ((v ushr 24) and 0xff).toByte()169    }170171    private fun writeU32BE(buf: ByteArray, off: Int, v: Long) {172        buf[off] = ((v ushr 24) and 0xff).toByte()173        buf[off + 1] = ((v ushr 16) and 0xff).toByte()174        buf[off + 2] = ((v ushr 8) and 0xff).toByte()175        buf[off + 3] = (v and 0xff).toByte()176    }177178    private fun writeU16LE(buf: ByteArray, off: Int, v: Int) {179        buf[off] = (v and 0xff).toByte()180        buf[off + 1] = ((v ushr 8) and 0xff).toByte()181    }182}