zkv 0.1.2

//! `.zkv` 加密容器:文件头、解锁/保存/锁定。对应 PRD §3、§4。
//!
//! 文件格式(小端,见 PRD §4,头长 58 字节):
//! ```text
//! [4 Magic "ZKV1"][1 ver][1 flags][4 m_kib][4 t_cost][4 p_cost][16 salt][24 nonce][N ciphertext]
//! ```
//!
//! 对外接口:
//! - [`VaultHeader`] / [`HEADER_LEN`] / [`MAGIC`]
//! - [`create`] / [`unlock`] / [`save`]:使用默认 KDF 参数(64MiB/3/4)的生产入口。
//! - [`create_with_params`] / [`unlock_with_params`] / [`save_with_params`]:可注入 KDF 参数,
//!   **测试**用 `KdfParams{4096,1,1}` 加速(见模块测试)。
//!
//! ## 测试 KDF 加速策略
//! 默认 `KdfParams::default()`(64MiB/3/4)每次派生约几百 ms,多个集成测试会很慢。
//! 因此 `create`/`unlock`/`save` 均拆出 `_with_params` 版本,核心逻辑只实现一次;
//! 测试统一用小参数 `KdfParams{4096,1,1}` 派生(与 crypto.rs 单测一致)。
//! `VaultHeader` 里已保存 params,故 `unlock` 只需调用方传一次小 params 即可正确派生。
//!
//! 原子写:写 `<path>.tmp`(0600)→ `fs::rename` → 目标 0600。
//!
//! 规则(L2):仅依赖 [`crate::error`]/[`crate::crypto`]/[`crate::db`],不引用 store/search/app/ui。

use std::fs;
use std::path::{Path, PathBuf};

use crate::crypto::{self, Encrypted, KdfParams, MasterKey};
use crate::db::Database;
use crate::error::{Error, Result};

/// 文件魔数 `ZKV1`。
pub const MAGIC: [u8; 4] = *b"ZKV1";
/// 固定头长度:4+1+1+4+4+4+16+24 = 58 字节。
pub const HEADER_LEN: usize = 58;
/// 当前容器版本。
pub const VERSION: u8 = 1;

/// `.zkv` 文件头。小端序,布局见 PRD §4。
#[derive(Debug, Clone)]
pub struct VaultHeader {
    /// 版本(当前固定 1)。
    pub version: u8,
    /// 保留标志位(当前置 0)。
    pub flags: u8,
    /// Argon2id 参数。
    pub kdf: KdfParams,
    /// 16 字节 salt。
    pub salt: [u8; 16],
    /// 本次加密的 24 字节 XChaCha20 nonce。
    pub nonce: [u8; 24],
}

impl VaultHeader {
    /// 序列化为 58 字节(小端)。
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut out = Vec::with_capacity(HEADER_LEN);
        out.extend_from_slice(&MAGIC);
        out.push(self.version);
        out.push(self.flags);
        out.extend_from_slice(&self.kdf.m_kib.to_le_bytes());
        out.extend_from_slice(&self.kdf.t_cost.to_le_bytes());
        out.extend_from_slice(&self.kdf.p_cost.to_le_bytes());
        out.extend_from_slice(&self.salt);
        out.extend_from_slice(&self.nonce);
        debug_assert_eq!(out.len(), HEADER_LEN);
        out
    }

    /// 从字节解析头(至少 58 字节)。校验 Magic/Version,不符返回 [`Error::CorruptFile`]。
    pub fn parse(bytes: &[u8]) -> Result<VaultHeader> {
        if bytes.len() < HEADER_LEN {
            return Err(Error::CorruptFile(format!(
                "header too short: {} < {HEADER_LEN}",
                bytes.len()
            )));
        }
        if bytes[0..4] != MAGIC {
            return Err(Error::CorruptFile("bad magic".to_string()));
        }
        let version = bytes[4];
        if version != VERSION {
            return Err(Error::CorruptFile(format!(
                "unsupported version: {version}"
            )));
        }
        let flags = bytes[5];
        let m_kib = u32::from_le_bytes([bytes[6], bytes[7], bytes[8], bytes[9]]);
        let t_cost = u32::from_le_bytes([bytes[10], bytes[11], bytes[12], bytes[13]]);
        let p_cost = u32::from_le_bytes([bytes[14], bytes[15], bytes[16], bytes[17]]);
        let mut salt = [0u8; 16];
        salt.copy_from_slice(&bytes[18..34]);
        let mut nonce = [0u8; 24];
        nonce.copy_from_slice(&bytes[34..58]);

        Ok(VaultHeader {
            version,
            flags,
            kdf: KdfParams {
                m_kib,
                t_cost,
                p_cost,
            },
            salt,
            nonce,
        })
    }
}

