mithril_signer/store/
mktree_store_sqlite.rs

1use std::{iter::repeat_n, sync::Arc};
2
3use anyhow::Context;
4use mithril_common::{
5    crypto_helper::{Bytes, MKTreeLeafIndexer, MKTreeLeafPosition, MKTreeNode, MKTreeStorer},
6    StdResult,
7};
8
9/// A Merkle tree store with Sqlite backend
10/// * This store is used to store the Merkle tree nodes in a Sqlite database with in memory mode.
11/// * This store is not suited for proving, it is only used to store the compute the root of a Merkle tree.
12/// * This store is slower than the in memory store but uses less memory which is important to minimize the signer footprint on the SPO infrastructure.
13pub struct MKTreeStoreSqlite {
14    inner_store: sqlite::ConnectionThreadSafe,
15}
16
17impl MKTreeStoreSqlite {
18    fn build() -> StdResult<Self> {
19        Ok(Self {
20            inner_store: Self::create_connection()?,
21        })
22    }
23
24    fn create_connection() -> StdResult<sqlite::ConnectionThreadSafe> {
25        let connection = sqlite::Connection::open_thread_safe(":memory:")?;
26        connection.execute("pragma shrink_memory; pragma temp_store = FILE;")?;
27        connection.execute(
28            "create table merkle_tree (
29                position integer, 
30                element blob, 
31                primary key (position)
32            )",
33        )?;
34
35        Ok(connection)
36    }
37
38    fn get_element_at_position(&self, position: u64) -> StdResult<Option<Arc<MKTreeNode>>> {
39        let query = "SELECT element FROM merkle_tree WHERE position = ?";
40        let mut statement = self.inner_store.prepare(query)?;
41        statement.bind((1, position as i64)).unwrap();
42        let result = if let Ok(sqlite::State::Row) = statement.next() {
43            Some(Arc::new(MKTreeNode::new(
44                statement.read::<Bytes, _>("element")?,
45            )))
46        } else {
47            None
48        };
49
50        Ok(result)
51    }
52
53    fn insert_elements_from_position(
54        &self,
55        position: u64,
56        elements: Vec<Arc<MKTreeNode>>,
57    ) -> StdResult<()> {
58        let values_columns: Vec<&str> = repeat_n("(?, ?)", elements.len()).collect();
59        let values: Vec<sqlite::Value> = elements
60            .into_iter()
61            .enumerate()
62            .flat_map(|(i, elem)| {
63                vec![
64                    sqlite::Value::Integer((position + i as u64) as i64),
65                    sqlite::Value::Binary((**elem).to_vec()),
66                ]
67            })
68            .collect();
69        let query = format!(
70            "INSERT INTO merkle_tree(position, element) VALUES {}",
71            values_columns.join(", ")
72        );
73        let mut statement = self.inner_store.prepare(query)?;
74        statement.bind::<&[(_, sqlite::Value)]>(
75            values
76                .into_iter()
77                .enumerate()
78                .map(|(i, v)| (i + 1, v))
79                .collect::<Vec<_>>()
80                .as_slice(),
81        )?;
82        statement.next()?;
83
84        Ok(())
85    }
86}
87
88impl Clone for MKTreeStoreSqlite {
89    fn clone(&self) -> Self {
90        unimplemented!("Clone is not implemented for MKTreeStoreSqlite")
91    }
92}
93
94impl MKTreeLeafIndexer for MKTreeStoreSqlite {
95    fn set_leaf_position(&self, _pos: MKTreeLeafPosition, _node: Arc<MKTreeNode>) -> StdResult<()> {
96        Ok(())
97    }
98
99    fn get_leaf_position(&self, _node: &MKTreeNode) -> Option<MKTreeLeafPosition> {
100        unimplemented!("get_leaf_position is not implemented for MKTreeStoreSqlite")
101    }
102
103    fn total_leaves(&self) -> usize {
104        unimplemented!("total_leaves is not implemented for MKTreeStoreSqlite")
105    }
106
107    fn leaves(&self) -> Vec<MKTreeNode> {
108        unimplemented!("leaves is not implemented for MKTreeStoreSqlite")
109    }
110}
111
112impl MKTreeStorer for MKTreeStoreSqlite {
113    fn build() -> StdResult<Self> {
114        Self::build()
115    }
116
117    fn get_elem(&self, pos: u64) -> StdResult<Option<Arc<MKTreeNode>>> {
118        self.get_element_at_position(pos).with_context(|| {
119            format!("MKTreeStoreSqlite failed to retrieve element at position {pos}")
120        })
121    }
122
123    fn append(&self, pos: u64, elems: Vec<Arc<MKTreeNode>>) -> StdResult<()> {
124        self.insert_elements_from_position(pos, elems)
125            .with_context(|| {
126                format!("MKTreeStoreSqlite failed to insert elements from position {pos}")
127            })
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use mithril_common::crypto_helper::MKTree;
134
135    use super::*;
136
137    #[test]
138    fn test_golden_merkle_root() {
139        let leaves = vec!["golden-1", "golden-2", "golden-3", "golden-4", "golden-5"];
140        let mktree =
141            MKTree::<MKTreeStoreSqlite>::new(&leaves).expect("MKTree creation should not fail");
142        let mkroot = mktree
143            .compute_root()
144            .expect("MKRoot generation should not fail");
145
146        assert_eq!(
147            "3bbced153528697ecde7345a22e50115306478353619411523e804f2323fd921",
148            mkroot.to_hex()
149        );
150    }
151}
mithril_signer/store/mktree_store_sqlite.rs

mithril_signer/store/
mktree_store_sqlite.rs
