ping-openmls-sdk-mls-store 0.6.13

//! `PersistentMlsProvider` — the OpenMLS provider the SDK actually uses.

use std::sync::Arc;

use openmls_rust_crypto::{MemoryStorage, RustCrypto};
use openmls_traits::OpenMlsProvider;
use serde::{Deserialize, Serialize};

use crate::{backend::StorageBackend, Error, Result};

#[cfg(not(target_arch = "wasm32"))]
use crate::sqlite::SqliteFile;

use crate::AsyncBlobStore;

/// Versioned wire-format envelope for the persisted MLS state. Bumped only on
/// incompatible serialization changes; today's `v = 1` is "raw `MemoryStorage` HashMap
/// pairs."
const PERSIST_VERSION: u8 = 1;

#[derive(Serialize, Deserialize)]
struct PersistedBlob {
    v: u8,
    pairs: Vec<(serde_bytes::ByteBuf, serde_bytes::ByteBuf)>,
}

/// The SDK's OpenMLS provider. Wraps the in-memory `MemoryStorage` so OpenMLS' working
/// set stays fast, and adds a persistent `checkpoint()` that flushes the working set
/// to the configured [`StorageBackend`].
///
/// `Arc<PersistentMlsProvider>` is what `MessagingClient` holds and passes into
/// OpenMLS APIs. All methods are safe to call concurrently — the inner `MemoryStorage`
/// guards itself with an `RwLock`, and SQLite ops are serialised by a separate `Mutex`.
pub struct PersistentMlsProvider {
    crypto: RustCrypto,
    storage: MemoryStorage,
    backend: PersistentBackend,
}

impl std::fmt::Debug for PersistentMlsProvider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PersistentMlsProvider")
            .field("backend", &self.backend.kind())
            .finish()
    }
}

enum PersistentBackend {
    Memory,
    #[cfg(not(target_arch = "wasm32"))]
    Sqlite(SqliteFile),
    /// Host-backed async blob storage (universal). The provider keeps the
    /// `MemoryStorage` working set in memory and round-trips a CBOR
    /// snapshot of the whole map through the host-supplied trait via
    /// `checkpoint_async` / `open_async`.
    AsyncBlob(Arc<dyn AsyncBlobStore>),
}

impl PersistentBackend {
    fn kind(&self) -> &'static str {
        match self {
            Self::Memory => "memory",
            #[cfg(not(target_arch = "wasm32"))]
            Self::Sqlite(_) => "sqlite",
            Self::AsyncBlob(_) => "async_blob",
        }
    }
}

impl PersistentMlsProvider {
    /// Open (or create) a provider against `backend`. On first use against a
    /// previously-checkpointed file, the saved working set is loaded back into
    /// memory; otherwise the `MemoryStorage` starts empty.
    ///
    /// Synchronous path — supports [`StorageBackend::Memory`] and
    /// [`StorageBackend::Sqlite`]. The WASM-only
    /// [`StorageBackend::IndexedDb`] variant cannot be opened synchronously
    /// because IndexedDB is async-only; callers on WASM must use
    /// [`open_async`](Self::open_async) instead.
    pub fn open(backend: StorageBackend) -> Result<Arc<Self>> {
        let crypto = RustCrypto::default();
        let storage = MemoryStorage::default();
        let backend = match backend {
            StorageBackend::Memory => PersistentBackend::Memory,
            #[cfg(not(target_arch = "wasm32"))]
            StorageBackend::Sqlite {
                path,
                encryption_key,
            } => {
                let file = SqliteFile::open(&path, encryption_key.as_ref())?;
                if let Some(bytes) = file.read_blob()? {
                    load_into_memstore(&storage, &bytes)?;
                }
                PersistentBackend::Sqlite(file)
            }
            StorageBackend::AsyncBlob { blob_store: _ } => {
                return Err(Error::Io(
                    "StorageBackend::AsyncBlob requires PersistentMlsProvider::open_async — \
                     the sync open() path can't drive async I/O"
                        .into(),
                ));
            }
        };
        Ok(Arc::new(Self {
            crypto,
            storage,
            backend,
        }))
    }