/// 创建一个新的空 `.zkv` 库(默认 KDF 参数)。
pub fn create(path: &Path, passphrase: &str) -> Result<()> {
    create_with_params(path, passphrase, &KdfParams::default())
}

/// 创建一个新的空 `.zkv` 库,使用指定 KDF 参数(测试用小参数加速)。
pub fn create_with_params(path: &Path, passphrase: &str, kdf: &KdfParams) -> Result<()> {
    let db = Database::open_in_memory()?;
    let plaintext = db.dump_bytes()?;
    let salt = crypto::gen_salt();
    let key = crypto::derive_key(passphrase.as_bytes(), &salt, kdf)?;
    encrypt_and_write(path, &key, kdf, salt, plaintext)
}

/// 解锁现有 `.zkv` 库(默认 KDF 参数),返回 `Database`。
///
/// 等价于调用 [`unlock_full`](返回 db/key/salt)后丢弃 key/salt。
pub fn unlock(path: &Path, passphrase: &str) -> Result<Database> {
    let (db, _key, _kdf, _salt) = unlock_full(path, passphrase)?;
    Ok(db)
}

/// 解锁现有 `.zkv` 库,返回 `(Database, MasterKey, KdfParams, salt)`。
///
/// 解析文件头 → `derive_key`(用头中的 KDF 参数)→ `decrypt` → `Database::from_bytes`,
/// 并把派生出的 [`MasterKey`]、`header.kdf` 与 `header.salt` 一并返回,供调用方缓存复用,
/// 避免后续 `save` 时重复 Argon2id 派生且回写一致的头参数。
pub fn unlock_full(
    path: &Path,
    passphrase: &str,
) -> Result<(Database, MasterKey, KdfParams, [u8; 16])> {
    let file = fs::read(path)?;
    let header = VaultHeader::parse(&file)?;
    let key = crypto::derive_key(passphrase.as_bytes(), &header.salt, &header.kdf)?;
    let (db, kdf, salt) = decrypt_body(&file, &key)?;
    Ok((db, key, kdf, salt))
}

/// 用**已派生**的 key 解锁 `.zkv` 库(**跳过 Argon2id 派生**)。
///
/// 供 agent 守护进程命中缓存时使用:客户端从 agent 拿到派生密钥后,直接现读文件、
/// 解密、还原 `Database`,省去昂贵的 KDF(百毫秒级)。
///
/// 取值传递 key(因 [`MasterKey`] 非 `Clone`)并原样返回,供调用方继续 `save` 复用。
///
/// **注意**:若文件已被外部用别的口令(新 salt)重加密,AEAD 校验失败 ⇒ [`Error::BadPassphrase`],
/// 调用方应据此回退到 [`unlock_full`] 重新派生(见 `cli::Unlocked::unlock` 的 fail-closed 路径)。
pub fn unlock_with_key(
    path: &Path,
    key: MasterKey,
) -> Result<(Database, MasterKey, KdfParams, [u8; 16])> {
    let file = fs::read(path)?;
    let (db, kdf, salt) = decrypt_body(&file, &key)?;
    Ok((db, key, kdf, salt))
}

