dpp-crypto 0.20.0

use std::collections::HashMap;

use aes_gcm::{
    Aes256Gcm, Nonce,
    aead::{Aead, KeyInit, consts::U12},
};
use ed25519_dalek::SigningKey;
use rand::Rng;
use sha2::{Digest, Sha256};
use zeroize::Zeroize;

use std::path::PathBuf;

use super::crypto::{compute_envelope_hmac, derive_aes_key_sha256};
use super::store::{KeyRecord, KeyRecordMap, KeyStore, default_algorithm};

/// A unique temp-file path for a keystore under test.
fn temp_path(label: &str) -> PathBuf {
    std::env::temp_dir().join(format!(
        "test-keystore-{label}-{}.json",
        uuid::Uuid::now_v7()
    ))
}

fn temp_store() -> KeyStore {
    KeyStore::open(temp_path("default"), "test-passphrase").expect("open key store")
}

/// A fresh [`KeyStore`] plus its backing path, for tests that need to read
/// back or reopen the store file.
fn temp_store_at(label: &str, passphrase: &str) -> (KeyStore, PathBuf) {
    let path = temp_path(label);
    let store = KeyStore::open(&path, passphrase).expect("open key store");
    (store, path)
}

#[test]
fn generate_then_load_roundtrip() {
    let store = temp_store();
    let generated = store.generate_key("issuer-1").expect("generate key");
    let loaded = store.load_key("issuer-1").expect("load key");
    assert_eq!(generated.fingerprint, loaded.fingerprint);
    assert_eq!(
        generated.verifying_key.as_bytes(),
        loaded.verifying_key.as_bytes()
    );
}

#[test]
fn stored_bytes_differ_from_plaintext() {
    let (store, path) = temp_store_at("enc", "test-passphrase");
    let key = store.generate_key("issuer-enc").expect("generate");

    let raw_file = std::fs::read_to_string(&path).expect("read file");
    let plaintext_hex = hex::encode(key.signing_key.as_bytes());
    assert!(
        !raw_file.contains(&plaintext_hex),
        "plaintext key bytes must not appear in the store file"
    );
}

#[test]
fn store_file_contains_argon2id_kdf_marker() {
    let (store, path) = temp_store_at("kdf", "test-passphrase");
    store.generate_key("issuer-kdf").expect("generate");

    let raw_file = std::fs::read_to_string(&path).expect("read file");
    assert!(
        raw_file.contains("argon2id"),
        "store file must contain argon2id KDF marker"
    );
    assert!(
        raw_file.contains("salt"),
        "store file must contain salt field"
    );
}

#[test]
fn reopen_store_from_disk_with_argon2id() {
    let (store, path) = temp_store_at("reopen", "my-passphrase");
    let generated = store.generate_key("issuer-reopen").expect("generate");
    drop(store);

    // Re-open from disk.
    let store2 = KeyStore::open(&path, "my-passphrase").expect("reopen");
    let loaded = store2.load_key("issuer-reopen").expect("load");
    assert_eq!(generated.fingerprint, loaded.fingerprint);
}