    /// Async open path. On WASM the [`StorageBackend::IndexedDb`] variant
    /// asks the host to read its snapshot blob and loads it into
    /// `MemoryStorage` before returning — so the very first call to
    /// `MlsGroup::load(crypto.storage(), ...)` after re-init finds the
    /// group state intact. Non-WASM targets delegate to the sync path so
    /// hosts can use a single code path.
    pub async fn open_async(backend: StorageBackend) -> Result<Arc<Self>> {
        // Fast-path: the non-async-blob variants don't need any async
        // work, so forward them to the sync `open`. Only `AsyncBlob`
        // requires the async read.
        if let StorageBackend::AsyncBlob { blob_store } = backend {
            let crypto = RustCrypto::default();
            let storage = MemoryStorage::default();
            let bytes = blob_store.read_blob().await.map_err(Error::Io)?;
            if let Some(bytes) = bytes {
                load_into_memstore(&storage, &bytes)?;
            }
            return Ok(Arc::new(Self {
                crypto,
                storage,
                backend: PersistentBackend::AsyncBlob(blob_store),
            }));
        }
        Self::open(backend)
    }

    /// Flush the in-memory MLS state to the configured backend. No-op for
    /// `StorageBackend::Memory`. Called by `MessagingClient` after every state-changing
    /// op so cold-start re-init picks up where we left off.
    ///
    /// Synchronous variant — works for Memory + Sqlite. WASM IndexedDb
    /// callers must use [`checkpoint_async`](Self::checkpoint_async).
    pub fn checkpoint(&self) -> Result<()> {
        match &self.backend {
            PersistentBackend::Memory => Ok(()),
            #[cfg(not(target_arch = "wasm32"))]
            PersistentBackend::Sqlite(file) => {
                let bytes = serialize_memstore(&self.storage)?;
                file.write_blob(&bytes)
            }
            PersistentBackend::AsyncBlob(_) => Err(Error::Io(
                "PersistentBackend::AsyncBlob requires checkpoint_async — \
                 the sync checkpoint() path can't drive async I/O"
                    .into(),
            )),
        }
    }

    /// Async flush. Hosts using `StorageBackend::AsyncBlob` MUST call
    /// this (rather than the sync `checkpoint`) after every
    /// state-changing op. The sync-backed variants (Memory, Sqlite)
    /// delegate to the sync path under the hood; the `.await` is a
    /// no-op for them.
    pub async fn checkpoint_async(&self) -> Result<()> {
        match &self.backend {
            PersistentBackend::Memory => Ok(()),
            #[cfg(not(target_arch = "wasm32"))]
            PersistentBackend::Sqlite(file) => {
                let bytes = serialize_memstore(&self.storage)?;
                file.write_blob(&bytes)
            }
            PersistentBackend::AsyncBlob(blob_store) => {
                let bytes = serialize_memstore(&self.storage)?;
                blob_store.write_blob(bytes).await.map_err(Error::Io)
            }
        }
    }

    /// Cheap inspection helper for tests — returns the current number of (key, value)
    /// pairs in the working set. Not exposed in production code paths.
    #[cfg(any(test, feature = "test-utils"))]
    pub fn memstore_len(&self) -> usize {
        self.storage.values.read().unwrap().len()
    }

    /// [CR-7] Enumerate every (raw_storage_key, raw_storage_value) pair in the working
    /// set whose key references the given MLS group id.
    ///
    /// OpenMLS' `MemoryStorage` keys are formatted as `label || serde_json(natural_key)
    /// || u16(version)`. For group-scoped entries the natural key is the `GroupId`,
    /// which is JSON-serialised as a number-array. We match by substring against that
    /// JSON form. False positives are vanishingly unlikely — a 16-byte group id's JSON
    /// encoding is ~50 bytes of array literal, which won't collide with leaf-keyed
    /// material like `EncryptionKeyPair` (32-byte hex blob) or `SignatureKeyPair`.
    ///
    /// Returns owned copies of both halves so the caller can hold them past the read
    /// guard's scope. This is what `Conversation::export_state_snapshot` consumes.
    pub fn group_scoped_entries(&self, group_id_bytes: &[u8]) -> Vec<(Vec<u8>, Vec<u8>)> {
        // Pre-compute the JSON-encoded form of the GroupId — that's how OpenMLS lays
        // it out inside storage keys.
        let needle = serde_json::to_vec(&group_id_bytes).unwrap_or_default();
        let guard = self.storage.values.read().unwrap();
        guard
            .iter()
            .filter(|(k, _)| contains_subsequence(k, &needle))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect()
    }

