quicproquo/crates/quicproquo-client/src/client/state.rs

use std::path::{Path, PathBuf};
use std::sync::Arc;

use anyhow::Context;
use argon2::{Algorithm, Argon2, Params, Version};
use chacha20poly1305::{
    aead::{Aead, KeyInit},
    ChaCha20Poly1305, Key, Nonce,
};
use rand::RngCore;
use serde::{Deserialize, Serialize};

use quicproquo_core::{DiskKeyStore, GroupMember, HybridKeypair, HybridKeypairBytes, IdentityKeypair};

/// Magic bytes for encrypted client state files.
const STATE_MAGIC: &[u8; 4] = b"QPCE";
const STATE_SALT_LEN: usize = 16;
const STATE_NONCE_LEN: usize = 12;

#[derive(Serialize, Deserialize)]
pub struct StoredState {
    pub identity_seed: [u8; 32],
    pub group: Option<Vec<u8>>,
    /// Post-quantum hybrid keypair (X25519 + ML-KEM-768). `None` for state created before hybrid was added.
    #[serde(default)]
    pub hybrid_key: Option<HybridKeypairBytes>,
    /// Cached member public keys for group participants.
    #[serde(default)]
    pub member_keys: Vec<Vec<u8>>,
}

impl StoredState {
    pub fn into_parts(self, state_path: &Path) -> anyhow::Result<(GroupMember, Option<HybridKeypair>)> {
        let identity = Arc::new(IdentityKeypair::from_seed(self.identity_seed));
        let group = self
            .group
            .map(|bytes| bincode::deserialize(&bytes).context("decode group"))
            .transpose()?;
        let key_store = DiskKeyStore::persistent(keystore_path(state_path))?;
        let hybrid = self.hybrid_key.is_some();
        let member = GroupMember::new_with_state(identity, key_store, group, hybrid);

        let hybrid_kp = self
            .hybrid_key
            .map(|bytes| HybridKeypair::from_bytes(&bytes).context("decode hybrid key"))
            .transpose()?;

        Ok((member, hybrid_kp))
    }

    pub fn from_parts(member: &GroupMember, hybrid_kp: Option<&HybridKeypair>) -> anyhow::Result<Self> {
        let group = member
            .group_ref()
            .map(|g| bincode::serialize(g).context("serialize group"))
            .transpose()?;

        Ok(Self {
            identity_seed: member.identity_seed(),
            group,
            hybrid_key: hybrid_kp.map(|kp| kp.to_bytes()),
            member_keys: Vec::new(),
        })
    }
}

/// Argon2id parameters for client state key derivation (auditable; matches argon2 crate defaults).
/// - Memory: 19 MiB (m_cost = 19*1024 KiB)
/// - Time: 2 iterations
/// - Parallelism: 1 lane
const ARGON2_STATE_M_COST: u32 = 19 * 1024;
const ARGON2_STATE_T_COST: u32 = 2;
const ARGON2_STATE_P_COST: u32 = 1;

/// Derive a 32-byte key from a password and salt using Argon2id with explicit parameters.
fn derive_state_key(password: &str, salt: &[u8]) -> anyhow::Result<[u8; 32]> {
    let params = Params::new(ARGON2_STATE_M_COST, ARGON2_STATE_T_COST, ARGON2_STATE_P_COST, Some(32))
        .map_err(|e| anyhow::anyhow!("argon2 params: {e}"))?;
    let argon2 = Argon2::new(Algorithm::Argon2id, Version::default(), params);
    let mut key = [0u8; 32];
    argon2
        .hash_password_into(password.as_bytes(), salt, &mut key)
        .map_err(|e| anyhow::anyhow!("argon2 key derivation failed: {e}"))?;
    Ok(key)
}

