Skip to main content

xmtp_db/
sql_key_store.rs

1use xmtp_common::{ErrorCode, RetryableError, retryable};
2
3use self::transactions::MutableTransactionConnection;
4use crate::{ConnectionExt, TransactionalKeyStore, XmtpMlsStorageProvider};
5
6use bincode;
7use diesel::{
8    prelude::*,
9    sql_types::Binary,
10    {RunQueryDsl, sql_query},
11};
12use openmls_traits::storage::*;
13use serde::Serialize;
14use xmtp_configuration::OPENMLS_KV_TARGET;
15
16#[cfg(any(feature = "test-utils", test))]
17pub mod mock;
18mod transactions;
19
20const SELECT_QUERY: &str =
21    "SELECT value_bytes FROM openmls_key_value WHERE key_bytes = ? AND version = ?";
22const REPLACE_QUERY: &str =
23    "REPLACE INTO openmls_key_value (key_bytes, version, value_bytes) VALUES (?, ?, ?)";
24const UPDATE_QUERY: &str =
25    "UPDATE openmls_key_value SET value_bytes = ? WHERE key_bytes = ? AND version = ?";
26const DELETE_QUERY: &str = "DELETE FROM openmls_key_value WHERE key_bytes = ? AND version = ?";
27
28#[cfg(feature = "test-utils")]
29#[derive(
30    Selectable, Queryable, QueryableByName, Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash,
31)]
32#[diesel(table_name = crate::schema::openmls_key_value)]
33pub struct OpenMlsKeyValue {
34    pub version: i32,
35    pub key_bytes: Vec<u8>,
36    pub value_bytes: Vec<u8>,
37}
38
39#[cfg(feature = "test-utils")]
40impl OpenMlsKeyValue {
41    pub fn hash_all(conn: &mut SqliteConnection) -> Result<Vec<u8>, diesel::result::Error> {
42        use crate::schema::openmls_key_value;
43        use xmtp_common::Sha2Digest;
44        let values = openmls_key_value::table
45            .order(openmls_key_value::version.asc())
46            .order(openmls_key_value::key_bytes.asc())
47            .load_iter::<OpenMlsKeyValue, _>(conn)?;
48
49        let mut hasher = xmtp_common::Sha256Digest::new();
50        for (i, result) in values.enumerate() {
51            let value = result?;
52            hasher.update(b"version");
53            hasher.update(value.version.to_be_bytes());
54            hasher.update(b"key_bytes");
55            hasher.update(&value.key_bytes);
56            hasher.update(b"value_bytes");
57            hasher.update(&value.value_bytes);
58            hasher.update(b"index");
59            hasher.update(i.to_be_bytes());
60            hasher.update(b"\n");
61        }
62        Ok(hasher.finalize().to_vec())
63    }
64}
65
66#[derive(QueryableByName, Debug, Clone, PartialEq, Eq)]
67#[diesel(table_name = openmls_key_value)]
68struct StorageData {
69    #[diesel(sql_type = Binary)]
70    value_bytes: Vec<u8>,
71}
72
73impl TransactionalKeyStore for diesel::SqliteConnection {
74    type Store<'a>
75        = SqlKeyStore<MutableTransactionConnection<'a>>
76    where
77        Self: 'a;
78
79    fn key_store<'a>(&'a mut self) -> Self::Store<'a> {
80        SqlKeyStore::new_transactional(self)
81    }
82}
83
84#[derive(Clone)]
85pub struct SqlKeyStore<T> {
86    // Directly wrap the DbConnection which is a SqliteConnection in this case
87    conn: T,
88}
89
90// Test-only instrumentation (compiled out of release/production builds): counts
91// openmls KV read round-trips (`select_query`) so the metadata
92// read-amplification fix can be measured. Every `read`/`read_list`/
93// `group_context` bottoms out in `select_query`, so this is the single point
94// that observes an actual storage read.
95//
96// The counter is a **tokio task-local**, established for the duration of a
97// [`count_kv_reads`] call. Being task-scoped (not thread- or process-scoped) it
98// follows the measured work across `.await` points and across worker threads on
99// a multi-threaded runtime, and stays isolated from any other concurrently
100// running task. Reads performed outside a `count_kv_reads` scope — production
101// code and unrelated tests — are ignored, so the measurement is deterministic
102// regardless of runtime flavor or whether the measured accessors are sync or
103// async.
104// Native-only: the sole consumer is the native metadata read-amplification test,
105// and the task-local relies on the (native-only) `tokio` dep enabled by
106// `test-utils`.
107#[cfg(all(any(test, feature = "test-utils"), not(target_arch = "wasm32")))]
108tokio::task_local! {
109    static KV_READS: std::cell::Cell<u64>;
110}
111
112/// Record one KV read round-trip against the enclosing [`count_kv_reads`] scope.
113/// Outside such a scope this is a no-op.
114#[cfg(all(any(test, feature = "test-utils"), not(target_arch = "wasm32")))]
115fn record_kv_read() {
116    let _ = KV_READS.try_with(|c| c.set(c.get() + 1));
117}
118
119/// Run `f` with a fresh KV-read counter scoped to the current task and return
120/// its result together with the number of `select_query` round-trips it made.
121#[cfg(all(any(test, feature = "test-utils"), not(target_arch = "wasm32")))]
122pub fn count_kv_reads<R>(f: impl FnOnce() -> R) -> (R, u64) {
123    KV_READS.sync_scope(std::cell::Cell::new(0), || {
124        let result = f();
125        let count = KV_READS.with(|c| c.get());
126        (result, count)
127    })
128}
129
130impl<A> SqlKeyStore<A> {
131    pub fn new(conn: A) -> Self {
132        Self { conn }
133    }
134}
135
136impl<'a> SqlKeyStore<SqliteConnection> {
137    pub fn new_transactional(
138        conn: &'a mut SqliteConnection,
139    ) -> SqlKeyStore<MutableTransactionConnection<'a>> {
140        SqlKeyStore {
141            conn: MutableTransactionConnection::new(conn),
142        }
143    }
144}
145
146// refactor to use diesel directly
147impl<C> SqlKeyStore<C>
148where
149    C: ConnectionExt,
150{
151    fn select_query<const VERSION: u16>(
152        &self,
153        storage_key: &Vec<u8>,
154    ) -> Result<Vec<StorageData>, crate::ConnectionError> {
155        #[cfg(all(any(test, feature = "test-utils"), not(target_arch = "wasm32")))]
156        record_kv_read();
157        self.conn.raw_query(|conn| {
158            sql_query(SELECT_QUERY)
159                .bind::<diesel::sql_types::Binary, _>(&storage_key)
160                .bind::<diesel::sql_types::Integer, _>(VERSION as i32)
161                .load(conn)
162        })
163    }
164
165    fn replace_query<const VERSION: u16>(
166        &self,
167        storage_key: &Vec<u8>,
168        value: &[u8],
169    ) -> Result<usize, crate::ConnectionError> {
170        self.conn.raw_query(|conn| {
171            sql_query(REPLACE_QUERY)
172                .bind::<diesel::sql_types::Binary, _>(&storage_key)
173                .bind::<diesel::sql_types::Integer, _>(VERSION as i32)
174                .bind::<diesel::sql_types::Binary, _>(&value)
175                .execute(conn)
176        })
177    }
178
179    fn update_query<const VERSION: u16>(
180        &self,
181        storage_key: &Vec<u8>,
182        modified_data: &Vec<u8>,
183    ) -> Result<usize, crate::ConnectionError> {
184        self.conn.raw_query(|conn| {
185            sql_query(UPDATE_QUERY)
186                .bind::<diesel::sql_types::Binary, _>(&modified_data)
187                .bind::<diesel::sql_types::Binary, _>(&storage_key)
188                .bind::<diesel::sql_types::Integer, _>(VERSION as i32)
189                .execute(conn)
190        })
191    }
192
193    pub fn write<const VERSION: u16>(
194        &self,
195        label: &[u8],
196        key: &[u8],
197        value: &[u8],
198    ) -> Result<(), <Self as StorageProvider<CURRENT_VERSION>>::Error> {
199        let storage_key = build_key_from_vec::<VERSION>(label, key.to_vec());
200        let _ = self.replace_query::<VERSION>(&storage_key, value)?;
201        Ok(())
202    }
203
204    pub fn append<const VERSION: u16>(
205        &self,
206        label: &[u8],
207        key: &[u8],
208        value: &[u8],
209    ) -> Result<(), <Self as StorageProvider<CURRENT_VERSION>>::Error> {
210        tracing::trace!("append {}", String::from_utf8_lossy(label));
211
212        let storage_key = build_key_from_vec::<VERSION>(label, key.to_vec());
213        let data = self.select_query::<VERSION>(&storage_key)?;
214
215        if let Some(entry) = data.into_iter().next() {
216            // The value in the storage is an array of array of bytes
217            let mut deserialized = deserialize_bincode::<Vec<Vec<u8>>>(label, &entry.value_bytes)?;
218            deserialized.push(value.to_vec());
219            let modified_data = bincode::serialize(&deserialized)?;
220            let _ = self.update_query::<VERSION>(&storage_key, &modified_data)?;
221            Ok(())
222        } else {
223            // Add a first entry
224            let value_bytes = &bincode::serialize(&vec![value])?;
225            let _ = self.replace_query::<VERSION>(&storage_key, value_bytes)?;
226
227            Ok(())
228        }
229    }
230
231    pub fn remove_item<const VERSION: u16>(
232        &self,
233        label: &[u8],
234        key: &[u8],
235        value: &[u8],
236    ) -> Result<(), <Self as StorageProvider<CURRENT_VERSION>>::Error> {
237        tracing::trace!("remove_item {}", String::from_utf8_lossy(label));
238
239        let storage_key = build_key_from_vec::<VERSION>(label, key.to_vec());
240        let data: Vec<StorageData> = self.select_query::<VERSION>(&storage_key)?;
241
242        if let Some(entry) = data.into_iter().next() {
243            // The value in the storage is an array of array of bytes.
244            let mut deserialized = deserialize_bincode::<Vec<Vec<u8>>>(label, &entry.value_bytes)?;
245            let vpos = deserialized.iter().position(|v| v == value);
246
247            if let Some(pos) = vpos {
248                deserialized.remove(pos);
249            }
250            let modified_data = bincode::serialize(&deserialized)
251                .map_err(|_| SqlKeyStoreError::SerializationError)?;
252
253            let _ = self.update_query::<VERSION>(&storage_key, &modified_data)?;
254            Ok(())
255        } else {
256            // Add a first entry
257            let value_bytes =
258                bincode::serialize(&[value]).map_err(|_| SqlKeyStoreError::SerializationError)?;
259            let _ = self.replace_query::<VERSION>(&storage_key, &value_bytes)?;
260            Ok(())
261        }
262    }
263
264    pub fn read<const VERSION: u16, V: Entity<VERSION>>(
265        &self,
266        label: &[u8],
267        key: &[u8],
268    ) -> Result<Option<V>, <Self as StorageProvider<CURRENT_VERSION>>::Error> {
269        let storage_key = build_key_from_vec::<VERSION>(label, key.to_vec());
270
271        let data = self.select_query::<VERSION>(&storage_key)?;
272
273        if let Some(entry) = data.into_iter().next() {
274            let deserialized = deserialize_bincode::<V>(label, &entry.value_bytes)?;
275
276            Ok(Some(deserialized))
277        } else {
278            Ok(None)
279        }
280    }
281
282    pub fn read_list<const VERSION: u16, V: Entity<VERSION>>(
283        &self,
284        label: &[u8],
285        key: &[u8],
286    ) -> Result<Vec<V>, <Self as StorageProvider<CURRENT_VERSION>>::Error> {
287        let storage_key = build_key_from_vec::<VERSION>(label, key.to_vec());
288        let results = self.select_query::<VERSION>(&storage_key)?;
289
290        if let Some(entry) = results.into_iter().next() {
291            let list = deserialize_bincode::<Vec<Vec<u8>>>(label, &entry.value_bytes)?;
292
293            // Read the values from the bytes in the list
294            let mut deserialized_list = Vec::with_capacity(list.len());
295            for v in list {
296                deserialized_list.push(deserialize_bincode::<V>(label, &v)?);
297            }
298            Ok(deserialized_list)
299        } else {
300            Ok(vec![])
301        }
302    }
303
304    pub fn delete<const VERSION: u16>(
305        &self,
306        label: &[u8],
307        key: &[u8],
308    ) -> Result<(), <Self as StorageProvider<CURRENT_VERSION>>::Error> {
309        let storage_key = build_key_from_vec::<VERSION>(label, key.to_vec());
310        self.conn.raw_query(|conn| {
311            sql_query(DELETE_QUERY)
312                .bind::<diesel::sql_types::Binary, _>(&storage_key)
313                .bind::<diesel::sql_types::Integer, _>(VERSION as i32)
314                .execute(conn)
315        })?;
316        Ok(())
317    }
318}
319
320/// Errors thrown by the key store.
321/// General error type for Mls Storage Trait
322#[derive(thiserror::Error, Debug, ErrorCode)]
323pub enum SqlKeyStoreError {
324    /// Unsupported value type.
325    ///
326    /// Key store does not allow storing serialized values. Not retryable.
327    #[error("The key store does not allow storing serialized values.")]
328    UnsupportedValueTypeBytes,
329    /// Unsupported method.
330    ///
331    /// PSK operations not supported by this key store. Not retryable.
332    #[error("Updating is not supported by this key store.")]
333    UnsupportedMethod,
334    /// Serialization error.
335    ///
336    /// Failed to serialize value for key store. Not retryable.
337    #[error("Error serializing value.")]
338    SerializationError,
339    /// Value not found.
340    ///
341    /// Requested key not in OpenMLS key store. Not retryable.
342    #[error("Value does not exist.")]
343    NotFound,
344    /// Database error.
345    ///
346    /// Underlying Diesel database error. May be retryable.
347    #[error("database error: {0}")]
348    Storage(#[from] diesel::result::Error),
349    /// Connection error.
350    ///
351    /// Database connection error. Retryable.
352    #[error("connection {0}")]
353    Connection(#[from] crate::ConnectionError),
354}
355
356impl RetryableError for SqlKeyStoreError {
357    fn is_retryable(&self) -> bool {
358        use SqlKeyStoreError::*;
359        match self {
360            Storage(err) => retryable!(err),
361            SerializationError => false,
362            UnsupportedMethod => false,
363            UnsupportedValueTypeBytes => false,
364            NotFound => false,
365            Connection(c) => retryable!(c),
366        }
367    }
368}
369
370const KEY_PACKAGE_LABEL: &[u8] = b"KeyPackage";
371const ENCRYPTION_KEY_PAIR_LABEL: &[u8] = b"EncryptionKeyPair";
372const SIGNATURE_KEY_PAIR_LABEL: &[u8] = b"SignatureKeyPair";
373const EPOCH_KEY_PAIRS_LABEL: &[u8] = b"EpochKeyPairs";
374pub const KEY_PACKAGE_REFERENCES: &[u8] = b"KeyPackageReferences";
375pub const KEY_PACKAGE_WRAPPER_PRIVATE_KEY: &[u8] = b"KeyPackageWrapperPrivateKey";
376pub const COMMIT_LOG_SIGNER_PRIVATE_KEY: &[u8] = b"CommitLogSignerPrivateKey";
377
378// related to PublicGroup
379const TREE_LABEL: &[u8] = b"Tree";
380const GROUP_CONTEXT_LABEL: &[u8] = b"GroupContext";
381const INTERIM_TRANSCRIPT_HASH_LABEL: &[u8] = b"InterimTranscriptHash";
382const CONFIRMATION_TAG_LABEL: &[u8] = b"ConfirmationTag";
383
384// related to CoreGroup
385const OWN_LEAF_NODE_INDEX_LABEL: &[u8] = b"OwnLeafNodeIndex";
386const EPOCH_SECRETS_LABEL: &[u8] = b"EpochSecrets";
387const MESSAGE_SECRETS_LABEL: &[u8] = b"MessageSecrets";
388
389// related to MlsGroup
390const JOIN_CONFIG_LABEL: &[u8] = b"MlsGroupJoinConfig";
391const OWN_LEAF_NODES_LABEL: &[u8] = b"OwnLeafNodes";
392const GROUP_STATE_LABEL: &[u8] = b"GroupState";
393const QUEUED_PROPOSAL_LABEL: &[u8] = b"QueuedProposal";
394const PROPOSAL_QUEUE_REFS_LABEL: &[u8] = b"ProposalQueueRefs";
395const RESUMPTION_PSK_STORE_LABEL: &[u8] = b"ResumptionPskStore";
396
397// related to ApplicationExportTree
398const APPLICATION_EXPORT_TREE_LABEL: &[u8] = b"ApplicationExportTree";
399
400impl<C> StorageProvider<CURRENT_VERSION> for SqlKeyStore<C>
401where
402    C: ConnectionExt,
403{
404    type Error = SqlKeyStoreError;
405
406    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id), proposal_ref = %hex_kv(proposal_ref)), err)]
407    fn queue_proposal<
408        GroupId: traits::GroupId<CURRENT_VERSION>,
409        ProposalRef: traits::ProposalRef<CURRENT_VERSION>,
410        QueuedProposal: traits::QueuedProposal<CURRENT_VERSION>,
411    >(
412        &self,
413        group_id: &GroupId,
414        proposal_ref: &ProposalRef,
415        proposal: &QueuedProposal,
416    ) -> Result<(), Self::Error> {
417        // write proposal to key (group_id, proposal_ref)
418        let key = bincode::serialize(&(group_id, proposal_ref))?;
419        let value = bincode::serialize(proposal)?;
420        self.write::<CURRENT_VERSION>(QUEUED_PROPOSAL_LABEL, &key, &value)?;
421
422        // update proposal list for group_id
423        let key = build_key::<CURRENT_VERSION, &GroupId>(PROPOSAL_QUEUE_REFS_LABEL, group_id)?;
424        let value = bincode::serialize(proposal_ref)?;
425        self.append::<CURRENT_VERSION>(PROPOSAL_QUEUE_REFS_LABEL, &key, &value)?;
426
427        Ok(())
428    }
429
430    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
431    fn write_tree<
432        GroupId: traits::GroupId<CURRENT_VERSION>,
433        TreeSync: traits::TreeSync<CURRENT_VERSION>,
434    >(
435        &self,
436        group_id: &GroupId,
437        tree: &TreeSync,
438    ) -> Result<(), Self::Error> {
439        let key = build_key::<CURRENT_VERSION, &GroupId>(TREE_LABEL, group_id)?;
440        let value = bincode::serialize(&tree)?;
441        self.write::<CURRENT_VERSION>(TREE_LABEL, &key, &value)
442    }
443
444    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
445    fn write_interim_transcript_hash<
446        GroupId: traits::GroupId<CURRENT_VERSION>,
447        InterimTranscriptHash: traits::InterimTranscriptHash<CURRENT_VERSION>,
448    >(
449        &self,
450        group_id: &GroupId,
451        interim_transcript_hash: &InterimTranscriptHash,
452    ) -> Result<(), Self::Error> {
453        let key = build_key::<CURRENT_VERSION, &GroupId>(INTERIM_TRANSCRIPT_HASH_LABEL, group_id)?;
454        let value = bincode::serialize(&interim_transcript_hash)?;
455        let _ = self.write::<CURRENT_VERSION>(INTERIM_TRANSCRIPT_HASH_LABEL, &key, &value);
456
457        Ok(())
458    }
459
460    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id), group_context = %hex_kv(group_context)), err)]
461    fn write_context<
462        GroupId: traits::GroupId<CURRENT_VERSION>,
463        GroupContext: traits::GroupContext<CURRENT_VERSION>,
464    >(
465        &self,
466        group_id: &GroupId,
467        group_context: &GroupContext,
468    ) -> Result<(), Self::Error> {
469        let key = build_key::<CURRENT_VERSION, &GroupId>(GROUP_CONTEXT_LABEL, group_id)?;
470        let value = bincode::serialize(&group_context)?;
471
472        self.write::<CURRENT_VERSION>(GROUP_CONTEXT_LABEL, &key, &value)
473    }
474
475    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
476    fn write_confirmation_tag<
477        GroupId: traits::GroupId<CURRENT_VERSION>,
478        ConfirmationTag: traits::ConfirmationTag<CURRENT_VERSION>,
479    >(
480        &self,
481        group_id: &GroupId,
482        confirmation_tag: &ConfirmationTag,
483    ) -> Result<(), Self::Error> {
484        let key = build_key::<CURRENT_VERSION, &GroupId>(CONFIRMATION_TAG_LABEL, group_id)?;
485        let value = bincode::serialize(&confirmation_tag)?;
486
487        self.write::<CURRENT_VERSION>(CONFIRMATION_TAG_LABEL, &key, &value)
488    }
489
490    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(public_key = %hex_kv(public_key)), err)]
491    fn write_signature_key_pair<
492        SignaturePublicKey: traits::SignaturePublicKey<CURRENT_VERSION>,
493        SignatureKeyPair: traits::SignatureKeyPair<CURRENT_VERSION>,
494    >(
495        &self,
496        public_key: &SignaturePublicKey,
497        signature_key_pair: &SignatureKeyPair,
498    ) -> Result<(), Self::Error> {
499        let key = build_key::<CURRENT_VERSION, &SignaturePublicKey>(
500            SIGNATURE_KEY_PAIR_LABEL,
501            public_key,
502        )?;
503        let value = bincode::serialize(&signature_key_pair)?;
504
505        self.write::<CURRENT_VERSION>(SIGNATURE_KEY_PAIR_LABEL, &key, &value)
506    }
507
508    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
509    fn queued_proposal_refs<
510        GroupId: traits::GroupId<CURRENT_VERSION>,
511        ProposalRef: traits::ProposalRef<CURRENT_VERSION>,
512    >(
513        &self,
514        group_id: &GroupId,
515    ) -> Result<Vec<ProposalRef>, Self::Error> {
516        let key = build_key::<CURRENT_VERSION, &GroupId>(PROPOSAL_QUEUE_REFS_LABEL, group_id)?;
517        self.read_list(PROPOSAL_QUEUE_REFS_LABEL, &key)
518    }
519
520    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
521    fn queued_proposals<
522        GroupId: traits::GroupId<CURRENT_VERSION>,
523        ProposalRef: traits::ProposalRef<CURRENT_VERSION>,
524        QueuedProposal: traits::QueuedProposal<CURRENT_VERSION>,
525    >(
526        &self,
527        group_id: &GroupId,
528    ) -> Result<Vec<(ProposalRef, QueuedProposal)>, Self::Error> {
529        let key = build_key::<CURRENT_VERSION, &GroupId>(PROPOSAL_QUEUE_REFS_LABEL, group_id)?;
530        let refs: Vec<ProposalRef> = self.read_list(PROPOSAL_QUEUE_REFS_LABEL, &key)?;
531
532        refs.into_iter()
533            .map(|proposal_ref| -> Result<_, _> {
534                let key = bincode::serialize(&(group_id, &proposal_ref))?;
535                match self.read(QUEUED_PROPOSAL_LABEL, &key)? {
536                    Some(proposal) => Ok((proposal_ref, proposal)),
537                    None => Err(SqlKeyStoreError::NotFound),
538                }
539            })
540            .collect::<Result<Vec<_>, _>>()
541    }
542
543    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
544    fn tree<
545        GroupId: traits::GroupId<CURRENT_VERSION>,
546        TreeSync: traits::TreeSync<CURRENT_VERSION>,
547    >(
548        &self,
549        group_id: &GroupId,
550    ) -> Result<Option<TreeSync>, Self::Error> {
551        let key = build_key::<CURRENT_VERSION, &GroupId>(TREE_LABEL, group_id)?;
552
553        self.read(TREE_LABEL, &key)
554    }
555
556    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
557    fn group_context<
558        GroupId: traits::GroupId<CURRENT_VERSION>,
559        GroupContext: traits::GroupContext<CURRENT_VERSION>,
560    >(
561        &self,
562        group_id: &GroupId,
563    ) -> Result<Option<GroupContext>, Self::Error> {
564        let key = build_key::<CURRENT_VERSION, &GroupId>(GROUP_CONTEXT_LABEL, group_id)?;
565
566        self.read(GROUP_CONTEXT_LABEL, &key)
567    }
568
569    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
570    fn interim_transcript_hash<
571        GroupId: traits::GroupId<CURRENT_VERSION>,
572        InterimTranscriptHash: traits::InterimTranscriptHash<CURRENT_VERSION>,
573    >(
574        &self,
575        group_id: &GroupId,
576    ) -> Result<Option<InterimTranscriptHash>, Self::Error> {
577        let key = build_key::<CURRENT_VERSION, &GroupId>(INTERIM_TRANSCRIPT_HASH_LABEL, group_id)?;
578
579        self.read(INTERIM_TRANSCRIPT_HASH_LABEL, &key)
580    }
581
582    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
583    fn confirmation_tag<
584        GroupId: traits::GroupId<CURRENT_VERSION>,
585        ConfirmationTag: traits::ConfirmationTag<CURRENT_VERSION>,
586    >(
587        &self,
588        group_id: &GroupId,
589    ) -> Result<Option<ConfirmationTag>, Self::Error> {
590        let key = build_key::<CURRENT_VERSION, &GroupId>(CONFIRMATION_TAG_LABEL, group_id)?;
591
592        self.read(CONFIRMATION_TAG_LABEL, &key)
593    }
594
595    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(public_key = %hex_kv(public_key)), err)]
596    fn signature_key_pair<
597        SignaturePublicKey: traits::SignaturePublicKey<CURRENT_VERSION>,
598        SignatureKeyPair: traits::SignatureKeyPair<CURRENT_VERSION>,
599    >(
600        &self,
601        public_key: &SignaturePublicKey,
602    ) -> Result<Option<SignatureKeyPair>, Self::Error> {
603        let key = build_key::<CURRENT_VERSION, &SignaturePublicKey>(
604            SIGNATURE_KEY_PAIR_LABEL,
605            public_key,
606        )?;
607
608        self.read(SIGNATURE_KEY_PAIR_LABEL, &key)
609    }
610
611    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(hash_ref = %hex_kv(hash_ref), key_package = %hex_kv(key_package)), err)]
612    fn write_key_package<
613        HashReference: traits::HashReference<CURRENT_VERSION>,
614        KeyPackage: traits::KeyPackage<CURRENT_VERSION>,
615    >(
616        &self,
617        hash_ref: &HashReference,
618        key_package: &KeyPackage,
619    ) -> Result<(), Self::Error> {
620        let key = build_key::<CURRENT_VERSION, &HashReference>(KEY_PACKAGE_LABEL, hash_ref)?;
621        let value = bincode::serialize(&key_package)?;
622
623        // Store the key package
624        self.write::<CURRENT_VERSION>(KEY_PACKAGE_LABEL, &key, &value)
625    }
626
627    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(_psk_id = %hex_kv(_psk_id)), err)]
628    fn write_psk<
629        PskId: traits::PskId<CURRENT_VERSION>,
630        PskBundle: traits::PskBundle<CURRENT_VERSION>,
631    >(
632        &self,
633        _psk_id: &PskId,
634        _psk: &PskBundle,
635    ) -> Result<(), Self::Error> {
636        Ok(())
637    }
638
639    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(public_key = %hex_kv(public_key)), err)]
640    fn write_encryption_key_pair<
641        EncryptionKey: traits::EncryptionKey<CURRENT_VERSION>,
642        HpkeKeyPair: traits::HpkeKeyPair<CURRENT_VERSION>,
643    >(
644        &self,
645        public_key: &EncryptionKey,
646        key_pair: &HpkeKeyPair,
647    ) -> Result<(), Self::Error> {
648        let key =
649            build_key::<CURRENT_VERSION, &EncryptionKey>(ENCRYPTION_KEY_PAIR_LABEL, public_key)?;
650
651        self.write::<CURRENT_VERSION>(
652            ENCRYPTION_KEY_PAIR_LABEL,
653            &key,
654            &bincode::serialize(key_pair)?,
655        )
656    }
657
658    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(hash_ref = %hex_kv(hash_ref)), err)]
659    fn key_package<
660        HashReference: traits::HashReference<CURRENT_VERSION>,
661        KeyPackage: traits::KeyPackage<CURRENT_VERSION>,
662    >(
663        &self,
664        hash_ref: &HashReference,
665    ) -> Result<Option<KeyPackage>, Self::Error> {
666        let key = build_key::<CURRENT_VERSION, &HashReference>(KEY_PACKAGE_LABEL, hash_ref)?;
667
668        self.read(KEY_PACKAGE_LABEL, &key)
669    }
670
671    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(_psk_id = %hex_kv(_psk_id)), err)]
672    fn psk<PskBundle: traits::PskBundle<CURRENT_VERSION>, PskId: traits::PskId<CURRENT_VERSION>>(
673        &self,
674        _psk_id: &PskId,
675    ) -> Result<Option<PskBundle>, Self::Error> {
676        Ok(None)
677    }
678
679    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(public_key = %hex_kv(public_key)), err)]
680    fn encryption_key_pair<
681        HpkeKeyPair: traits::HpkeKeyPair<CURRENT_VERSION>,
682        EncryptionKey: traits::EncryptionKey<CURRENT_VERSION>,
683    >(
684        &self,
685        public_key: &EncryptionKey,
686    ) -> Result<Option<HpkeKeyPair>, Self::Error> {
687        let key =
688            build_key::<CURRENT_VERSION, &EncryptionKey>(ENCRYPTION_KEY_PAIR_LABEL, public_key)?;
689
690        self.read(ENCRYPTION_KEY_PAIR_LABEL, &key)
691    }
692
693    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(public_key = %hex_kv(public_key)), err)]
694    fn delete_signature_key_pair<
695        SignaturePublicKey: traits::SignaturePublicKey<CURRENT_VERSION>,
696    >(
697        &self,
698        public_key: &SignaturePublicKey,
699    ) -> Result<(), Self::Error> {
700        let key = build_key::<CURRENT_VERSION, &SignaturePublicKey>(
701            SIGNATURE_KEY_PAIR_LABEL,
702            public_key,
703        )?;
704
705        self.delete::<CURRENT_VERSION>(SIGNATURE_KEY_PAIR_LABEL, &key)
706    }
707
708    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(public_key = %hex_kv(public_key)), err)]
709    fn delete_encryption_key_pair<EncryptionKey: traits::EncryptionKey<CURRENT_VERSION>>(
710        &self,
711        public_key: &EncryptionKey,
712    ) -> Result<(), Self::Error> {
713        let key =
714            build_key::<CURRENT_VERSION, &EncryptionKey>(ENCRYPTION_KEY_PAIR_LABEL, public_key)?;
715
716        self.delete::<CURRENT_VERSION>(ENCRYPTION_KEY_PAIR_LABEL, &key)
717    }
718
719    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(hash_ref = %hex_kv(hash_ref)), err)]
720    fn delete_key_package<HashReference: traits::HashReference<CURRENT_VERSION>>(
721        &self,
722        hash_ref: &HashReference,
723    ) -> Result<(), Self::Error> {
724        let key = build_key::<CURRENT_VERSION, &HashReference>(KEY_PACKAGE_LABEL, hash_ref)?;
725        self.delete::<CURRENT_VERSION>(KEY_PACKAGE_LABEL, &key)
726    }
727
728    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(_psk_id = %hex_kv(_psk_id)), err)]
729    fn delete_psk<PskKey: traits::PskId<CURRENT_VERSION>>(
730        &self,
731        _psk_id: &PskKey,
732    ) -> Result<(), Self::Error> {
733        Err(SqlKeyStoreError::UnsupportedMethod)
734    }
735
736    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
737    fn group_state<
738        GroupState: traits::GroupState<CURRENT_VERSION>,
739        GroupId: traits::GroupId<CURRENT_VERSION>,
740    >(
741        &self,
742        group_id: &GroupId,
743    ) -> Result<Option<GroupState>, Self::Error> {
744        let key = build_key::<CURRENT_VERSION, &GroupId>(GROUP_STATE_LABEL, group_id)?;
745
746        self.read(GROUP_STATE_LABEL, &key)
747    }
748
749    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
750    fn write_group_state<
751        GroupState: traits::GroupState<CURRENT_VERSION>,
752        GroupId: traits::GroupId<CURRENT_VERSION>,
753    >(
754        &self,
755        group_id: &GroupId,
756        group_state: &GroupState,
757    ) -> Result<(), Self::Error> {
758        let key = build_key::<CURRENT_VERSION, &GroupId>(GROUP_STATE_LABEL, group_id)?;
759
760        self.write::<CURRENT_VERSION>(GROUP_STATE_LABEL, &key, &bincode::serialize(group_state)?)
761    }
762
763    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
764    fn delete_group_state<GroupId: traits::GroupId<CURRENT_VERSION>>(
765        &self,
766        group_id: &GroupId,
767    ) -> Result<(), Self::Error> {
768        let key = build_key::<CURRENT_VERSION, &GroupId>(GROUP_STATE_LABEL, group_id)?;
769
770        self.delete::<CURRENT_VERSION>(GROUP_STATE_LABEL, &key)
771    }
772
773    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
774    fn message_secrets<
775        GroupId: traits::GroupId<CURRENT_VERSION>,
776        MessageSecrets: traits::MessageSecrets<CURRENT_VERSION>,
777    >(
778        &self,
779        group_id: &GroupId,
780    ) -> Result<Option<MessageSecrets>, Self::Error> {
781        let key = build_key::<CURRENT_VERSION, &GroupId>(MESSAGE_SECRETS_LABEL, group_id)?;
782
783        self.read(MESSAGE_SECRETS_LABEL, &key)
784    }
785
786    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
787    fn write_message_secrets<
788        GroupId: traits::GroupId<CURRENT_VERSION>,
789        MessageSecrets: traits::MessageSecrets<CURRENT_VERSION>,
790    >(
791        &self,
792        group_id: &GroupId,
793        message_secrets: &MessageSecrets,
794    ) -> Result<(), Self::Error> {
795        let key = build_key::<CURRENT_VERSION, &GroupId>(MESSAGE_SECRETS_LABEL, group_id)?;
796
797        self.write::<CURRENT_VERSION>(
798            MESSAGE_SECRETS_LABEL,
799            &key,
800            &bincode::serialize(message_secrets)?,
801        )
802    }
803
804    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
805    fn delete_message_secrets<GroupId: traits::GroupId<CURRENT_VERSION>>(
806        &self,
807        group_id: &GroupId,
808    ) -> Result<(), Self::Error> {
809        let key = build_key::<CURRENT_VERSION, &GroupId>(MESSAGE_SECRETS_LABEL, group_id)?;
810
811        self.delete::<CURRENT_VERSION>(MESSAGE_SECRETS_LABEL, &key)
812    }
813
814    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
815    fn resumption_psk_store<
816        GroupId: traits::GroupId<CURRENT_VERSION>,
817        ResumptionPskStore: traits::ResumptionPskStore<CURRENT_VERSION>,
818    >(
819        &self,
820        group_id: &GroupId,
821    ) -> Result<Option<ResumptionPskStore>, Self::Error> {
822        self.read(RESUMPTION_PSK_STORE_LABEL, &bincode::serialize(group_id)?)
823    }
824
825    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
826    fn write_resumption_psk_store<
827        GroupId: traits::GroupId<CURRENT_VERSION>,
828        ResumptionPskStore: traits::ResumptionPskStore<CURRENT_VERSION>,
829    >(
830        &self,
831        group_id: &GroupId,
832        resumption_psk_store: &ResumptionPskStore,
833    ) -> Result<(), Self::Error> {
834        self.write::<CURRENT_VERSION>(
835            RESUMPTION_PSK_STORE_LABEL,
836            &bincode::serialize(group_id)?,
837            &bincode::serialize(resumption_psk_store)?,
838        )
839    }
840
841    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
842    fn delete_all_resumption_psk_secrets<GroupId: traits::GroupId<CURRENT_VERSION>>(
843        &self,
844        group_id: &GroupId,
845    ) -> Result<(), Self::Error> {
846        self.delete::<CURRENT_VERSION>(RESUMPTION_PSK_STORE_LABEL, &bincode::serialize(group_id)?)
847    }
848
849    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
850    fn own_leaf_index<
851        GroupId: traits::GroupId<CURRENT_VERSION>,
852        LeafNodeIndex: traits::LeafNodeIndex<CURRENT_VERSION>,
853    >(
854        &self,
855        group_id: &GroupId,
856    ) -> Result<Option<LeafNodeIndex>, Self::Error> {
857        let key = build_key::<CURRENT_VERSION, &GroupId>(OWN_LEAF_NODE_INDEX_LABEL, group_id)?;
858        self.read(OWN_LEAF_NODE_INDEX_LABEL, &key)
859    }
860
861    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id), own_leaf_index = %hex_kv(own_leaf_index)), err)]
862    fn write_own_leaf_index<
863        GroupId: traits::GroupId<CURRENT_VERSION>,
864        LeafNodeIndex: traits::LeafNodeIndex<CURRENT_VERSION>,
865    >(
866        &self,
867        group_id: &GroupId,
868        own_leaf_index: &LeafNodeIndex,
869    ) -> Result<(), Self::Error> {
870        let key = build_key::<CURRENT_VERSION, &GroupId>(OWN_LEAF_NODE_INDEX_LABEL, group_id)?;
871        self.write::<CURRENT_VERSION>(
872            OWN_LEAF_NODE_INDEX_LABEL,
873            &key,
874            &bincode::serialize(own_leaf_index)?,
875        )
876    }
877
878    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
879    fn delete_own_leaf_index<GroupId: traits::GroupId<CURRENT_VERSION>>(
880        &self,
881        group_id: &GroupId,
882    ) -> Result<(), Self::Error> {
883        let key = build_key::<CURRENT_VERSION, &GroupId>(OWN_LEAF_NODE_INDEX_LABEL, group_id)?;
884        self.delete::<CURRENT_VERSION>(OWN_LEAF_NODE_INDEX_LABEL, &key)
885    }
886
887    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
888    fn group_epoch_secrets<
889        GroupId: traits::GroupId<CURRENT_VERSION>,
890        GroupEpochSecrets: traits::GroupEpochSecrets<CURRENT_VERSION>,
891    >(
892        &self,
893        group_id: &GroupId,
894    ) -> Result<Option<GroupEpochSecrets>, Self::Error> {
895        let key = build_key::<CURRENT_VERSION, &GroupId>(EPOCH_SECRETS_LABEL, group_id)?;
896        self.read(EPOCH_SECRETS_LABEL, &key)
897    }
898
899    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
900    fn write_group_epoch_secrets<
901        GroupId: traits::GroupId<CURRENT_VERSION>,
902        GroupEpochSecrets: traits::GroupEpochSecrets<CURRENT_VERSION>,
903    >(
904        &self,
905        group_id: &GroupId,
906        group_epoch_secrets: &GroupEpochSecrets,
907    ) -> Result<(), Self::Error> {
908        let key = build_key::<CURRENT_VERSION, &GroupId>(EPOCH_SECRETS_LABEL, group_id)?;
909        self.write::<CURRENT_VERSION>(
910            EPOCH_SECRETS_LABEL,
911            &key,
912            &bincode::serialize(group_epoch_secrets)?,
913        )
914    }
915
916    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
917    fn delete_group_epoch_secrets<GroupId: traits::GroupId<CURRENT_VERSION>>(
918        &self,
919        group_id: &GroupId,
920    ) -> Result<(), Self::Error> {
921        let key = build_key::<CURRENT_VERSION, &GroupId>(EPOCH_SECRETS_LABEL, group_id)?;
922        self.delete::<CURRENT_VERSION>(EPOCH_SECRETS_LABEL, &key)
923    }
924
925    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id), epoch = %hex_kv(epoch), leaf_index = %leaf_index), err)]
926    fn write_encryption_epoch_key_pairs<
927        GroupId: traits::GroupId<CURRENT_VERSION>,
928        EpochKey: traits::EpochKey<CURRENT_VERSION>,
929        HpkeKeyPair: traits::HpkeKeyPair<CURRENT_VERSION>,
930    >(
931        &self,
932        group_id: &GroupId,
933        epoch: &EpochKey,
934        leaf_index: u32,
935        key_pairs: &[HpkeKeyPair],
936    ) -> Result<(), Self::Error> {
937        let key = epoch_key_pairs_id(group_id, epoch, leaf_index)?;
938        let value = bincode::serialize(key_pairs)?;
939        tracing::trace!("Writing encryption epoch key pairs");
940
941        self.write::<CURRENT_VERSION>(EPOCH_KEY_PAIRS_LABEL, &key, &value)
942    }
943
944    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id), epoch = %hex_kv(epoch), leaf_index = %leaf_index), err)]
945    fn encryption_epoch_key_pairs<
946        GroupId: traits::GroupId<CURRENT_VERSION>,
947        EpochKey: traits::EpochKey<CURRENT_VERSION>,
948        HpkeKeyPair: traits::HpkeKeyPair<CURRENT_VERSION>,
949    >(
950        &self,
951        group_id: &GroupId,
952        epoch: &EpochKey,
953        leaf_index: u32,
954    ) -> Result<Vec<HpkeKeyPair>, Self::Error> {
955        tracing::trace!("Reading encryption epoch key pairs");
956
957        let key = epoch_key_pairs_id(group_id, epoch, leaf_index)?;
958        let storage_key = build_key_from_vec::<CURRENT_VERSION>(EPOCH_KEY_PAIRS_LABEL, key);
959        tracing::trace!("  key: {}", hex::encode(&storage_key));
960
961        let query = "SELECT value_bytes FROM openmls_key_value WHERE key_bytes = ? AND version = ?";
962
963        let data: Vec<StorageData> = self.conn.raw_query(|conn| {
964            sql_query(query)
965                .bind::<diesel::sql_types::Binary, _>(&storage_key)
966                .bind::<diesel::sql_types::Integer, _>(CURRENT_VERSION as i32)
967                .load(conn)
968        })?;
969
970        if let Some(entry) = data.into_iter().next() {
971            deserialize_bincode::<Vec<HpkeKeyPair>>(EPOCH_KEY_PAIRS_LABEL, &entry.value_bytes)
972        } else {
973            Ok(vec![])
974        }
975    }
976
977    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id), epoch = %hex_kv(epoch), leaf_index = %leaf_index), err)]
978    fn delete_encryption_epoch_key_pairs<
979        GroupId: traits::GroupId<CURRENT_VERSION>,
980        EpochKey: traits::EpochKey<CURRENT_VERSION>,
981    >(
982        &self,
983        group_id: &GroupId,
984        epoch: &EpochKey,
985        leaf_index: u32,
986    ) -> Result<(), Self::Error> {
987        let key = epoch_key_pairs_id(group_id, epoch, leaf_index)?;
988
989        self.delete::<CURRENT_VERSION>(EPOCH_KEY_PAIRS_LABEL, &key)
990    }
991
992    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
993    fn clear_proposal_queue<
994        GroupId: traits::GroupId<CURRENT_VERSION>,
995        ProposalRef: traits::ProposalRef<CURRENT_VERSION>,
996    >(
997        &self,
998        group_id: &GroupId,
999    ) -> Result<(), Self::Error> {
1000        let key = build_key::<CURRENT_VERSION, &GroupId>(PROPOSAL_QUEUE_REFS_LABEL, group_id)?;
1001        let proposal_refs: Vec<ProposalRef> = self.read_list(PROPOSAL_QUEUE_REFS_LABEL, &key)?;
1002
1003        for proposal_ref in proposal_refs {
1004            let key = bincode::serialize(&(group_id, proposal_ref))?;
1005            self.delete::<CURRENT_VERSION>(QUEUED_PROPOSAL_LABEL, &key)?;
1006        }
1007
1008        let key = build_key::<CURRENT_VERSION, &GroupId>(PROPOSAL_QUEUE_REFS_LABEL, group_id)?;
1009
1010        self.delete::<CURRENT_VERSION>(PROPOSAL_QUEUE_REFS_LABEL, &key)
1011    }
1012
1013    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
1014    fn mls_group_join_config<
1015        GroupId: traits::GroupId<CURRENT_VERSION>,
1016        MlsGroupJoinConfig: traits::MlsGroupJoinConfig<CURRENT_VERSION>,
1017    >(
1018        &self,
1019        group_id: &GroupId,
1020    ) -> Result<Option<MlsGroupJoinConfig>, Self::Error> {
1021        let key = build_key::<CURRENT_VERSION, &GroupId>(JOIN_CONFIG_LABEL, group_id)?;
1022
1023        self.read(JOIN_CONFIG_LABEL, &key)
1024    }
1025
1026    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id), config = %hex_kv(config)), err)]
1027    fn write_mls_join_config<
1028        GroupId: traits::GroupId<CURRENT_VERSION>,
1029        MlsGroupJoinConfig: traits::MlsGroupJoinConfig<CURRENT_VERSION>,
1030    >(
1031        &self,
1032        group_id: &GroupId,
1033        config: &MlsGroupJoinConfig,
1034    ) -> Result<(), Self::Error> {
1035        let key = build_key::<CURRENT_VERSION, &GroupId>(JOIN_CONFIG_LABEL, group_id)?;
1036        let value = bincode::serialize(config)?;
1037
1038        self.write::<CURRENT_VERSION>(JOIN_CONFIG_LABEL, &key, &value)
1039    }
1040
1041    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
1042    fn own_leaf_nodes<
1043        GroupId: traits::GroupId<CURRENT_VERSION>,
1044        LeafNode: traits::LeafNode<CURRENT_VERSION>,
1045    >(
1046        &self,
1047        group_id: &GroupId,
1048    ) -> Result<Vec<LeafNode>, Self::Error> {
1049        let key = build_key::<CURRENT_VERSION, &GroupId>(OWN_LEAF_NODES_LABEL, group_id)?;
1050
1051        self.read_list(OWN_LEAF_NODES_LABEL, &key)
1052    }
1053
1054    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id), leaf_node = %hex_kv(leaf_node)), err)]
1055    fn append_own_leaf_node<
1056        GroupId: traits::GroupId<CURRENT_VERSION>,
1057        LeafNode: traits::LeafNode<CURRENT_VERSION>,
1058    >(
1059        &self,
1060        group_id: &GroupId,
1061        leaf_node: &LeafNode,
1062    ) -> Result<(), Self::Error> {
1063        let key = build_key::<CURRENT_VERSION, &GroupId>(OWN_LEAF_NODES_LABEL, group_id)?;
1064        let value = bincode::serialize(leaf_node)?;
1065
1066        self.append::<CURRENT_VERSION>(OWN_LEAF_NODES_LABEL, &key, &value)
1067    }
1068
1069    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
1070    fn delete_own_leaf_nodes<GroupId: traits::GroupId<CURRENT_VERSION>>(
1071        &self,
1072        group_id: &GroupId,
1073    ) -> Result<(), Self::Error> {
1074        let key = build_key::<CURRENT_VERSION, &GroupId>(OWN_LEAF_NODES_LABEL, group_id)?;
1075        self.delete::<CURRENT_VERSION>(OWN_LEAF_NODES_LABEL, &key)
1076    }
1077
1078    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
1079    fn delete_group_config<GroupId: traits::GroupId<CURRENT_VERSION>>(
1080        &self,
1081        group_id: &GroupId,
1082    ) -> Result<(), Self::Error> {
1083        let key = build_key::<CURRENT_VERSION, &GroupId>(JOIN_CONFIG_LABEL, group_id)?;
1084        self.delete::<CURRENT_VERSION>(JOIN_CONFIG_LABEL, &key)
1085    }
1086
1087    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
1088    fn delete_tree<GroupId: traits::GroupId<CURRENT_VERSION>>(
1089        &self,
1090        group_id: &GroupId,
1091    ) -> Result<(), Self::Error> {
1092        let key = build_key::<CURRENT_VERSION, &GroupId>(TREE_LABEL, group_id)?;
1093
1094        self.delete::<CURRENT_VERSION>(TREE_LABEL, &key)
1095    }
1096
1097    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
1098    fn delete_confirmation_tag<GroupId: traits::GroupId<CURRENT_VERSION>>(
1099        &self,
1100        group_id: &GroupId,
1101    ) -> Result<(), Self::Error> {
1102        let key = build_key::<CURRENT_VERSION, &GroupId>(CONFIRMATION_TAG_LABEL, group_id)?;
1103
1104        self.delete::<CURRENT_VERSION>(CONFIRMATION_TAG_LABEL, &key)
1105    }
1106
1107    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
1108    fn delete_context<GroupId: traits::GroupId<CURRENT_VERSION>>(
1109        &self,
1110        group_id: &GroupId,
1111    ) -> Result<(), Self::Error> {
1112        let key = build_key::<CURRENT_VERSION, &GroupId>(GROUP_CONTEXT_LABEL, group_id)?;
1113
1114        self.delete::<CURRENT_VERSION>(GROUP_CONTEXT_LABEL, &key)
1115    }
1116
1117    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
1118    fn delete_interim_transcript_hash<GroupId: traits::GroupId<CURRENT_VERSION>>(
1119        &self,
1120        group_id: &GroupId,
1121    ) -> Result<(), Self::Error> {
1122        let key = build_key::<CURRENT_VERSION, &GroupId>(INTERIM_TRANSCRIPT_HASH_LABEL, group_id)?;
1123
1124        self.delete::<CURRENT_VERSION>(INTERIM_TRANSCRIPT_HASH_LABEL, &key)
1125    }
1126
1127    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id), proposal_ref = %hex_kv(proposal_ref)), err)]
1128    fn remove_proposal<
1129        GroupId: traits::GroupId<CURRENT_VERSION>,
1130        ProposalRef: traits::ProposalRef<CURRENT_VERSION>,
1131    >(
1132        &self,
1133        group_id: &GroupId,
1134        proposal_ref: &ProposalRef,
1135    ) -> Result<(), Self::Error> {
1136        // Delete the proposal ref
1137        let key = build_key::<CURRENT_VERSION, &GroupId>(PROPOSAL_QUEUE_REFS_LABEL, group_id)?;
1138        let value = bincode::serialize(proposal_ref)?;
1139        self.remove_item::<CURRENT_VERSION>(PROPOSAL_QUEUE_REFS_LABEL, &key, &value)?;
1140
1141        // Delete the proposal
1142        let key = bincode::serialize(&(group_id, proposal_ref))?;
1143        self.delete::<CURRENT_VERSION>(QUEUED_PROPOSAL_LABEL, &key)
1144    }
1145
1146    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
1147    fn write_application_export_tree<
1148        GroupId: traits::GroupId<CURRENT_VERSION>,
1149        ApplicationExportTree: traits::ApplicationExportTree<CURRENT_VERSION>,
1150    >(
1151        &self,
1152        group_id: &GroupId,
1153        application_export_tree: &ApplicationExportTree,
1154    ) -> Result<(), Self::Error> {
1155        let key = build_key::<CURRENT_VERSION, &GroupId>(APPLICATION_EXPORT_TREE_LABEL, group_id)?;
1156        self.write::<CURRENT_VERSION>(
1157            APPLICATION_EXPORT_TREE_LABEL,
1158            &key,
1159            &bincode::serialize(application_export_tree)?,
1160        )
1161    }
1162
1163    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
1164    fn application_export_tree<
1165        GroupId: traits::GroupId<CURRENT_VERSION>,
1166        ApplicationExportTree: traits::ApplicationExportTree<CURRENT_VERSION>,
1167    >(
1168        &self,
1169        group_id: &GroupId,
1170    ) -> Result<Option<ApplicationExportTree>, Self::Error> {
1171        let key = build_key::<CURRENT_VERSION, &GroupId>(APPLICATION_EXPORT_TREE_LABEL, group_id)?;
1172        self.read(APPLICATION_EXPORT_TREE_LABEL, &key)
1173    }
1174
1175    #[tracing::instrument(skip_all, target = OPENMLS_KV_TARGET, fields(group_id = %hex_kv(group_id)), err)]
1176    fn delete_application_export_tree<
1177        GroupId: traits::GroupId<CURRENT_VERSION>,
1178        ApplicationExportTree: traits::ApplicationExportTree<CURRENT_VERSION>,
1179    >(
1180        &self,
1181        group_id: &GroupId,
1182    ) -> Result<(), Self::Error> {
1183        let key = build_key::<CURRENT_VERSION, &GroupId>(APPLICATION_EXPORT_TREE_LABEL, group_id)?;
1184        self.delete::<CURRENT_VERSION>(APPLICATION_EXPORT_TREE_LABEL, &key)
1185    }
1186}
1187
1188/// Hex bincode of `v` for tracing fields. Empty string on serialize fail (field exprs can't propagate).
1189fn hex_kv<T: Serialize + ?Sized>(v: &T) -> String {
1190    bincode::serialize(v).map(hex::encode).unwrap_or_default()
1191}
1192
1193/// Build a key with version and label.
1194fn build_key_from_vec<const V: u16>(label: &[u8], key: Vec<u8>) -> Vec<u8> {
1195    let mut key_out = label.to_vec();
1196    key_out.extend_from_slice(&key);
1197    key_out.extend_from_slice(&u16::to_be_bytes(V));
1198    key_out
1199}
1200
1201/// Build a key with version and label.
1202fn build_key<const V: u16, K: Serialize>(
1203    label: &[u8],
1204    key: K,
1205) -> Result<Vec<u8>, SqlKeyStoreError> {
1206    let key_vec = bincode::serialize(&key)?;
1207    Ok(build_key_from_vec::<V>(label, key_vec))
1208}
1209
1210/// Bincode decode → `SerializationError`. With `deserialize-paths` feature + `openmls_kv` target
1211/// enabled, error path re-decodes via `serde_path_to_error` to log failing field path.
1212fn deserialize_bincode<'de, T>(
1213    #[cfg_attr(not(feature = "deserialize-paths"), allow(unused_variables))] label: &[u8],
1214    bytes: &'de [u8],
1215) -> Result<T, SqlKeyStoreError>
1216where
1217    T: serde::Deserialize<'de>,
1218{
1219    match bincode::deserialize::<T>(bytes) {
1220        Ok(val) => Ok(val),
1221        Err(_orig_err) => {
1222            #[cfg(feature = "deserialize-paths")]
1223            if tracing::event_enabled!(target: OPENMLS_KV_TARGET, tracing::Level::ERROR) {
1224                use bincode::Options;
1225                let opts = bincode::DefaultOptions::new()
1226                    .with_fixint_encoding()
1227                    .allow_trailing_bytes();
1228                let mut de = bincode::de::Deserializer::from_slice(bytes, opts);
1229                match serde_path_to_error::deserialize::<_, T>(&mut de) {
1230                    Ok(_) => {
1231                        // Second pass succeeded — bincode option mismatch; log err without path.
1232                        tracing::error!(
1233                            target: OPENMLS_KV_TARGET,
1234                            label = %String::from_utf8_lossy(label),
1235                            type_name = std::any::type_name::<T>(),
1236                            bytes_len = bytes.len(),
1237                            error = %_orig_err,
1238                            "bincode deserialize failed (no path captured)",
1239                        );
1240                    }
1241                    Err(e) => {
1242                        let path = e.path().to_string();
1243                        let inner = e.into_inner();
1244                        tracing::error!(
1245                            target: OPENMLS_KV_TARGET,
1246                            label = %String::from_utf8_lossy(label),
1247                            path = %path,
1248                            type_name = std::any::type_name::<T>(),
1249                            bytes_len = bytes.len(),
1250                            error = %inner,
1251                            "bincode deserialize failed",
1252                        );
1253                    }
1254                }
1255            }
1256            Err(SqlKeyStoreError::SerializationError)
1257        }
1258    }
1259}
1260
1261fn epoch_key_pairs_id(
1262    group_id: &impl traits::GroupId<CURRENT_VERSION>,
1263    epoch: &impl traits::EpochKey<CURRENT_VERSION>,
1264    leaf_index: u32,
1265) -> Result<Vec<u8>, SqlKeyStoreError> {
1266    let mut key = bincode::serialize(group_id)?;
1267    key.extend_from_slice(&bincode::serialize(epoch)?);
1268    key.extend_from_slice(&bincode::serialize(&leaf_index)?);
1269    Ok(key)
1270}
1271
1272impl From<bincode::Error> for SqlKeyStoreError {
1273    fn from(_: bincode::Error) -> Self {
1274        Self::SerializationError
1275    }
1276}
1277
1278#[cfg(any(test, feature = "test-utils"))]
1279impl SqlKeyStore<crate::test_utils::MemoryStorage> {
1280    pub fn kv_pairs(&self) -> String {
1281        self.conn.key_value_pairs()
1282    }
1283
1284    pub fn kv_pairs_utf8(&self) -> String {
1285        self.conn.key_value_pairs_utf8()
1286    }
1287}
1288
1289#[cfg(test)]
1290pub(crate) mod tests {
1291    use openmls::group::GroupId;
1292    use openmls_basic_credential::{SignatureKeyPair, StorageId};
1293    use openmls_traits::{
1294        OpenMlsProvider,
1295        storage::{
1296            CURRENT_VERSION, Entity, Key, StorageProvider,
1297            traits::{self},
1298        },
1299    };
1300    use serde::{Deserialize, Serialize};
1301
1302    use super::SqlKeyStore;
1303    use crate::TransactionOutcome::{Continue, Rollback};
1304    use crate::encrypted_store::MlsProviderExt;
1305    use crate::{
1306        XmtpTestDb, sql_key_store::SqlKeyStoreError, xmtp_openmls_provider::XmtpOpenMlsProvider,
1307    };
1308    use xmtp_cryptography::configuration::CIPHERSUITE;
1309
1310    #[xmtp_common::test]
1311    async fn store_read_delete() {
1312        let store = crate::TestDb::create_persistent_store(None).await;
1313        let conn = store.conn();
1314        let key_store = SqlKeyStore::new(conn);
1315
1316        let signature_keys = SignatureKeyPair::new(CIPHERSUITE.signature_algorithm()).unwrap();
1317        let public_key = StorageId::from(signature_keys.to_public_vec());
1318        assert!(
1319            key_store
1320                .signature_key_pair::<StorageId, SignatureKeyPair>(&public_key)
1321                .unwrap()
1322                .is_none()
1323        );
1324
1325        key_store
1326            .write_signature_key_pair::<StorageId, SignatureKeyPair>(&public_key, &signature_keys)
1327            .unwrap();
1328
1329        assert!(
1330            key_store
1331                .signature_key_pair::<StorageId, SignatureKeyPair>(&public_key)
1332                .unwrap()
1333                .is_some()
1334        );
1335
1336        key_store
1337            .delete_signature_key_pair::<StorageId>(&public_key)
1338            .unwrap();
1339
1340        assert!(
1341            key_store
1342                .signature_key_pair::<StorageId, SignatureKeyPair>(&public_key)
1343                .unwrap()
1344                .is_none()
1345        );
1346    }
1347
1348    #[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone)]
1349    struct Proposal(Vec<u8>);
1350    impl traits::QueuedProposal<CURRENT_VERSION> for Proposal {}
1351    impl Entity<CURRENT_VERSION> for Proposal {}
1352
1353    #[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Copy)]
1354    struct ProposalRef(usize);
1355    impl traits::ProposalRef<CURRENT_VERSION> for ProposalRef {}
1356    impl Key<CURRENT_VERSION> for ProposalRef {}
1357    impl Entity<CURRENT_VERSION> for ProposalRef {}
1358
1359    #[xmtp_common::test(unwrap_try = true)]
1360    async fn test_read_write() {
1361        let store = crate::TestDb::create_persistent_store(None).await;
1362        let conn = store.conn();
1363        let mls_store = SqlKeyStore::new(conn);
1364        let provider = XmtpOpenMlsProvider::new(mls_store);
1365        let key_store = provider.key_store();
1366
1367        let raw_value = vec![3u8; 32];
1368        let group_1 = bincode::serialize(&[1u8; 32])?;
1369        let group_2 = bincode::serialize(&[2u8; 32])?;
1370        let value_1 = bincode::serialize(&raw_value)?;
1371
1372        key_store.write::<CURRENT_VERSION>(
1373            crate::sql_key_store::COMMIT_LOG_SIGNER_PRIVATE_KEY,
1374            &group_1,
1375            &value_1,
1376        )?;
1377
1378        // Query on a value that hasn't been written
1379        let result = key_store.read::<CURRENT_VERSION, Vec<u8>>(
1380            crate::sql_key_store::COMMIT_LOG_SIGNER_PRIVATE_KEY,
1381            &group_2,
1382        );
1383        assert!(result.is_ok(), "{}", result.err().unwrap());
1384        assert!(result.unwrap().is_none());
1385
1386        let result = key_store.read::<CURRENT_VERSION, Vec<u8>>(
1387            crate::sql_key_store::COMMIT_LOG_SIGNER_PRIVATE_KEY,
1388            &group_1,
1389        );
1390        assert!(result.is_ok(), "{}", result.err().unwrap());
1391        assert_eq!(result.unwrap(), Some(raw_value));
1392    }
1393
1394    #[xmtp_common::test(unwrap_try = true)]
1395    async fn transaction_commit_persists_rollback_does_not_and_error_propagates() {
1396        use crate::{
1397            StorageError, TransactionOutcome, traits::TransactionalKeyStore,
1398            xmtp_openmls_provider::XmtpMlsStorageProvider,
1399        };
1400
1401        let store = crate::TestDb::create_persistent_store(None).await;
1402        let conn = store.conn();
1403        let key_store = SqlKeyStore::new(conn);
1404
1405        let committed_key = bincode::serialize(&[10u8; 32])?;
1406        let rolled_back_key = bincode::serialize(&[11u8; 32])?;
1407        let value = bincode::serialize(&vec![7u8; 16])?;
1408
1409        let is_present = |k: &[u8]| {
1410            key_store
1411                .read::<CURRENT_VERSION, Vec<u8>>(
1412                    crate::sql_key_store::COMMIT_LOG_SIGNER_PRIVATE_KEY,
1413                    k,
1414                )
1415                .unwrap()
1416                .is_some()
1417        };
1418
1419        // Commit: a value written inside a committing transaction is visible after.
1420        let outcome = key_store
1421            .transaction(|conn| {
1422                conn.key_store().write::<CURRENT_VERSION>(
1423                    crate::sql_key_store::COMMIT_LOG_SIGNER_PRIVATE_KEY,
1424                    &committed_key,
1425                    &value,
1426                )?;
1427                Ok::<_, StorageError>(Continue(()))
1428            })
1429            .unwrap();
1430        assert!(matches!(outcome, Continue(())));
1431        assert!(is_present(&committed_key), "commit must persist");
1432
1433        // Rollback: returns Ok(Rollback), NOT an Err, and the write is discarded.
1434        let outcome = key_store
1435            .transaction(|conn| {
1436                conn.key_store().write::<CURRENT_VERSION>(
1437                    crate::sql_key_store::COMMIT_LOG_SIGNER_PRIVATE_KEY,
1438                    &rolled_back_key,
1439                    &value,
1440                )?;
1441                Ok::<TransactionOutcome<()>, StorageError>(Rollback)
1442            })
1443            .unwrap();
1444        assert!(matches!(outcome, Rollback));
1445        assert!(!is_present(&rolled_back_key), "rollback must not persist");
1446
1447        // Real error: a closure returning Err propagates as Err (and rolls back).
1448        let result: Result<TransactionOutcome<()>, StorageError> =
1449            key_store.transaction(|_conn| Err(StorageError::DbSerialize));
1450        assert!(matches!(result, Err(StorageError::DbSerialize)));
1451    }
1452
1453    #[xmtp_common::test]
1454    async fn list_append_remove() {
1455        let store = crate::TestDb::create_persistent_store(None).await;
1456        let conn = store.conn();
1457        let mls_store = SqlKeyStore::new(conn);
1458        let provider = XmtpOpenMlsProvider::new(mls_store);
1459        let group_id = GroupId::random(provider.rand());
1460        let proposals = (0..10)
1461            .map(|i| Proposal(format!("TestProposal{i}").as_bytes().to_vec()))
1462            .collect::<Vec<_>>();
1463
1464        // Store proposals
1465        for (i, proposal) in proposals.iter().enumerate() {
1466            provider
1467                .storage()
1468                .queue_proposal::<GroupId, ProposalRef, Proposal>(
1469                    &group_id,
1470                    &ProposalRef(i),
1471                    proposal,
1472                )
1473                .expect("Failed to queue proposal");
1474        }
1475
1476        tracing::trace!("Finished with queued proposals");
1477        // Read proposal refs
1478        let proposal_refs_read: Vec<ProposalRef> = provider
1479            .storage()
1480            .queued_proposal_refs(&group_id)
1481            .expect("Failed to read proposal refs");
1482        assert_eq!(
1483            (0..10).map(ProposalRef).collect::<Vec<_>>(),
1484            proposal_refs_read
1485        );
1486
1487        // Read proposals
1488        let proposals_read: Vec<(ProposalRef, Proposal)> =
1489            provider.storage().queued_proposals(&group_id).unwrap();
1490        let proposals_expected: Vec<(ProposalRef, Proposal)> = (0..10)
1491            .map(ProposalRef)
1492            .zip(proposals.clone().into_iter())
1493            .collect();
1494        assert_eq!(proposals_expected, proposals_read);
1495
1496        // Remove proposal 5
1497        provider
1498            .storage()
1499            .remove_proposal(&group_id, &ProposalRef(5))
1500            .unwrap();
1501
1502        let proposal_refs_read: Vec<ProposalRef> =
1503            provider.storage().queued_proposal_refs(&group_id).unwrap();
1504        let mut expected = (0..10).map(ProposalRef).collect::<Vec<_>>();
1505        expected.remove(5);
1506        assert_eq!(expected, proposal_refs_read);
1507
1508        let proposals_read: Vec<(ProposalRef, Proposal)> =
1509            provider.storage().queued_proposals(&group_id).unwrap();
1510        let mut proposals_expected: Vec<(ProposalRef, Proposal)> = (0..10)
1511            .map(ProposalRef)
1512            .zip(proposals.clone().into_iter())
1513            .collect();
1514        proposals_expected.remove(5);
1515        assert_eq!(proposals_expected, proposals_read);
1516
1517        // Clear all proposals
1518        provider
1519            .storage()
1520            .clear_proposal_queue::<GroupId, ProposalRef>(&group_id)
1521            .unwrap();
1522        let proposal_refs_read: Result<Vec<ProposalRef>, SqlKeyStoreError> =
1523            provider.storage().queued_proposal_refs(&group_id);
1524        assert!(proposal_refs_read.unwrap().is_empty());
1525
1526        let proposals_read: Result<Vec<(ProposalRef, Proposal)>, SqlKeyStoreError> =
1527            provider.storage().queued_proposals(&group_id);
1528        assert!(proposals_read.unwrap().is_empty());
1529    }
1530
1531    #[xmtp_common::test]
1532    async fn group_state() {
1533        let store = crate::TestDb::create_persistent_store(None).await;
1534        let conn = store.conn();
1535        let store = SqlKeyStore::new(conn);
1536        let provider = XmtpOpenMlsProvider::new(store);
1537
1538        #[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Copy)]
1539        struct GroupState(usize);
1540        impl traits::GroupState<CURRENT_VERSION> for GroupState {}
1541        impl Entity<CURRENT_VERSION> for GroupState {}
1542
1543        let group_id = GroupId::random(provider.rand());
1544
1545        // Group state
1546        provider
1547            .storage()
1548            .write_group_state(&group_id, &GroupState(77))
1549            .unwrap();
1550
1551        // Read group state
1552        let group_state: Option<GroupState> = provider.storage().group_state(&group_id).unwrap();
1553        assert_eq!(GroupState(77), group_state.unwrap());
1554    }
1555
1556    #[xmtp_common::test]
1557    async fn application_export_tree() {
1558        let store = crate::TestDb::create_persistent_store(None).await;
1559        let conn = store.conn();
1560        let store = SqlKeyStore::new(conn);
1561        let provider = XmtpOpenMlsProvider::new(store);
1562
1563        #[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone)]
1564        struct AppExportTree {
1565            nodes: Vec<Vec<u8>>,
1566            leaf_count: u32,
1567            root_hash: [u8; 32],
1568        }
1569        impl traits::ApplicationExportTree<CURRENT_VERSION> for AppExportTree {}
1570        impl Entity<CURRENT_VERSION> for AppExportTree {}
1571
1572        fn random_tree(rand: &impl openmls_traits::random::OpenMlsRand) -> AppExportTree {
1573            let leaf_count = 3 + (rand.random_vec(1).unwrap()[0] % 10) as u32;
1574            let nodes: Vec<Vec<u8>> = (0..leaf_count)
1575                .map(|_| rand.random_vec(64).unwrap())
1576                .collect();
1577            AppExportTree {
1578                nodes,
1579                leaf_count,
1580                root_hash: xmtp_common::rand_array(),
1581            }
1582        }
1583
1584        let group_id_1 = GroupId::random(provider.rand());
1585        let group_id_2 = GroupId::random(provider.rand());
1586        let tree_1 = random_tree(provider.rand());
1587        let tree_2 = random_tree(provider.rand());
1588
1589        // Read before write should return None
1590        let result: Option<AppExportTree> = provider
1591            .storage()
1592            .application_export_tree(&group_id_1)
1593            .unwrap();
1594        assert!(result.is_none());
1595
1596        // Write tree for group 1
1597        provider
1598            .storage()
1599            .write_application_export_tree(&group_id_1, &tree_1)
1600            .unwrap();
1601
1602        // Read back group 1
1603        let result: Option<AppExportTree> = provider
1604            .storage()
1605            .application_export_tree(&group_id_1)
1606            .unwrap();
1607        assert_eq!(result.unwrap(), tree_1);
1608
1609        // Group 2 should still be None
1610        let result: Option<AppExportTree> = provider
1611            .storage()
1612            .application_export_tree(&group_id_2)
1613            .unwrap();
1614        assert!(result.is_none());
1615
1616        // Write tree for group 2
1617        provider
1618            .storage()
1619            .write_application_export_tree(&group_id_2, &tree_2)
1620            .unwrap();
1621        let result: Option<AppExportTree> = provider
1622            .storage()
1623            .application_export_tree(&group_id_2)
1624            .unwrap();
1625        assert_eq!(result.unwrap(), tree_2);
1626
1627        // Overwrite group 1 with new data
1628        let tree_1_updated = random_tree(provider.rand());
1629        provider
1630            .storage()
1631            .write_application_export_tree(&group_id_1, &tree_1_updated)
1632            .unwrap();
1633        let result: Option<AppExportTree> = provider
1634            .storage()
1635            .application_export_tree(&group_id_1)
1636            .unwrap();
1637        assert_eq!(result.unwrap(), tree_1_updated);
1638
1639        // Group 2 should be unaffected
1640        let result: Option<AppExportTree> = provider
1641            .storage()
1642            .application_export_tree(&group_id_2)
1643            .unwrap();
1644        assert_eq!(result.unwrap(), tree_2);
1645
1646        // Delete group 1
1647        provider
1648            .storage()
1649            .delete_application_export_tree::<GroupId, AppExportTree>(&group_id_1)
1650            .unwrap();
1651        let result: Option<AppExportTree> = provider
1652            .storage()
1653            .application_export_tree(&group_id_1)
1654            .unwrap();
1655        assert!(result.is_none());
1656
1657        // Group 2 should still exist
1658        let result: Option<AppExportTree> = provider
1659            .storage()
1660            .application_export_tree(&group_id_2)
1661            .unwrap();
1662        assert_eq!(result.unwrap(), tree_2);
1663
1664        // Delete group 2
1665        provider
1666            .storage()
1667            .delete_application_export_tree::<GroupId, AppExportTree>(&group_id_2)
1668            .unwrap();
1669        let result: Option<AppExportTree> = provider
1670            .storage()
1671            .application_export_tree(&group_id_2)
1672            .unwrap();
1673        assert!(result.is_none());
1674
1675        // Deleting again should not error
1676        provider
1677            .storage()
1678            .delete_application_export_tree::<GroupId, AppExportTree>(&group_id_1)
1679            .unwrap();
1680    }
1681}