quicproquo/crates/quicproquo-core/src/keystore.rs

use std::{
    collections::HashMap,
    fs,
    path::{Path, PathBuf},
    sync::RwLock,
};

use openmls_traits::key_store::{MlsEntity, OpenMlsKeyStore};

/// A disk-backed key store implementing `OpenMlsKeyStore`.
///
/// In-memory when `path` is `None`; otherwise flushes the entire map to disk on
/// every store/delete so HPKE init keys survive process restarts.
///
/// # Serialization
///
/// Uses bincode for both individual MLS entity values and the outer HashMap
/// container. This is required because OpenMLS types use bincode-compatible
/// serialization, and `HashMap<Vec<u8>, Vec<u8>>` requires a binary format
/// (JSON mandates string keys).
///
/// # Persistence security
///
/// When `path` is set, file permissions are restricted to owner-only (0o600)
/// on Unix platforms, since the store may contain HPKE private keys.
#[derive(Debug)]
pub struct DiskKeyStore {
    path: Option<PathBuf>,
    values: RwLock<HashMap<Vec<u8>, Vec<u8>>>,
}

#[derive(thiserror::Error, Debug, PartialEq, Eq)]
pub enum DiskKeyStoreError {
    #[error("serialization error")]
    Serialization,
    #[error("io error: {0}")]
    Io(String),
}

impl DiskKeyStore {
    /// In-memory keystore (no persistence).
    pub fn ephemeral() -> Self {
        Self {
            path: None,
            values: RwLock::new(HashMap::new()),
        }
    }

    /// Persistent keystore backed by `path`. Creates an empty store if missing.
    pub fn persistent(path: impl AsRef<Path>) -> Result<Self, DiskKeyStoreError> {
        let path = path.as_ref().to_path_buf();
        let values = if path.exists() {
            let bytes = fs::read(&path).map_err(|e| DiskKeyStoreError::Io(e.to_string()))?;
            if bytes.is_empty() {
                HashMap::new()
            } else {
                bincode::deserialize(&bytes)
                    .map_err(|_| DiskKeyStoreError::Serialization)?
            }
        } else {
            HashMap::new()
        };

        let store = Self {
            path: Some(path),
            values: RwLock::new(values),
        };

        // Set restrictive file permissions on the keystore file.
        store.set_file_permissions()?;

        Ok(store)
    }

    fn flush(&self) -> Result<(), DiskKeyStoreError> {
        let Some(path) = &self.path else {
            return Ok(());
        };
        let values = self.values.read().map_err(|_| DiskKeyStoreError::Io("lock poisoned".into()))?;
        let bytes = bincode::serialize(&*values).map_err(|_| DiskKeyStoreError::Serialization)?;
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).map_err(|e| DiskKeyStoreError::Io(e.to_string()))?;
        }
        fs::write(path, &bytes).map_err(|e| DiskKeyStoreError::Io(e.to_string()))?;
        self.set_file_permissions()?;
        Ok(())
    }

    /// Restrict file permissions to owner-only (0o600) on Unix.
    #[cfg(unix)]
    fn set_file_permissions(&self) -> Result<(), DiskKeyStoreError> {
        use std::os::unix::fs::PermissionsExt;
        if let Some(path) = &self.path {
            if path.exists() {
                let perms = std::fs::Permissions::from_mode(0o600);
                fs::set_permissions(path, perms)
                    .map_err(|e| DiskKeyStoreError::Io(format!("set permissions: {e}")))?;
            }
        }
        Ok(())
    }

    #[cfg(not(unix))]
    fn set_file_permissions(&self) -> Result<(), DiskKeyStoreError> {
        Ok(())
    }
}

impl Default for DiskKeyStore {
    fn default() -> Self {
        Self::ephemeral()
    }
}

impl OpenMlsKeyStore for DiskKeyStore {
    type Error = DiskKeyStoreError;

    fn store<V: MlsEntity>(&self, k: &[u8], v: &V) -> Result<(), Self::Error> {
        let value = bincode::serialize(v).map_err(|_| DiskKeyStoreError::Serialization)?;
        let mut values = self.values.write().map_err(|_| DiskKeyStoreError::Io("lock poisoned".into()))?;
        values.insert(k.to_vec(), value);
        drop(values);
        self.flush()
    }

    fn read<V: MlsEntity>(&self, k: &[u8]) -> Option<V> {
        let values = match self.values.read() {
            Ok(v) => v,
            Err(_) => return None,
        };
        values
            .get(k)
            .and_then(|bytes| bincode::deserialize(bytes).ok())
    }

    fn delete<V: MlsEntity>(&self, k: &[u8]) -> Result<(), Self::Error> {
        let mut values = self.values.write().map_err(|_| DiskKeyStoreError::Io("lock poisoned".into()))?;
        values.remove(k);
        drop(values);
        self.flush()
    }
}