#[test]
fn legacy_sha256_store_can_be_opened_and_migrated() {
    // Simulate a V1 (raw HashMap) key store created with SHA-256.
    let path = temp_path("legacy");

    // Create a legacy-format store by directly writing the old format.
    let passphrase = "legacy-pass";
    let legacy_key = derive_aes_key_sha256(passphrase);
    let cipher = Aes256Gcm::new(&legacy_key);

    let signing_key = SigningKey::generate(&mut crate::os_rng());
    let verifying_key = signing_key.verifying_key();
    let fingerprint = hex::encode(Sha256::digest(verifying_key.as_bytes()));

    let mut nonce_bytes = [0u8; 12];
    crate::os_rng().fill_bytes(&mut nonce_bytes);
    let nonce = <&Nonce<U12>>::from(&nonce_bytes);
    let mut raw = signing_key.to_bytes();
    let encrypted = cipher.encrypt(nonce, raw.as_ref()).expect("encrypt");
    raw.zeroize();

    let record = KeyRecord {
        encrypted_signing_key: encrypted,
        nonce: nonce_bytes.to_vec(),
        fingerprint: fingerprint.clone(),
        verifying_key_hex: hex::encode(verifying_key.as_bytes()),
        revoked: false,
        algorithm: default_algorithm(),
    };
    let mut map = HashMap::new();
    map.insert("legacy-key".to_string(), record);
    let bytes = serde_json::to_vec(&map).expect("serialize");
    std::fs::write(&path, bytes).expect("write");

    // `open` refuses it, by name, and says where to go.
    // `expect_err` would need `KeyStore: Debug`, and it deliberately has none —
    // the struct holds a passphrase-derived cipher.
    let Err(refused) = KeyStore::open(&path, passphrase) else {
        panic!("a legacy-KDF store must not open through `open`");
    };
    let message = refused.to_string();
    assert!(message.contains("legacy SHA-256 KDF"), "{message}");
    assert!(message.contains("open_and_migrate"), "{message}");

    // The migration door opens it, upgrades it, and hands back a store that
    // satisfies `open` on its own terms.
    let store = KeyStore::open_and_migrate(&path, passphrase).expect("migrate legacy store");
    let loaded = store.load_key("legacy-key").expect("load legacy key");
    assert_eq!(loaded.fingerprint, fingerprint);

    let raw_file = std::fs::read_to_string(&path).expect("read");
    assert!(raw_file.contains("argon2id"), "file must be migrated");

    // And it is now openable strictly — which is the property that says the
    // upgrade actually finished rather than merely ran.
    KeyStore::open(&path, passphrase).expect("migrated store opens strictly");
}

#[test]
fn malformed_nonce_returns_error_not_panic() {
    // A legacy V1 store (raw HashMap, no HMAC) whose record carries a truncated
    // nonce. Opening must succeed; loading must return an error rather than
    // panicking on the 12-byte nonce check.
    let path = temp_path("badnonce");

    let record = KeyRecord {
        encrypted_signing_key: vec![0u8; 48],
        nonce: vec![0u8; 11], // truncated — not 12 bytes
        fingerprint: "deadbeef".to_string(),
        verifying_key_hex: format!("{:064x}", 0),
        revoked: false,
        algorithm: default_algorithm(),
    };
    let mut map: KeyRecordMap = HashMap::new();
    map.insert("bad-key".to_string(), record);
    std::fs::write(&path, serde_json::to_vec(&map).expect("serialize")).expect("write");

    // Reached through the migration door, since a raw-map store is legacy by
    // definition. Migration itself fails on the malformed record, so the check
    // is on the permissive open that precedes it.
    let store = KeyStore::open_permissively(&path, "any-pass").expect("open legacy store");
    assert!(
        store.load_key("bad-key").is_err(),
        "malformed nonce must return Err, not panic"
    );
}

#[test]
fn archive_key_creates_archived_entry() {
    let store = temp_store();
    store.generate_key("issuer-arc").expect("generate");
    store.archive_key("issuer-arc").expect("archive");
    let archived = store.load_archived_keys("issuer-arc");
    assert_eq!(archived.len(), 1, "expected one archived key");
}

#[test]
fn load_archived_keys_empty_before_rotation() {
    let store = temp_store();
    store.generate_key("issuer-noarc").expect("generate");
    let archived = store.load_archived_keys("issuer-noarc");
    assert!(
        archived.is_empty(),
        "no archived keys before first rotation"
    );
}

// ── Gap 7: atomic rotation + revocation ───────────────────────────────────

#[test]
fn rotate_key_archives_old_and_installs_new() {
    let store = temp_store();
    let k1 = store.generate_key("iss").expect("gen");
    let k2 = store.rotate_key("iss").expect("rotate");
    assert_ne!(
        k1.fingerprint, k2.fingerprint,
        "rotation installs a new key"
    );
    assert_eq!(
        store.load_key("iss").unwrap().fingerprint,
        k2.fingerprint,
        "current key is the new one"
    );
    let archived = store.load_archived_keys("iss");
    assert_eq!(archived.len(), 1);
    assert_eq!(archived[0].fingerprint, k1.fingerprint);
    assert!(
        !archived[0].revoked,
        "hygiene rotation keeps the old key valid"
    );
}