/// Encrypt `plaintext` with the QPCE format: magic(4) | salt(16) | nonce(12) | ciphertext.
pub fn encrypt_state(password: &str, plaintext: &[u8]) -> anyhow::Result<Vec<u8>> {
    let mut salt = [0u8; STATE_SALT_LEN];
    rand::rngs::OsRng.fill_bytes(&mut salt);

    let mut nonce_bytes = [0u8; STATE_NONCE_LEN];
    rand::rngs::OsRng.fill_bytes(&mut nonce_bytes);

    let key = zeroize::Zeroizing::new(derive_state_key(password, &salt)?);
    let cipher = ChaCha20Poly1305::new(Key::from_slice(&*key));
    let nonce = Nonce::from_slice(&nonce_bytes);

    let ciphertext = cipher
        .encrypt(nonce, plaintext)
        .map_err(|e| anyhow::anyhow!("state encryption failed: {e}"))?;

    let mut out = Vec::with_capacity(4 + STATE_SALT_LEN + STATE_NONCE_LEN + ciphertext.len());
    out.extend_from_slice(STATE_MAGIC);
    out.extend_from_slice(&salt);
    out.extend_from_slice(&nonce_bytes);
    out.extend_from_slice(&ciphertext);
    Ok(out)
}

/// Decrypt a QPCE-formatted state file.
pub fn decrypt_state(password: &str, data: &[u8]) -> anyhow::Result<Vec<u8>> {
    let header_len = 4 + STATE_SALT_LEN + STATE_NONCE_LEN;
    anyhow::ensure!(
        data.len() > header_len,
        "encrypted state file too short ({} bytes)",
        data.len()
    );

    let salt = &data[4..4 + STATE_SALT_LEN];
    let nonce_bytes = &data[4 + STATE_SALT_LEN..header_len];
    let ciphertext = &data[header_len..];

    let key = zeroize::Zeroizing::new(derive_state_key(password, salt)?);
    let cipher = ChaCha20Poly1305::new(Key::from_slice(&*key));
    let nonce = Nonce::from_slice(nonce_bytes);

    let plaintext = cipher
        .decrypt(nonce, ciphertext)
        .map_err(|_| anyhow::anyhow!("state decryption failed (wrong password?)"))?;

    Ok(plaintext)
}

/// Returns true if raw bytes begin with the QPCE magic header.
pub fn is_encrypted_state(bytes: &[u8]) -> bool {
    bytes.len() >= 4 && &bytes[..4] == STATE_MAGIC
}

pub fn load_or_init_state(path: &Path, password: Option<&str>) -> anyhow::Result<StoredState> {
    if path.exists() {
        let mut state = load_existing_state(path, password)?;
        // Generate hybrid keypair if missing (upgrade from older state).
        if state.hybrid_key.is_none() {
            state.hybrid_key = Some(HybridKeypair::generate().to_bytes());
            write_state(path, &state, password)?;
        }
        return Ok(state);
    }

    let identity = IdentityKeypair::generate();
    let hybrid_kp = HybridKeypair::generate();
    let key_store = DiskKeyStore::persistent(keystore_path(path))?;
    let member = GroupMember::new_with_state(Arc::new(identity), key_store, None, false);
    let state = StoredState::from_parts(&member, Some(&hybrid_kp))?;
    write_state(path, &state, password)?;
    Ok(state)
}

pub fn load_existing_state(path: &Path, password: Option<&str>) -> anyhow::Result<StoredState> {
    let bytes = std::fs::read(path).with_context(|| format!("read state file {path:?}"))?;

    if is_encrypted_state(&bytes) {
        let pw = password
            .context("state file is encrypted (QPCE); a password is required to decrypt it")?;
        let plaintext = decrypt_state(pw, &bytes)?;
        bincode::deserialize(&plaintext).context("decode encrypted state")
    } else {
        bincode::deserialize(&bytes).context("decode state")
    }
}

pub fn save_state(
    path: &Path,
    member: &GroupMember,
    hybrid_kp: Option<&HybridKeypair>,
    password: Option<&str>,
) -> anyhow::Result<()> {
    let state = StoredState::from_parts(member, hybrid_kp)?;
    write_state(path, &state, password)
}