/// 内部:用已派生 key 解密「整文件字节」(含 58B 头)→ 还原 `Database` + 头里的 kdf/salt。
///
/// 被 [`unlock_full`](先派生后调用)与 [`unlock_with_key`](直接调用)共用,
/// 确保两者解密路径完全一致。
fn decrypt_body(file: &[u8], key: &MasterKey) -> Result<(Database, KdfParams, [u8; 16])> {
    let header = VaultHeader::parse(file)?;
    let enc = Encrypted {
        nonce: header.nonce,
        ciphertext: file[HEADER_LEN..].to_vec(),
    };
    let plaintext = crypto::decrypt(key, &enc)?;
    let db = Database::from_bytes(&plaintext)?;
    Ok((db, header.kdf, header.salt))
}

/// 解锁现有 `.zkv` 库,使用指定 KDF 参数。
///
/// 实际上 `unlock` 总是读取文件头中的 KDF 参数来派生,故此入口仅用于显式语义/测试对照,
/// 其实现等价于 [`unlock`](同名函数)(params 来自文件头)。
pub fn unlock_with_params(path: &Path, passphrase: &str, _kdf: &KdfParams) -> Result<Database> {
    // 文件头自带 kdf,忽略传入参数,保持与 unlock 一致。
    unlock(path, passphrase)
}

/// 保存(覆盖)`.zkv` 库(默认 KDF 参数)。
pub fn save(path: &Path, passphrase: &str, db: &Database) -> Result<()> {
    save_with_params(path, passphrase, db, &KdfParams::default())
}

/// 保存(覆盖)`.zkv` 库,使用指定 KDF 参数。
///
/// **salt 复用现有文件头中的 salt**(保持同一派生密钥);每次保存生成**新 nonce**。
/// 若文件不存在则等价于新建(用传入 kdf + 新随机 salt)。
pub fn save_with_params(
    path: &Path,
    passphrase: &str,
    db: &Database,
    kdf: &KdfParams,
) -> Result<()> {
    let plaintext = db.dump_bytes()?;

    // 复用现有文件的 salt(若有),保持同一派生密钥。
    let salt = match fs::read(path) {
        Ok(file) => match VaultHeader::parse(&file) {
            Ok(h) => h.salt,
            Err(_) => crypto::gen_salt(),
        },
        Err(_) => crypto::gen_salt(),
    };

    let key = crypto::derive_key(passphrase.as_bytes(), &salt, kdf)?;
    encrypt_and_write(path, &key, kdf, salt, plaintext)
}

/// 保存(覆盖)`.zkv` 库,复用已派生的 [`MasterKey`]、salt 与 KDF 参数。
///
/// 与 [`save_with_params`]( 不同:不重复 Argon2id 派生、也不重读整个文件取 salt,
/// 直接用传入 key 做 AEAD(微秒级)。每次保存仍生成**新 nonce**。
///
/// `kdf` 仅用于回写文件头(下次 `unlock` 读取以派生),派生本身已由调用方完成。
/// 为保持解锁闭环一致性,必须传入与生成该 key 时相同的 KDF 参数。
pub fn save_with_key(
    path: &Path,
    key: &MasterKey,
    kdf: &KdfParams,
    salt: [u8; 16],
    db: &Database,
) -> Result<()> {
    let plaintext = db.dump_bytes()?;
    encrypt_and_write(path, key, kdf, salt, plaintext)
}

/// 修改主口令:验旧口令 → 用新口令 + 新随机 salt 重新加密整库。
///
/// 流程:
/// 1. 用旧口令解锁(`unlock_full`,错则 [`Error::BadPassphrase`]),取 db 与 kdf。
/// 2. 生成新随机 salt,用新口令 + 原 KDF 参数(不改强度)派生新 key。
/// 3. dump 明文 → 用新 key 加密(新 nonce)→ 原子写(新 salt 进文件头)。
///
/// KDF 参数沿用原库的;新 salt 保证换口令后密钥彻底变化,旧口令再无法解锁。
pub fn change_passphrase(path: &Path, old_pass: &str, new_pass: &str) -> Result<()> {
    // 1. 验旧口令 + 取 db 与 kdf。
    let (db, _old_key, kdf, _old_salt) = unlock_full(path, old_pass)?;
    // 2. 新随机 salt + 用新口令派生新 key(沿用原库 kdf 参数)。
    let new_salt = crypto::gen_salt();
    let new_key = crypto::derive_key(new_pass.as_bytes(), &new_salt, &kdf)?;
    // 3. dump 明文 → 用新 key 加密 + 原子写(新 salt 进文件头)。
    let plaintext = db.dump_bytes()?;
    encrypt_and_write(path, &new_key, &kdf, new_salt, plaintext)
}