    /// [CR-7] Bulk-insert key/value pairs into the working set. Used by
    /// `Conversation::import_state_snapshot` to replay an exporter's snapshot into a
    /// fresh provider before calling `MlsGroup::load`.
    ///
    /// Existing entries with the same key are overwritten — that's what we want when
    /// applying a more-recent snapshot. The caller is responsible for triggering a
    /// [`checkpoint`](Self::checkpoint) afterwards so the imported state hits disk.
    pub fn import_entries(&self, entries: Vec<(Vec<u8>, Vec<u8>)>) -> Result<()> {
        let mut guard = self
            .storage
            .values
            .write()
            .map_err(|e| Error::Io(format!("memstore lock: {e}")))?;
        for (k, v) in entries {
            guard.insert(k, v);
        }
        Ok(())
    }
}

fn contains_subsequence(haystack: &[u8], needle: &[u8]) -> bool {
    if needle.is_empty() || haystack.len() < needle.len() {
        return false;
    }
    haystack
        .windows(needle.len())
        .any(|window| window == needle)
}

impl OpenMlsProvider for PersistentMlsProvider {
    type CryptoProvider = RustCrypto;
    type RandProvider = RustCrypto;
    type StorageProvider = MemoryStorage;

    fn crypto(&self) -> &Self::CryptoProvider {
        &self.crypto
    }
    fn rand(&self) -> &Self::RandProvider {
        &self.crypto
    }
    fn storage(&self) -> &Self::StorageProvider {
        &self.storage
    }
}

#[cfg_attr(target_arch = "wasm32", allow(dead_code))]
fn serialize_memstore(s: &MemoryStorage) -> Result<Vec<u8>> {
    let map = s
        .values
        .read()
        .map_err(|e| Error::Io(format!("memstore lock: {e}")))?;
    let pairs: Vec<(serde_bytes::ByteBuf, serde_bytes::ByteBuf)> = map
        .iter()
        .map(|(k, v)| {
            (
                serde_bytes::ByteBuf::from(k.clone()),
                serde_bytes::ByteBuf::from(v.clone()),
            )
        })
        .collect();
    let blob = PersistedBlob {
        v: PERSIST_VERSION,
        pairs,
    };
    let mut bytes = Vec::with_capacity(map.values().map(Vec::len).sum::<usize>() + 64);
    ciborium::ser::into_writer(&blob, &mut bytes)
        .map_err(|e| Error::Codec(format!("encode blob: {e}")))?;
    Ok(bytes)
}

#[cfg_attr(target_arch = "wasm32", allow(dead_code))]
fn load_into_memstore(s: &MemoryStorage, bytes: &[u8]) -> Result<()> {
    let blob: PersistedBlob =
        ciborium::de::from_reader(bytes).map_err(|e| Error::Codec(format!("decode blob: {e}")))?;
    if blob.v != PERSIST_VERSION {
        return Err(Error::Codec(format!(
            "persisted blob v={} not supported (this SDK supports v={})",
            blob.v, PERSIST_VERSION
        )));
    }
    let mut map = s
        .values
        .write()
        .map_err(|e| Error::Io(format!("memstore lock: {e}")))?;
    map.clear();
    for (k, v) in blob.pairs {
        map.insert(k.into_vec(), v.into_vec());
    }
    Ok(())
}

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

    #[test]
    fn memory_backend_checkpoint_is_noop() {
        let provider = PersistentMlsProvider::open(StorageBackend::Memory).expect("open");
        assert!(provider.checkpoint().is_ok());
    }

    #[test]
    fn memstore_round_trips_through_serialize_deserialize() {
        let s = MemoryStorage::default();
        s.values
            .write()
            .unwrap()
            .insert(b"k1".to_vec(), b"v1".to_vec());
        s.values
            .write()
            .unwrap()
            .insert(b"k2".to_vec(), b"v2".to_vec());

        let bytes = serialize_memstore(&s).expect("serialize");

        let s2 = MemoryStorage::default();
        load_into_memstore(&s2, &bytes).expect("load");
        let map = s2.values.read().unwrap();
        assert_eq!(
            map.get(&b"k1".to_vec()).map(|v| v.as_slice()),
            Some(b"v1".as_ref())
        );
        assert_eq!(
            map.get(&b"k2".to_vec()).map(|v| v.as_slice()),
            Some(b"v2".as_ref())
        );
        assert_eq!(map.len(), 2);
    }
}