armdb 0.7.0

#![cfg(feature = "encryption")]

#[cfg(feature = "var-collections")]
use armdb::VarTree;
use armdb::{Config, ConstTree, DbError};
use tempfile::tempdir;

fn encrypted_config(key: [u8; 32]) -> Config {
    let mut config = Config::test();
    config.encryption_key = Some(key);
    config
}

#[test]
fn test_encrypted_const_tree_basic() {
    let dir = tempdir().unwrap();
    let key = [0x42u8; 32];
    let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), encrypted_config(key)).unwrap();

    // Basic CRUD
    for i in 0..100u64 {
        let k = i.to_be_bytes();
        let v = (i * 10).to_be_bytes();
        tree.put(&k, &v).unwrap();
    }

    for i in 0..100u64 {
        let k = i.to_be_bytes();
        let expected = (i * 10).to_be_bytes();
        assert_eq!(tree.get(&k).unwrap(), expected);
    }

    // Delete
    let k = 50u64.to_be_bytes();
    tree.delete(&k).unwrap();
    assert!(tree.get(&k).is_none());
    assert_eq!(tree.len(), 99);
}

#[test]
fn test_encrypted_recovery() {
    let dir = tempdir().unwrap();
    let key = [0x42u8; 32];

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), encrypted_config(key)).unwrap();

        for i in 0..100u64 {
            let k = i.to_be_bytes();
            let v = (i * 10).to_be_bytes();
            tree.put(&k, &v).unwrap();
        }

        tree.close().unwrap();
    }

    // Reopen with same key
    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), encrypted_config(key)).unwrap();

        assert_eq!(tree.len(), 100);

        for i in 0..100u64 {
            let k = i.to_be_bytes();
            let expected = (i * 10).to_be_bytes();
            assert_eq!(
                tree.get(&k).unwrap(),
                expected,
                "encrypted recovery mismatch at key {}",
                i
            );
        }
    }
}

/// Recovery must skip mid-file page padding in an *encrypted* data file.
///
/// An encrypted shard is page-aligned, so `flush()` pads the partial page to a
/// 4096 boundary and the next write lands at that boundary. The padding region
/// on disk is the ciphertext of zeros, not raw zeros — recovery only recognises
/// it after decrypting through the encrypted read path. Interleaving puts with
/// flushes produces several internal padding regions the recovery scan must
/// walk over without mistaking them for EOF, corruption, or bogus entries.
#[test]
fn test_encrypted_recovery_mid_file_padding() {
    let dir = tempdir().unwrap();
    let key = [0x42u8; 32];
    let mut config = encrypted_config(key);
    config.shard_count = 1;
    config.max_file_size = 1 << 20; // keep all pages in one data file

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config.clone()).unwrap();
        for i in 0..5u64 {
            tree.put(&i.to_be_bytes(), &(i * 10).to_be_bytes()).unwrap();
            // Final flush pads the partial page; the next put starts at the
            // next page boundary, leaving mid-file padding in between.
            tree.flush().unwrap();
        }
        tree.close().unwrap();
    }

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config).unwrap();
        assert_eq!(
            tree.len(),
            5,
            "all entries recovered across mid-file padding"
        );
        for i in 0..5u64 {
            assert_eq!(
                tree.get(&i.to_be_bytes()).unwrap(),
                (i * 10).to_be_bytes(),
                "encrypted mid-file-padding recovery mismatch at key {i}"
            );
        }
    }
}

#[test]
fn test_encrypted_wrong_key_fails() {
    let dir = tempdir().unwrap();
    let key_a = [0x42u8; 32];
    let key_b = [0x99u8; 32];

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), encrypted_config(key_a)).unwrap();

        for i in 0..50u64 {
            let k = i.to_be_bytes();
            let v = i.to_be_bytes();
            tree.put(&k, &v).unwrap();
        }

        tree.close().unwrap();
    }

    // Reopen with different key — should fail during recovery
    let result = ConstTree::<[u8; 8], 8>::open(dir.path(), encrypted_config(key_b));
    assert!(result.is_err(), "opening with wrong key should fail");
}

