khadhroony-solana-project/crates/ksp-wallet-lib/unit_tests/crypto.rs

// file: crates/ksp-wallet-lib/unit_tests/crypto.rs
// version: 2

use base64::Engine; // rust-rules: derive-import

const VECTOR: &[u8] = include_bytes!("../tests/fixtures/kspwallet_v1_crypto_vectors.json");

#[derive(serde::Deserialize)]
struct CryptoVector {
    password_utf8: std::string::String,
    salt: std::string::String,
    memory_kib: u32,
    iterations: u32,
    parallelism: u32,
    derived_key: std::string::String,
    content_key: std::string::String,
    nonce: std::string::String,
    aad: std::string::String,
    wrapped_key: std::string::String,
}

#[derive(serde::Deserialize)]
struct CryptoVectors {
    vector_version: u32,
    warning: std::string::String,
    argon2id_xchacha20poly1305_wrap: CryptoVector,
}

#[test]
fn deterministic_argon2id_and_xchacha_wrap_vector_matches_external_canary() -> ksp_core_lib::Result<()> {
    let vectors_result = serde_json::from_slice::<CryptoVectors>(VECTOR);
    let vectors = match vectors_result {
        std::result::Result::Ok(vectors) => vectors,
        std::result::Result::Err(_) => {
            return std::result::Result::Err(ksp_core_lib::Error::new(crate::ERROR_CODE_FORMAT_INVALID, "Wallet crypto test vector JSON is invalid"));
        },
    };
    assert_eq!(vectors.vector_version, 1);
    assert!(vectors.warning.contains("TEST ONLY"));
    let vector = vectors.argon2id_xchacha20poly1305_wrap;
    let salt = match decode(vector.salt.as_str()) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let expected_derived = match decode_fixed::<32>(vector.derived_key.as_str()) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let content_key = match decode_fixed::<32>(vector.content_key.as_str()) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let nonce = match decode_fixed::<{ crate::KSPWALLET_V1_XCHACHA_NONCE_BYTES }>(vector.nonce.as_str()) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let aad = match decode(vector.aad.as_str()) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let expected_wrapped = match decode(vector.wrapped_key.as_str()) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };

    let derived_result = super::derive_argon2id(vector.password_utf8.as_bytes(), salt.as_slice(), vector.memory_kib, vector.iterations, vector.parallelism);
    let derived = match derived_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    assert_eq!(derived.as_bytes(), &expected_derived);

    let content = crate::SecretKeyV1::from_bytes(content_key);
    let wrapped_result = crate::wrap_key(&derived, &content, &nonce, aad.as_slice());
    let wrapped = match wrapped_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    assert_eq!(wrapped, expected_wrapped);

    let unwrapped_result = crate::unwrap_key(&derived, &nonce, aad.as_slice(), wrapped.as_slice());
    let unwrapped = match unwrapped_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    assert_eq!(unwrapped.as_bytes(), content.as_bytes());
    return std::result::Result::Ok(());
}

#[test]
fn xchacha_tampering_is_reported_as_generic_authentication_failure() -> ksp_core_lib::Result<()> {
    let key = crate::SecretKeyV1::from_bytes([0x11_u8; 32]);
    let nonce = [0x22_u8; crate::KSPWALLET_V1_XCHACHA_NONCE_BYTES];
    let plaintext = [0x33_u8; 32];
    let encrypted_result = crate::encrypt_bytes(&key, &nonce, b"kspwallet-test-aad", plaintext.as_slice());
    let mut encrypted = match encrypted_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    encrypted[0] ^= 1;
    let result = crate::decrypt_bytes(&key, &nonce, b"kspwallet-test-aad", encrypted.as_slice());
    assert!(result.is_err());
    let error = match result {
        std::result::Result::Err(error) => error,
        std::result::Result::Ok(_) => return std::result::Result::Ok(()),
    };
    assert_eq!(error.code(), crate::ERROR_CODE_AUTHENTICATION_FAILED);
    return std::result::Result::Ok(());
}

#[test]
fn secret_key_debug_is_redacted_and_random_sources_are_callable() -> ksp_core_lib::Result<()> {
    let key_result = crate::SecretKeyV1::random();
    let key = match key_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let nonce_result = crate::random_nonce();
    let nonce = match nonce_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let rendered = format!("{key:?}");
    assert_eq!(rendered, "SecretKeyV1(<redacted>)");
    assert_eq!(nonce.len(), crate::KSPWALLET_V1_XCHACHA_NONCE_BYTES);
    return std::result::Result::Ok(());
}

#[test]
#[ignore = "operator-only Argon2 calibration; run single-threaded with --nocapture"]
fn benchmark_argon2_creation_candidates() -> ksp_core_lib::Result<()> {
    let password = b"KSPWALLET-BENCHMARK-TEST-ONLY-PASSWORD";
    let salt = [0x5a_u8; 16];
    let candidates = [(65_536_u32, 3_u32, 1_u32), (131_072_u32, 3_u32, 1_u32), (262_144_u32, 3_u32, 1_u32)];
    for candidate in candidates {
        let started = std::time::Instant::now();
        let result = super::derive_argon2id(password, salt.as_slice(), candidate.0, candidate.1, candidate.2);
        match result {
            std::result::Result::Ok(_) => {},
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        }
        println!(
            "KSP Argon2id candidate memory_kib={} iterations={} parallelism={} elapsed_ms={}",
            candidate.0,
            candidate.1,
            candidate.2,
            started.elapsed().as_millis()
        );
    }
    return std::result::Result::Ok(());
}

fn decode(value: &str) -> ksp_core_lib::Result<std::vec::Vec<u8>> {
    let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(value.as_bytes());
    return match decoded {
        std::result::Result::Ok(value) => std::result::Result::Ok(value),
        std::result::Result::Err(_) => {
            std::result::Result::Err(ksp_core_lib::Error::new(crate::ERROR_CODE_FORMAT_INVALID, "Wallet crypto test vector Base64url is invalid"))
        },
    };
}

fn decode_fixed<const LENGTH: usize>(value: &str) -> ksp_core_lib::Result<[u8; LENGTH]> {
    let decoded = match decode(value) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let converted = <[u8; LENGTH]>::try_from(decoded.as_slice());
    return match converted {
        std::result::Result::Ok(value) => std::result::Result::Ok(value),
        std::result::Result::Err(_) => {
            std::result::Result::Err(ksp_core_lib::Error::new(crate::ERROR_CODE_FORMAT_INVALID, "Wallet crypto test vector has an invalid fixed-size field"))
        },
    };
}