#[test]
fn revoke_and_rotate_marks_old_key_revoked() {
    let store = temp_store();
    let k1 = store.generate_key("iss").expect("gen");
    let k2 = store.revoke_and_rotate("iss").expect("revoke+rotate");
    assert_ne!(k1.fingerprint, k2.fingerprint);
    assert!(!k2.revoked, "the new current key is not revoked");
    let archived = store.load_archived_keys("iss");
    assert_eq!(archived.len(), 1);
    assert_eq!(archived[0].fingerprint, k1.fingerprint);
    assert!(
        archived[0].revoked,
        "compromise rotation marks the old key revoked"
    );
}

#[test]
fn revoked_state_persists_across_reopen() {
    let (store, path) = temp_store_at("revoke", "rev");
    store.generate_key("iss").expect("gen");
    let revoked_fp = store.load_key("iss").unwrap().fingerprint.clone();
    store.revoke_and_rotate("iss").expect("revoke+rotate");
    drop(store);

    let store2 = KeyStore::open(&path, "rev").expect("reopen");
    let archived = store2.load_archived_keys("iss");
    assert!(
        archived
            .iter()
            .any(|k| k.fingerprint == revoked_fp && k.revoked),
        "revoked flag must survive a reopen (and the HMAC must still verify)"
    );
}

#[test]
fn store_file_contains_hmac() {
    let (store, path) = temp_store_at("hmac", "hmac-test");
    store.generate_key("issuer-hmac").expect("generate");

    let raw = std::fs::read_to_string(&path).expect("read file");
    assert!(
        raw.contains("\"hmac\""),
        "store file must contain HMAC field"
    );

    let envelope: serde_json::Value = serde_json::from_str(&raw).unwrap();
    let hmac_val = envelope["hmac"].as_str().unwrap();
    // HMAC-SHA256 produces 64 hex characters
    assert_eq!(hmac_val.len(), 64, "HMAC must be 64 hex chars");
}

#[test]
fn tampered_store_file_rejected_on_open() {
    let (store, path) = temp_store_at("tamper", "tamper-test");
    store.generate_key("issuer-tamper").expect("generate");
    drop(store);

    // Tamper with the file: change a fingerprint value.
    let mut raw = std::fs::read_to_string(&path).expect("read");
    // Replace first hex char of any fingerprint
    if let Some(pos) = raw.find("\"fingerprint\":\"") {
        let fp_start = pos + "\"fingerprint\":\"".len();
        let old_char = raw.as_bytes()[fp_start];
        let new_char = if old_char == b'a' { b'b' } else { b'a' };
        unsafe {
            raw.as_bytes_mut()[fp_start] = new_char;
        }
    }
    std::fs::write(&path, &raw).expect("write tampered file");

    // Re-opening should fail due to HMAC mismatch.
    let result = KeyStore::open(&path, "tamper-test");
    assert!(
        result.is_err(),
        "tampered store should fail integrity check"
    );
    let Err(e) = result else {
        panic!("tampered store should fail integrity check");
    };
    let err = e.to_string();
    assert!(
        err.contains("integrity") || err.contains("tamper"),
        "error should mention integrity, got: {err}"
    );
}

#[test]
fn reopen_with_hmac_succeeds() {
    let (store, path) = temp_store_at("hmac-ok", "hmac-ok");
    store.generate_key("issuer-ok").expect("generate");
    drop(store);

    // Re-open — should pass HMAC verification.
    let store2 = KeyStore::open(&path, "hmac-ok").expect("reopen");
    let loaded = store2.load_key("issuer-ok").expect("load");
    assert!(!loaded.fingerprint.is_empty());
}