#[test]
fn test_encrypted_no_key_fails() {
    let dir = tempdir().unwrap();
    let key = [0x42u8; 32];

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), encrypted_config(key)).unwrap();
        tree.put(&1u64.to_be_bytes(), &1u64.to_be_bytes()).unwrap();
        tree.close().unwrap();
    }

    // Reopen without encryption key — should fail
    let result = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test());
    assert!(
        matches!(&result, Err(DbError::Config(_))),
        "opening encrypted db without key should return Config error, got {:?}",
        result.as_ref().err()
    );
}

#[test]
fn test_unencrypted_with_key_fails() {
    let dir = tempdir().unwrap();

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();
        tree.put(&1u64.to_be_bytes(), &1u64.to_be_bytes()).unwrap();
        tree.close().unwrap();
    }

    // Reopen with encryption key on unencrypted db — should fail
    let key = [0x42u8; 32];
    let result = ConstTree::<[u8; 8], 8>::open(dir.path(), encrypted_config(key));
    assert!(
        matches!(&result, Err(DbError::Config(_))),
        "opening unencrypted db with key should return Config error, got {:?}",
        result.as_ref().err()
    );
}

#[test]
fn test_encrypted_compaction() {
    let dir = tempdir().unwrap();
    let key = [0x42u8; 32];
    let mut config = encrypted_config(key);
    config.max_file_size = 8192;
    config.write_buffer_size = 8192;
    config.compaction_threshold = 0.1;

    let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config).unwrap();

    // Write and overwrite to create dead bytes
    for i in 0..500u64 {
        let k = i.to_be_bytes();
        let v = i.to_be_bytes();
        tree.put(&k, &v).unwrap();
    }
    for i in 0..250u64 {
        let k = i.to_be_bytes();
        let v = (i + 1000).to_be_bytes();
        tree.put(&k, &v).unwrap();
    }

    tree.compact().unwrap();

    // Verify all data after compaction
    for i in 0..500u64 {
        let k = i.to_be_bytes();
        let expected = if i < 250 {
            (i + 1000).to_be_bytes()
        } else {
            i.to_be_bytes()
        };
        assert_eq!(
            tree.get(&k).unwrap(),
            expected,
            "encrypted compaction mismatch at key {}",
            i
        );
    }
}

/// Compaction must skip mid-file page padding inside an *encrypted* immutable
/// data file and still copy the entries that follow it. A seeding `flush()`
/// after the first put pads the first page; the subsequent writes fill and
/// rotate that file, so the now-immutable `000000.data` carries zero padding
/// (ciphertext of zeros) between its first entry and the page boundary. If the
/// encrypted compaction scan mis-handled that padding, every entry after the
/// padding would lose its `DiskLoc` and reads would fail.
#[test]
fn test_encrypted_compaction_skips_mid_file_padding() {
    let dir = tempdir().unwrap();
    let key = [0x42u8; 32];
    let mut config = encrypted_config(key);
    config.shard_count = 1;
    config.max_file_size = 8192;
    config.write_buffer_size = 8192;
    config.compaction_threshold = 0.1;

    let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config).unwrap();

    // Seed mid-file padding in the first page, then keep writing so the file
    // fills past the padding and rotates into an immutable file.
    tree.put(&0u64.to_be_bytes(), &0u64.to_be_bytes()).unwrap();
    tree.flush().unwrap();

    for i in 0..500u64 {
        tree.put(&i.to_be_bytes(), &i.to_be_bytes()).unwrap();
    }
    for i in 0..250u64 {
        tree.put(&i.to_be_bytes(), &(i + 1000).to_be_bytes())
            .unwrap();
    }

    tree.compact().unwrap();

    for i in 0..500u64 {
        let expected = if i < 250 {
            (i + 1000).to_be_bytes()
        } else {
            i.to_be_bytes()
        };
        assert_eq!(
            tree.get(&i.to_be_bytes()).unwrap(),
            expected,
            "encrypted compaction-over-padding mismatch at key {i}"
        );
    }
}

