quicproquo/crates/quicprochat-sdk/src/state.rs

//! Encrypted client state — identity key, MLS group, and hybrid KEM persistence.
//!
//! State is stored with the QPCE format:
//! `QPCE` magic (4 bytes) + salt (16 bytes) + nonce (12 bytes) + ciphertext.
//! Key derivation uses Argon2id (m=19456 KiB, t=2, p=1).

use std::path::Path;

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

use crate::error::SdkError;

/// 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;

/// Argon2id parameters for client state key derivation.
const ARGON2_STATE_M_COST: u32 = 19 * 1024;
const ARGON2_STATE_T_COST: u32 = 2;
const ARGON2_STATE_P_COST: u32 = 1;

/// Encrypted client state (identity + MLS group + hybrid keys).
#[derive(Serialize, Deserialize)]
pub struct StoredState {
    pub identity_seed: [u8; 32],
    pub group: Option<Vec<u8>>,
    pub hybrid_key: Option<Vec<u8>>,
    #[serde(default)]
    pub member_keys: Vec<Vec<u8>>,
}

/// Derive a 32-byte key from a password and salt using Argon2id.
fn derive_state_key(password: &str, salt: &[u8]) -> Result<Zeroizing<[u8; 32]>, SdkError> {
    let params = Params::new(ARGON2_STATE_M_COST, ARGON2_STATE_T_COST, ARGON2_STATE_P_COST, Some(32))
        .map_err(|e| SdkError::Crypto(format!("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| SdkError::Crypto(format!("argon2 key derivation failed: {e}")))?;
    Ok(Zeroizing::new(key))
}

/// Encrypt `plaintext` with the QPCE format: magic(4) | salt(16) | nonce(12) | ciphertext.
pub fn encrypt_state(password: &str, plaintext: &[u8]) -> Result<Vec<u8>, SdkError> {
    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 = 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| SdkError::Crypto(format!("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]) -> Result<Vec<u8>, SdkError> {
    let header_len = 4 + STATE_SALT_LEN + STATE_NONCE_LEN;
    if data.len() <= header_len {
        return Err(SdkError::Crypto(format!(
            "encrypted state file too short ({} bytes)",
            data.len()
        )));
    }

    if &data[..4] != STATE_MAGIC {
        return Err(SdkError::Crypto("invalid state file magic header".into()));
    }

    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 = derive_state_key(password, salt)?;
    let cipher = ChaCha20Poly1305::new(Key::from_slice(&*key));
    let nonce = Nonce::from_slice(nonce_bytes);

    cipher
        .decrypt(nonce, ciphertext)
        .map_err(|_| SdkError::Crypto("state decryption failed (wrong password?)".into()))
}

/// Save state to disk, optionally encrypted with `password`.
pub fn save_state(path: &Path, state: &StoredState, password: Option<&str>) -> Result<(), SdkError> {
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)
            .map_err(|e| SdkError::Storage(format!("create dir {}: {e}", parent.display())))?;
    }

    let plaintext = bincode::serialize(state)
        .map_err(|e| SdkError::Storage(format!("encode state: {e}")))?;

    let bytes = match password {
        Some(pw) => encrypt_state(pw, &plaintext)?,
        None => plaintext,
    };

    let tmp = path.with_extension("tmp");
    std::fs::write(&tmp, &bytes)
        .map_err(|e| SdkError::Storage(format!("write state temp {}: {e}", tmp.display())))?;
    std::fs::rename(&tmp, path)
        .map_err(|e| SdkError::Storage(format!("rename state: {e}")))?;
    Ok(())
}

/// Load state from disk, decrypting if necessary.
pub fn load_state(path: &Path, password: Option<&str>) -> Result<StoredState, SdkError> {
    let bytes = std::fs::read(path)
        .map_err(|e| SdkError::Storage(format!("read state file {}: {e}", path.display())))?;

    let is_encrypted = bytes.len() >= 4 && &bytes[..4] == STATE_MAGIC;

    if is_encrypted {
        let pw = password
            .ok_or_else(|| SdkError::Crypto("state file is encrypted; password required".into()))?;
        let plaintext = decrypt_state(pw, &bytes)?;
        bincode::deserialize(&plaintext)
            .map_err(|e| SdkError::Storage(format!("decode encrypted state: {e}")))
    } else {
        bincode::deserialize(&bytes)
            .map_err(|e| SdkError::Storage(format!("decode state: {e}")))
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
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_eq!(&encrypted[..4], STATE_MAGIC);
        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_serialize_roundtrip() {
        let state = StoredState {
            identity_seed: [42u8; 32],
            group: None,
            hybrid_key: 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.group.is_none());
        assert!(recovered.hybrid_key.is_none());
    }

    #[test]
    fn save_load_roundtrip() {
        let dir = std::env::temp_dir().join("qpc_sdk_state_test");
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("test.state");

        let state = StoredState {
            identity_seed: [7u8; 32],
            group: Some(vec![1, 2, 3]),
            hybrid_key: Some(vec![4, 5, 6]),
            member_keys: vec![vec![10, 11]],
        };

        // Test unencrypted.
        save_state(&path, &state, None).unwrap();
        let loaded = load_state(&path, None).unwrap();
        assert_eq!(loaded.identity_seed, state.identity_seed);
        assert_eq!(loaded.group, state.group);

        // Test encrypted.
        save_state(&path, &state, Some("pw")).unwrap();
        let loaded = load_state(&path, Some("pw")).unwrap();
        assert_eq!(loaded.identity_seed, state.identity_seed);
        assert_eq!(loaded.hybrid_key, state.hybrid_key);

        // Encrypted but no password fails.
        assert!(load_state(&path, None).is_err());

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn too_short_data_fails() {
        assert!(decrypt_state("pw", b"QPCE").is_err());
        assert!(decrypt_state("pw", &[]).is_err());
    }

    #[test]
    fn invalid_magic_fails() {
        let mut data = vec![0u8; 100];
        data[..4].copy_from_slice(b"NOPE");
        assert!(decrypt_state("pw", &data).is_err());
    }
}