/// Regression: HMAC must be identical regardless of HashMap insertion order.
/// Before the BTreeMap fix, two HashMaps with the same entries but different
/// insertion orders could iterate in different sequences and produce different
/// HMAC digests, falsely rejecting a valid store as tampered after restart.
#[test]
fn hmac_is_stable_across_map_insertion_order() {
    let integrity_key = [42u8; 32];

    let make_record = |n: u8| KeyRecord {
        encrypted_signing_key: vec![n; 48],
        nonce: vec![n; 12],
        fingerprint: format!("fp{n:02x}"),
        verifying_key_hex: format!("{:064x}", n),
        revoked: false,
        algorithm: default_algorithm(),
    };

    let keys_fwd = ["zebra", "alpha", "mango", "delta", "beta"];

    let mut map_a: KeyRecordMap = HashMap::new();
    for (i, k) in keys_fwd.iter().enumerate() {
        map_a.insert(k.to_string(), make_record(i as u8));
    }

    // Rebuild with the same key→value mapping but drained and re-inserted
    // in reverse order to try to get a different HashMap layout.
    let mut map_b: KeyRecordMap = HashMap::new();
    for (i, k) in keys_fwd.iter().enumerate() {
        map_b.insert(k.to_string(), make_record(i as u8));
    }
    let entries: Vec<_> = map_b.drain().collect();
    let mut map_b: KeyRecordMap = HashMap::new();
    for (k, v) in entries.into_iter().rev() {
        map_b.insert(k, v);
    }

    let test_salt = "dGVzdHNhbHQ="; // base64("testsalt")
    let hmac_a = compute_envelope_hmac(&integrity_key, "argon2id", test_salt, &map_a).unwrap();
    let hmac_b = compute_envelope_hmac(&integrity_key, "argon2id", test_salt, &map_b).unwrap();

    assert_eq!(
        hmac_a, hmac_b,
        "HMAC must be identical regardless of HashMap insertion/iteration order"
    );
}

// ── algorithm agility ─────────────────────────────────────────────────────────

/// The key record's on-disk shape is unchanged by giving `algorithm` a type:
/// it still serialises as the bare JOSE identifier, so an existing store file
/// is byte-compatible.
#[test]
fn algorithm_serialises_as_its_jose_identifier() {
    let record = KeyRecord::new(vec![1], vec![2], "fp".into(), "aa".into());
    let json = serde_json::to_value(&record).expect("serialise");
    assert_eq!(json["algorithm"], "EdDSA");
}

/// A store written before `algorithm` existed still loads, defaulting to the
/// only algorithm there has ever been.
#[test]
fn record_without_algorithm_field_defaults() {
    let record: KeyRecord = serde_json::from_value(serde_json::json!({
        "encrypted_signing_key": [1],
        "nonce": [2],
        "fingerprint": "fp",
        "verifying_key_hex": "aa",
    }))
    .expect("a pre-agility record must still deserialise");
    assert_eq!(record.algorithm, default_algorithm());
}

/// An *unrecognised* algorithm fails closed at deserialisation rather than
/// loading a key that nothing in this crate can safely sign or verify with.
/// Reaching the signing path and failing there would be a worse outcome: the
/// key would look usable right up to the point it was needed.
#[test]
fn record_with_unknown_algorithm_is_refused() {
    let result: Result<KeyRecord, _> = serde_json::from_value(serde_json::json!({
        "encrypted_signing_key": [1],
        "nonce": [2],
        "fingerprint": "fp",
        "verifying_key_hex": "aa",
        "algorithm": "ES256",
    }));
    assert!(
        result.is_err(),
        "a key record naming an algorithm this build does not implement must not load"
    );
}

/// The public half carries its algorithm, so a reader (the `did:web` builder)
/// never has to assume one to choose a JWK shape.
#[test]
fn public_key_info_carries_the_algorithm() {
    let store = temp_store();
    store.generate_key("alg-info").expect("generate");
    let info = store.public_key("alg-info").expect("public key");
    assert_eq!(info.algorithm, default_algorithm());
}

// ─── Legacy-shape refusal, and what each shape costs ─────────────────────────

/// Strip a field from a store file, producing an older shape on disk.
fn strip_field(path: &PathBuf, field: &str) {
    let raw = std::fs::read_to_string(path).expect("read");
    let mut envelope: serde_json::Value = serde_json::from_str(&raw).expect("parse");
    envelope
        .as_object_mut()
        .expect("envelope is an object")
        .remove(field);
    std::fs::write(path, serde_json::to_vec(&envelope).expect("serialize")).expect("write");
}

