flamenet-messenger-android/engine/src/test/kotlin/io/flamenet/e2e/SecureBytesTest.kt

package io.flamenet.e2e

import kotlin.test.Test
import kotlin.test.assertContentEquals
import kotlin.test.assertEquals
import kotlin.test.assertFailsWith
import kotlin.test.assertFalse
import kotlin.test.assertTrue

/**
 * Zeroing key material once it has been used (ios#39, Android half).
 */
class SecureBytesTest {
    private val passphrase = "a-long-enough-passphrase"
    private val salt = ByteArray(16) { it.toByte() }

    // ---- The failure that would matter most -----------------------------.

    /**
     * The pinned vectors already cover `authSecret` and `accountKeySeed`.
     */
    @Test
    fun `wiping the intermediate did not change what is derived`() {
        val (auth, account) = AccountKdf.signInSecrets(passphrase, salt)
        assertContentEquals(AccountKdf.authSecret(passphrase, salt), auth)
        assertContentEquals(AccountKdf.accountKeySeed(passphrase, salt), account)
    }

    /**
     * And neither is zeroes, which is what "derived from a wiped buffer" would
     * look like.
     */
    @Test
    fun `no derived secret is a run of zeroes`() {
        val (auth, account) = AccountKdf.signInSecrets(passphrase, salt)
        assertFalse(auth.all { it == 0.toByte() }, "the auth secret is all zeroes")
        assertFalse(account.all { it == 0.toByte() }, "the account seed is all zeroes")
        assertFalse(
            AccountKdf.identityKey(passphrase, salt).all { it == 0.toByte() },
            "the identity key is all zeroes",
        )
    }

    /** Derivation stays deterministic across calls, wiping or not. */
    @Test
    fun `derivation is still deterministic`() {
        assertContentEquals(
            AccountKdf.identityKey(passphrase, salt),
            AccountKdf.identityKey(passphrase, salt),
        )
    }

    /** Different contexts still separate, so the wipe did not collapse them. */
    @Test
    fun `the contexts still produce different secrets`() {
        val auth = AccountKdf.authSecret(passphrase, salt)
        val identity = AccountKdf.identityKey(passphrase, salt)
        val account = AccountKdf.accountKeySeed(passphrase, salt)
        assertFalse(auth.contentEquals(identity))
        assertFalse(auth.contentEquals(account))
        assertFalse(identity.contentEquals(account))
    }

    // ---- The wiping itself ----------------------------------------------.

    @Test
    fun `wipe zeroes the array it is given`() {
        val secret = ByteArray(32) { 0x41 }
        SecureBytes.wipe(secret)
        assertTrue(secret.all { it == 0.toByte() })
    }

    @Test
    fun `wiping nothing is not an error`() {
        SecureBytes.wipe(null)
        SecureBytes.wipe(ByteArray(0))
    }

    @Test
    fun `withSecret wipes on the way out`() {
        val secret = ByteArray(32) { 0x41 }
        val seen = SecureBytes.withSecret(secret) { it.copyOf() }
        assertTrue(seen.all { it == 0x41.toByte() }, "the body saw a wiped buffer")
        assertTrue(secret.all { it == 0.toByte() }, "the buffer was not wiped")
    }

    /**
     * The case the `finally` exists for.
     */
    @Test
    fun `withSecret wipes when the body throws`() {
        val secret = ByteArray(32) { 0x41 }
        assertFailsWith<IllegalStateException> {
            SecureBytes.withSecret(secret) { error("derivation failed half-way") }
        }
        assertTrue(secret.all { it == 0.toByte() }, "a throw left the material behind")
    }

    /**
     * The honest limit, asserted rather than left to the prose.
     */
    @Test
    fun `a copy taken beforehand survives the wipe`() {
        val secret = ByteArray(32) { 0x41 }
        val copy = SecureBytes.copyOf(secret)
        SecureBytes.wipe(secret)
        assertTrue(secret.all { it == 0.toByte() })
        assertTrue(
            copy.all { it == 0x41.toByte() },
            "the copy was wiped too, which would mean these share a buffer",
        )
    }

    @Test
    fun `withSecret returns what the body returned`() {
        assertEquals(7, SecureBytes.withSecret(ByteArray(4)) { 7 })
    }
}