#[test]
fn distinct_databases_same_key_produce_distinct_ciphertext() {
    // Две зашифрованные БД, один мастер-ключ, идентичные данные. Пер-БД соль
    // обязана сделать их шифртекст на диске разным (нет cross-database reuse).
    let key = [0x42u8; 32];
    let dir_a = tempdir().unwrap();
    let dir_b = tempdir().unwrap();
    for dir in [dir_a.path(), dir_b.path()] {
        let mut cfg = encrypted_config(key);
        cfg.shard_count = 1;
        let tree = ConstTree::<[u8; 8], 8>::open(dir, cfg).unwrap();
        tree.put(&1u64.to_be_bytes(), &7u64.to_be_bytes()).unwrap();
        tree.close().unwrap();
    }
    let data_a = std::fs::read(dir_a.path().join("shard_000/000001.data")).unwrap();
    let data_b = std::fs::read(dir_b.path().join("shard_000/000001.data")).unwrap();
    assert!(
        !data_a.is_empty() && !data_b.is_empty(),
        "data files must be non-empty"
    );
    assert_ne!(
        data_a, data_b,
        "same key + data but different per-DB salt must yield different ciphertext"
    );
}

#[test]
fn encrypted_multishard_reopen_reads_all() {
    let key = [0x77u8; 32];
    let dir = tempdir().unwrap();
    let mut cfg = encrypted_config(key);
    cfg.shard_count = 2;
    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), cfg.clone()).unwrap();
        for i in 0..200u64 {
            tree.put(&i.to_be_bytes(), &(i * 3).to_be_bytes()).unwrap();
        }
        tree.close().unwrap();
    }
    // Переоткрытие: db.salt читается, под-ключи выводятся заново, все шарды читаются.
    let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), cfg).unwrap();
    for i in 0..200u64 {
        assert_eq!(tree.get(&i.to_be_bytes()).unwrap(), (i * 3).to_be_bytes());
    }
}

#[test]
fn missing_salt_fails_closed() {
    let key = [0x42u8; 32];
    let dir = tempdir().unwrap();
    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), encrypted_config(key)).unwrap();
        tree.put(&1u64.to_be_bytes(), &2u64.to_be_bytes()).unwrap();
        tree.close().unwrap();
    }
    // Удаление соли (как у до-фикса зашифрованной БД) → fail-closed отказ открыться.
    std::fs::remove_file(dir.path().join("db.salt")).unwrap();
    let err = ConstTree::<[u8; 8], 8>::open(dir.path(), encrypted_config(key))
        .err()
        .unwrap();
    assert!(
        matches!(err, DbError::FormatMismatch(_)),
        "expected FormatMismatch, got {err:?}"
    );
}

#[test]
fn fresh_encrypted_db_writes_16_byte_salt() {
    let dir = tempdir().unwrap();
    let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), encrypted_config([0x42; 32])).unwrap();
    tree.close().unwrap();
    let meta = std::fs::metadata(dir.path().join("db.salt")).unwrap();
    assert_eq!(meta.len(), 16, "db.salt must be exactly 16 bytes");
}

#[cfg(feature = "var-collections")]
#[test]
fn test_encrypted_var_tree() {
    let dir = tempdir().unwrap();
    let key = [0x42u8; 32];
    let tree = VarTree::<[u8; 8]>::open(dir.path(), encrypted_config(key)).unwrap();

    for i in 0..100u64 {
        let k = i.to_be_bytes();
        let v = format!("encrypted_val_{i}");
        tree.put(&k, v.as_bytes()).unwrap();
    }

    for i in 0..100u64 {
        let k = i.to_be_bytes();
        let expected = format!("encrypted_val_{i}");
        assert_eq!(tree.get(&k).unwrap().as_ref(), expected.as_bytes());
    }

    // Delete and verify
    tree.delete(&50u64.to_be_bytes()).unwrap();
    assert!(tree.get(&50u64.to_be_bytes()).is_none());
    assert_eq!(tree.len(), 99);
}