/// A store with no integrity tag is refused, not opened quietly.
///
/// This is the one that mattered most. Without the tag, every plaintext field is
/// unauthenticated — including `revoked`, which `dpp-vc`'s `did:web` builder
/// reads to drop a compromised key from the published DID document. An attacker
/// with write access and no passphrase could flip it back and republish a
/// revoked key, and the old `open` accepted that file with an `info!` log.
#[test]
fn a_store_with_no_integrity_tag_is_refused() {
    let (store, path) = temp_store_at("nohmac", "pw");
    store.generate_key("iss").expect("generate");
    drop(store);

    strip_field(&path, "hmac");

    let Err(refused) = KeyStore::open(&path, "pw") else {
        panic!("a store with no integrity tag must not open");
    };
    let message = refused.to_string();
    assert!(message.contains("integrity tag"), "{message}");
    assert!(message.contains("open_and_migrate"), "{message}");

    // Repairable through the migration door, and strictly openable afterwards.
    KeyStore::open_and_migrate(&path, "pw").expect("upgrade restores the tag");
    KeyStore::open(&path, "pw").expect("and it opens strictly once repaired");
}

/// A genuine pre-binding (V3) store is refused, then upgraded.
///
/// Built by hand rather than by stripping a marker off a current store: a V3
/// store's records were sealed with **no** associated data, and a file that
/// merely claims to be V3 while holding bound records is not the thing under
/// test. The migration has to actually re-encrypt, so it has to actually decrypt
/// the old shape first.
#[test]
fn a_store_with_unbound_records_is_refused_then_upgraded() {
    let path = temp_path("unbound");
    let passphrase = "pw";

    let mut salt = [0u8; 16];
    crate::os_rng().fill_bytes(&mut salt);
    let cipher_key =
        super::crypto::derive_aes_key_argon2(passphrase, &salt).expect("derive aes key");
    let cipher = Aes256Gcm::new(&cipher_key);

    let signing_key = SigningKey::generate(&mut crate::os_rng());
    let verifying_key = signing_key.verifying_key();
    let fingerprint = hex::encode(Sha256::digest(verifying_key.as_bytes()));

    let mut nonce_bytes = [0u8; 12];
    crate::os_rng().fill_bytes(&mut nonce_bytes);
    let nonce = <&Nonce<U12>>::from(&nonce_bytes);
    let mut raw = signing_key.to_bytes();
    // No associated data — this is what V3 wrote.
    let encrypted = cipher.encrypt(nonce, raw.as_ref()).expect("encrypt");
    raw.zeroize();

    let mut keys: KeyRecordMap = HashMap::new();
    keys.insert(
        "iss".to_owned(),
        KeyRecord {
            encrypted_signing_key: encrypted,
            nonce: nonce_bytes.to_vec(),
            fingerprint: fingerprint.clone(),
            verifying_key_hex: hex::encode(verifying_key.as_bytes()),
            revoked: false,
            algorithm: default_algorithm(),
        },
    );

    let salt_b64 = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, salt);
    let integrity_key =
        super::crypto::derive_integrity_key(passphrase, &salt).expect("integrity key");
    let hmac = compute_envelope_hmac(&integrity_key, "argon2id", &salt_b64, &keys).expect("hmac");

    // A V3 envelope: argon2id, a valid tag, and no `version`.
    let envelope = serde_json::json!({
        "kdf": "argon2id",
        "salt": salt_b64,
        "hmac": hmac,
        "keys": keys,
    });
    std::fs::write(&path, serde_json::to_vec(&envelope).expect("serialize")).expect("write");

    let Err(refused) = KeyStore::open(&path, passphrase) else {
        panic!("a pre-binding store must not open strictly");
    };
    assert!(
        refused.to_string().contains("per-record binding"),
        "{refused}"
    );

    let upgraded = KeyStore::open_and_migrate(&path, passphrase).expect("upgrade");
    assert_eq!(
        upgraded.load_key("iss").expect("load").fingerprint,
        fingerprint,
        "the key survives the upgrade unchanged"
    );
    KeyStore::open(&path, passphrase).expect("opens strictly once upgraded");
}