/// 内部:用**已派生**的 key 加密(新 nonce 由 encrypt 生成,但需要它进文件头)→ 原子写。
/// 不再做任何 Argon2id 派生。
fn encrypt_and_write(
    path: &Path,
    key: &MasterKey,
    kdf: &KdfParams,
    salt: [u8; 16],
    plaintext: Vec<u8>,
) -> Result<()> {
    let enc = crypto::encrypt(key, &plaintext)?;
    let header = VaultHeader {
        version: VERSION,
        flags: 0,
        kdf: *kdf,
        salt,
        nonce: enc.nonce,
    };
    let mut out = header.to_bytes();
    out.extend_from_slice(&enc.ciphertext);
    atomic_write(path, &out)
}

/// 原子写:写到 `<path>.tmp`(0600)→ `fs::rename` → 目标 0600。失败删除临时文件。
fn atomic_write(path: &Path, data: &[u8]) -> Result<()> {
    let mut tmp = path.as_os_str().to_owned();
    tmp.push(".tmp");
    let tmp = PathBuf::from(tmp);

    // 写临时文件
    {
        let res = (|| -> Result<()> {
            let mut f = open_secure(&tmp)?;
            use std::io::Write;
            f.write_all(data)?;
            f.sync_all()?;
            Ok(())
        })();
        if res.is_err() {
            let _ = fs::remove_file(&tmp);
            return res;
        }
    }
    set_mode_0600(&tmp);

    // rename
    if let Err(e) = fs::rename(&tmp, path) {
        let _ = fs::remove_file(&tmp);
        return Err(Error::Io(e));
    }
    // 确保最终文件也是 0600(rename 保留临时文件权限,已设;再兜底)。
    set_mode_0600(path);
    // 持久化 rename 本身:fsync 父目录,避免崩溃时 rename 未落盘导致「丢失更新」。
    // best-effort,忽略不支持目录 fsync 的文件系统错误。
    fsync_parent(path);
    Ok(())
}

#[cfg(unix)]
fn open_secure(path: &Path) -> Result<std::fs::File> {
    use std::os::unix::fs::OpenOptionsExt;
    Ok(std::fs::OpenOptions::new()
        .write(true)
        .create_new(true)
        .truncate(true)
        .mode(0o600)
        .open(path)?)
}

#[cfg(not(unix))]
fn open_secure(path: &Path) -> Result<std::fs::File> {
    Ok(std::fs::OpenOptions::new()
        .write(true)
        .create_new(true)
        .truncate(true)
        .open(path)?)
}

#[cfg(unix)]
fn set_mode_0600(path: &Path) {
    use std::os::unix::fs::PermissionsExt;
    let _ = fs::set_permissions(path, fs::Permissions::from_mode(0o600));
}

#[cfg(not(unix))]
fn set_mode_0600(_path: &Path) {}

/// best-effort fsync 父目录(Unix),把 rename 目录项变更刷到磁盘,避免崩溃丢失更新。
///
/// 标准库 `File::sync_all()` 对一个打开的目录 fd,在 Linux 上会 fsync 该目录
/// (目录项/rename 落盘)。失败忽略:某些文件系统(如 tmpfs/网络 FS)不支持目录 fsync。
#[cfg(unix)]
fn fsync_parent(path: &Path) {
    let Some(parent) = path.parent() else {
        return;
    };
    if let Ok(dir) = std::fs::File::open(parent) {
        let _ = dir.sync_all();
    }
}