pub fn write_state(path: &Path, state: &StoredState, password: Option<&str>) -> anyhow::Result<()> {
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent).with_context(|| format!("create dir {parent:?}"))?;
    }
    let plaintext = bincode::serialize(state).context("encode state")?;

    let bytes = if let Some(pw) = password {
        encrypt_state(pw, &plaintext)?
    } else {
        plaintext
    };

    let tmp = path.with_extension("tmp");
    std::fs::write(&tmp, bytes).with_context(|| format!("write state temp {tmp:?}"))?;
    std::fs::rename(&tmp, path).with_context(|| format!("rename state {tmp:?} -> {path:?}"))?;
    Ok(())
}

pub fn decode_identity_key(hex_str: &str) -> anyhow::Result<Vec<u8>> {
    let bytes = super::hex::decode(hex_str)
        .map_err(|e| anyhow::anyhow!(e))
        .context("identity key must be hex")?;
    anyhow::ensure!(bytes.len() == 32, "identity key must be 32 bytes");
    Ok(bytes)
}

pub fn keystore_path(state_path: &Path) -> PathBuf {
    let mut path = state_path.to_path_buf();
    path.set_extension("ks");
    path
}

pub fn sha256(bytes: &[u8]) -> Vec<u8> {
    use sha2::{Digest, Sha256};
    Sha256::digest(bytes).to_vec()
}

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

    #[test]
    fn encrypt_decrypt_roundtrip() {
        let plaintext = b"test state data";
        let password = "test-password";
        let encrypted = encrypt_state(password, plaintext).unwrap();
        assert!(is_encrypted_state(&encrypted));
        let decrypted = decrypt_state(password, &encrypted).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn wrong_password_fails() {
        let plaintext = b"test state data";
        let encrypted = encrypt_state("correct", plaintext).unwrap();
        assert!(decrypt_state("wrong", &encrypted).is_err());
    }

    #[test]
    fn state_encrypt_decrypt_round_trip() {
        let state = StoredState {
            identity_seed: [42u8; 32],
            hybrid_key: None,
            group: None,
            member_keys: Vec::new(),
        };
        let password = "test-password";
        let plaintext = bincode::serialize(&state).unwrap();
        let encrypted = encrypt_state(password, &plaintext).unwrap();
        let decrypted = decrypt_state(password, &encrypted).unwrap();
        let recovered: StoredState = bincode::deserialize(&decrypted).unwrap();
        assert_eq!(recovered.identity_seed, state.identity_seed);
        assert!(recovered.hybrid_key.is_none());
        assert!(recovered.group.is_none());
    }

    #[test]
    fn state_encrypt_decrypt_with_hybrid_key() {
        use zeroize::Zeroizing;
        let state = StoredState {
            identity_seed: [7u8; 32],
            hybrid_key: Some(HybridKeypairBytes {
                x25519_sk: Zeroizing::new([1u8; 32]),
                mlkem_dk: Zeroizing::new(vec![3u8; 2400]),
                mlkem_ek: vec![4u8; 1184],
            }),
            group: None,
            member_keys: Vec::new(),
        };
        let password = "another-password";
        let plaintext = bincode::serialize(&state).unwrap();
        let encrypted = encrypt_state(password, &plaintext).unwrap();
        let decrypted = decrypt_state(password, &encrypted).unwrap();
        let recovered: StoredState = bincode::deserialize(&decrypted).unwrap();
        assert_eq!(recovered.identity_seed, state.identity_seed);
        assert!(recovered.hybrid_key.is_some());
    }

    #[test]
    fn state_wrong_password_fails() {
        let state = StoredState {
            identity_seed: [99u8; 32],
            hybrid_key: None,
            group: None,
            member_keys: Vec::new(),
        };
        let plaintext = bincode::serialize(&state).unwrap();
        let encrypted = encrypt_state("correct", &plaintext).unwrap();
        assert!(decrypt_state("wrong", &encrypted).is_err());
    }
}