/// A record's ciphertext cannot be grafted onto another record's identity.
///
/// The attack the associated data blocks: take the encrypted private key from
/// one record and file it under another record's `verifying_key_hex`,
/// `fingerprint` and `revoked` flag. Before binding, that decrypted cleanly and
/// the store happily reported the wrong public key — and, worse, the wrong
/// revocation state — for that private key. Now the AEAD refuses it.
///
/// The envelope HMAC also catches this, and both are wanted: the HMAC covers the
/// map as a whole, the binding covers each record on its own, and a control that
/// only works in aggregate fails differently from one that works per record.
#[test]
fn a_record_ciphertext_cannot_be_moved_onto_another_identity() {
    let (store, path) = temp_store_at("graft", "pw");
    store.generate_key("alice").expect("generate alice");
    store.generate_key("bob").expect("generate bob");
    drop(store);

    let raw = std::fs::read_to_string(&path).expect("read");
    let mut envelope: serde_json::Value = serde_json::from_str(&raw).expect("parse");
    let alice = envelope["keys"]["alice"].clone();
    let bob = envelope["keys"]["bob"].clone();

    // Bob's plaintext identity, Alice's encrypted private key.
    let mut grafted = bob.clone();
    grafted["encryptedSigningKey"] = alice["encryptedSigningKey"].clone();
    grafted["nonce"] = alice["nonce"].clone();
    envelope["keys"]["bob"] = grafted;

    // Recompute the envelope tag so the HMAC is not what rejects this — the
    // point is that the per-record binding rejects it independently.
    let keys: KeyRecordMap = serde_json::from_value(envelope["keys"].clone()).expect("keys parse");
    let salt_b64 = envelope["salt"].as_str().expect("salt").to_owned();
    let integrity_key = super::crypto::derive_integrity_key(
        "pw",
        &base64::Engine::decode(
            &base64::engine::general_purpose::STANDARD,
            salt_b64.as_str(),
        )
        .expect("salt b64"),
    )
    .expect("integrity key");
    envelope["hmac"] = serde_json::Value::String(
        compute_envelope_hmac(&integrity_key, "argon2id", &salt_b64, &keys).expect("hmac"),
    );
    std::fs::write(&path, serde_json::to_vec(&envelope).expect("serialize")).expect("write");

    let store = KeyStore::open(&path, "pw").expect("envelope is internally consistent");
    assert!(
        store.load_key("bob").is_err(),
        "a ciphertext bound to another fingerprint must not decrypt"
    );
    assert!(
        store.load_key("alice").is_ok(),
        "the untouched record still opens"
    );
}

/// Argon2id derivation is pinned to a known answer, so an upgrade of the
/// `argon2` crate cannot silently change it.
///
/// Every other keystore test derives a key and reads it back inside one
/// process, using the same `argon2` for both halves. A change in the KDF's
/// output would leave all of them green while making every store already on
/// disk unopenable — the one failure mode a round trip structurally cannot
/// see, and the reason this test is a fixture rather than another round trip.
///
/// The vectors were produced under `argon2` 0.5.3, the version in use while
/// the V2/V3 keystore format was written, so they are the bytes existing
/// stores were actually sealed with. Argon2id with explicit parameters is
/// fully specified (RFC 9106), so any correct implementation must reproduce
/// them. A difference here is either a real incompatibility or a deliberate
/// format change needing a keystore version bump and a migration path — never
/// something to re-baseline in place.
///
/// Parameters are fixed by `crypto.rs`: Argon2id, version 0x13, m=19456, t=2,
/// p=1, with the integrity key taking bytes 32..64 of a 64-byte output.
#[test]
fn argon2id_derivation_matches_its_known_answer() {
    let salt: [u8; 16] = [
        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
        0x0f,
    ];
    let passphrase = "correct horse battery staple";

    let aes = super::crypto::derive_aes_key_argon2(passphrase, &salt).expect("derive aes key");
    assert_eq!(
        hex::encode(aes),
        "818259b6310026a8e0dbac5d2e6927abcfdb07b32258fac4f61b18b80f929085",
        "Argon2id AES key derivation changed — every existing keystore is now unopenable"
    );

    let integrity =
        super::crypto::derive_integrity_key(passphrase, &salt).expect("derive integrity key");
    assert_eq!(
        hex::encode(integrity),
        "420feff9f780f2b13b889f7a6dcf263da5a4310898aa51ff4de932e6a0adcd11",
        "Argon2id integrity key derivation changed — every existing keystore fails its HMAC"
    );
}