#[cfg(not(unix))]
fn fsync_parent(_path: &Path) {}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::Error;

    /// 测试用小 KDF 参数(与 crypto.rs 单测一致),加速派生。
    fn fast_kdf() -> KdfParams {
        KdfParams {
            m_kib: 4_096,
            t_cost: 1,
            p_cost: 1,
        }
    }

    /// 唯一临时文件路径(每个测试独立)。
    fn tmp_path(tag: &str) -> std::path::PathBuf {
        use std::sync::atomic::{AtomicU64, Ordering};
        static C: AtomicU64 = AtomicU64::new(0);
        let n = C.fetch_add(1, Ordering::Relaxed);
        let mut p = std::env::temp_dir();
        p.push(format!("zkv_test_{tag}_{}_{}", std::process::id(), n));
        p
    }

    fn cleanup(p: &std::path::Path) {
        let _ = std::fs::remove_file(p);
        let mut t = p.as_os_str().to_owned();
        t.push(".tmp");
        let _ = std::fs::remove_file(std::path::PathBuf::from(t));
    }

    #[test]
    fn header_roundtrip_and_constants() {
        let kdf = fast_kdf();
        let h = VaultHeader {
            version: VERSION,
            flags: 0,
            kdf,
            salt: [0xA0u8; 16],
            nonce: [0xB1u8; 24],
        };
        let b = h.to_bytes();
        assert_eq!(b.len(), HEADER_LEN);
        assert_eq!(&b[0..4], &MAGIC);
        let h2 = VaultHeader::parse(&b).unwrap();
        assert_eq!(h2.version, VERSION);
        assert_eq!(h2.flags, 0);
        assert_eq!(h2.kdf, kdf);
        assert_eq!(h2.salt, [0xA0u8; 16]);
        assert_eq!(h2.nonce, [0xB1u8; 24]);
    }

    #[test]
    fn header_parse_rejects_bad_magic_and_short() {
        // 短
        assert!(matches!(
            VaultHeader::parse(&[0u8; 10]),
            Err(Error::CorruptFile(_))
        ));
        // 坏魔数
        let mut bad = vec![0u8; HEADER_LEN];
        bad[0..4].copy_from_slice(b"XXXX");
        assert!(matches!(
            VaultHeader::parse(&bad),
            Err(Error::CorruptFile(_))
        ));
    }

    #[test]
    fn create_unlock_roundtrip() {
        let p = tmp_path("cu");
        cleanup(&p);
        let kdf = fast_kdf();
        create_with_params(&p, "hunter2", &kdf).expect("create");
        let db = unlock(&p, "hunter2").expect("unlock");
        // 默认库应能查询 items 表(空)
        let cnt: i64 = db
            .conn()
            .query_row("SELECT COUNT(*) FROM items", [], |r| r.get::<_, i64>(0))
            .unwrap();
        assert_eq!(cnt, 0);
        cleanup(&p);
    }

    #[test]
    fn save_then_unlock_preserves_changes() {
        let p = tmp_path("save");
        cleanup(&p);
        let kdf = fast_kdf();
        create_with_params(&p, "pw", &kdf).unwrap();

        // 解锁 → 改 → 保存 → 再解锁
        {
            let db = unlock(&p, "pw").unwrap();
            db.conn()
                .execute(
                    "INSERT INTO items(type,title,data,search_text,created_at,updated_at)
                     VALUES ('note','T','{}','s',1,1)",
                    [],
                )
                .unwrap();
            save_with_params(&p, "pw", &db, &kdf).unwrap();
        }
        let db2 = unlock(&p, "pw").unwrap();
        let cnt: i64 = db2
            .conn()
            .query_row("SELECT COUNT(*) FROM items", [], |r| r.get::<_, i64>(0))
            .unwrap();
        assert_eq!(cnt, 1);
        cleanup(&p);
    }

    #[test]
    fn wrong_passphrase_yields_bad_passphrase() {
        let p = tmp_path("wrong");
        cleanup(&p);
        let kdf = fast_kdf();
        create_with_params(&p, "correct", &kdf).unwrap();
        let res = unlock(&p, "incorrect");
        assert!(
            matches!(res, Err(Error::BadPassphrase)),
            "expected BadPassphrase, got {:?}",
            res
        );
        cleanup(&p);
    }

    #[test]
    fn corrupted_magic_yields_corrupt_file() {
        let p = tmp_path("corrupt");
        cleanup(&p);
        let kdf = fast_kdf();
        create_with_params(&p, "pw", &kdf).unwrap();

        // 篡改魔数
        let mut data = std::fs::read(&p).unwrap();
        data[0] = 0x00;
        std::fs::write(&p, &data).unwrap();

        let res = unlock(&p, "pw");
        assert!(matches!(res, Err(Error::CorruptFile(_))));
        cleanup(&p);
    }

    #[test]
    fn save_each_time_new_nonce() {
        // 两次保存产生不同的 nonce(头中的 nonce 字段应不同)。
        let p = tmp_path("nonce");
        cleanup(&p);
        let kdf = fast_kdf();
        create_with_params(&p, "pw", &kdf).unwrap();
        let f1 = std::fs::read(&p).unwrap();
        let h1 = VaultHeader::parse(&f1).unwrap();

        {
            let db = unlock(&p, "pw").unwrap();
            save_with_params(&p, "pw", &db, &kdf).unwrap();
        }
        let f2 = std::fs::read(&p).unwrap();
        let h2 = VaultHeader::parse(&f2).unwrap();

        assert_ne!(h1.nonce, h2.nonce, "每次保存应生成新 nonce");
        // salt 保持不变
        assert_eq!(h1.salt, h2.salt);
        cleanup(&p);
    }

    #[test]
    fn change_passphrase_roundtrip() {
        // create(old) → 插条 item → change(old,new) →
        // unlock(old) Err(BadPassphrase)、unlock(new) Ok 且数据仍在、salt 变。
        let p = tmp_path("change");
        cleanup(&p);
        let kdf = fast_kdf();
        create_with_params(&p, "old", &kdf).unwrap();

        // 插一条 item。
        {
            let db = unlock(&p, "old").unwrap();
            db.conn()
                .execute(
                    "INSERT INTO items(type,title,data,search_text,created_at,updated_at)
                     VALUES ('password','GitHub','{\"password\":\"s3cret\"}','s3cret',1,1)",
                    [],
                )
                .unwrap();
            save_with_params(&p, "old", &db, &kdf).unwrap();
        }
        let before = std::fs::read(&p).unwrap();
        let h_before = VaultHeader::parse(&before).unwrap();

        // 改口令。
        change_passphrase(&p, "old", "new").unwrap();

        // 旧口令失败。
        assert!(
            matches!(unlock(&p, "old"), Err(Error::BadPassphrase)),
            "old passphrase must no longer unlock after change"
        );
        // 新口令成功且数据仍在。
        let db = unlock(&p, "new").unwrap();
        let cnt: i64 = db
            .conn()
            .query_row("SELECT COUNT(*) FROM items", [], |r| r.get::<_, i64>(0))
            .unwrap();
        assert_eq!(cnt, 1, "data should survive passphrase change");
        let title: String = db
            .conn()
            .query_row("SELECT title FROM items", [], |r| r.get::<_, String>(0))
            .unwrap();
        assert_eq!(title, "GitHub");

        // salt 变了。
        let after = std::fs::read(&p).unwrap();
        let h_after = VaultHeader::parse(&after).unwrap();
        assert_ne!(h_before.salt, h_after.salt, "salt must rotate on change");
        cleanup(&p);
    }

    #[test]
    fn change_passphrase_wrong_old_yields_bad_passphrase() {
        let p = tmp_path("change_wrong");
        cleanup(&p);
        let kdf = fast_kdf();
        create_with_params(&p, "correct", &kdf).unwrap();
        let res = change_passphrase(&p, "incorrect", "whatever");
        assert!(
            matches!(res, Err(Error::BadPassphrase)),
            "wrong old passphrase must yield BadPassphrase, got {:?}",
            res
        );
        // 库未被改写(原口令仍可用)。
        assert!(unlock(&p, "correct").is_ok());
        cleanup(&p);
    }
}