penumbra_sdk_view/storage/
sct.rs

1use std::ops::Range;
2
3use anyhow::Context as _;
4use genawaiter::{rc::gen, yield_};
5use r2d2_sqlite::rusqlite::Transaction;
6
7use core::fmt::Debug;
8use penumbra_sdk_tct::{
9    storage::{Read, StoredPosition, Write},
10    structure::Hash,
11    Forgotten, Position, StateCommitment,
12};
13
14#[derive(Debug)]
15pub struct TreeStore<'a, 'c: 'a>(pub &'a mut Transaction<'c>);
16
17impl Read for TreeStore<'_, '_> {
18    type Error = anyhow::Error;
19
20    type HashesIter<'a>
21        = Box<dyn Iterator<Item = Result<(Position, u8, Hash), Self::Error>> + 'a>
22    where
23        Self: 'a;
24
25    type CommitmentsIter<'a>
26        = Box<dyn Iterator<Item = Result<(Position, StateCommitment), Self::Error>> + 'a>
27    where
28        Self: 'a;
29
30    fn position(&mut self) -> Result<StoredPosition, Self::Error> {
31        let mut stmt = self
32            .0
33            .prepare_cached("SELECT position FROM sct_position LIMIT 1")
34            .context("failed to prepare position query")?;
35        let position = stmt
36            .query_row::<Option<u64>, _, _>([], |row| row.get("position"))
37            .context("failed to query position")?
38            .map(Position::from)
39            .into();
40        Ok(position)
41    }
42
43    fn forgotten(&mut self) -> Result<Forgotten, Self::Error> {
44        let mut stmt = self
45            .0
46            .prepare_cached("SELECT forgotten FROM sct_forgotten LIMIT 1")
47            .context("failed to prepare forgotten query")?;
48        let forgotten = stmt
49            .query_row::<u64, _, _>([], |row| row.get("forgotten"))
50            .context("failed to query forgotten")?
51            .into();
52        Ok(forgotten)
53    }
54
55    fn hash(&mut self, position: Position, height: u8) -> Result<Option<Hash>, Self::Error> {
56        let position = u64::from(position) as i64;
57
58        let mut stmt = self
59            .0
60            .prepare_cached(
61                "SELECT hash FROM sct_hashes WHERE position = ?1 AND height = ?2 LIMIT 1",
62            )
63            .context("failed to prepare hash query")?;
64        let bytes = stmt
65            .query_row::<Option<Vec<u8>>, _, _>((&position, &height), |row| row.get("hash"))
66            .context("failed to query hash")?;
67
68        bytes
69            .map(|bytes| {
70                <[u8; 32]>::try_from(bytes)
71                    .map_err(|_| anyhow::anyhow!("hash was of incorrect length"))
72                    .and_then(|array| {
73                        if let Ok(hash) = Hash::from_bytes(array) {
74                            Ok(hash)
75                        } else {
76                            Err(anyhow::anyhow!("Failed to create Hash from bytes"))
77                        }
78                    })
79            })
80            .transpose()
81    }
82
83    fn hashes(&mut self) -> Self::HashesIter<'_> {
84        // The iterator has to *own* the stmt because the rows borrow from it, so we use the
85        // `genawaiter` crate to shove the entire preparation of the iterator into an (implicit)
86        // async block, which handles the desuguaring to properly own the stmt for us.
87        Box::new(
88            gen!({
89                let mut stmt = match self
90                    .0
91                    .prepare_cached("SELECT position, height, hash FROM sct_hashes")
92                    .context("failed to prepare hashes query")
93                {
94                    Ok(stmt) => stmt,
95                    // If an error happens while preparing the statement, shove it inside the first returned
96                    // item of the iterator, because we can't return an outer error:
97                    Err(e) => {
98                        yield_!(Err(e));
99                        return;
100                    }
101                };
102
103                let rows = match stmt
104                    .query_and_then([], |row| {
105                        let position: i64 = row.get("position")?;
106                        let height: u8 = row.get("height")?;
107                        let hash: Vec<u8> = row.get("hash")?;
108                        let hash = <[u8; 32]>::try_from(hash)
109                            .map_err(|_| anyhow::anyhow!("hash was of incorrect length"))
110                            .and_then(|array| {
111                                Hash::from_bytes(array).map_err(|e| {
112                                    // Explicitly convert any error to anyhow::Error
113                                    anyhow::Error::msg(format!("Error converting hash: {}", e))
114                                })
115                            })?;
116                        anyhow::Ok((Position::from(position as u64), height, hash))
117                    })
118                    .context("couldn't query database")
119                {
120                    Ok(rows) => rows,
121                    // If an error happens while querying the database, shove it inside the first
122                    // returned item of the iterator, because we can't return an outer error:
123                    Err(e) => {
124                        yield_!(Err(e));
125                        return;
126                    }
127                };
128
129                // Actually iterate through the rows:
130                for row in rows {
131                    yield_!(row);
132                }
133            })
134            .into_iter(),
135        )
136    }
137
138    fn commitment(&mut self, position: Position) -> Result<Option<StateCommitment>, Self::Error> {
139        let position = u64::from(position) as i64;
140
141        let mut stmt = self
142            .0
143            .prepare_cached("SELECT commitment FROM sct_commitments WHERE position = ?1 LIMIT 1")
144            .context("failed to prepare commitment query")?;
145
146        let bytes = stmt
147            .query_row::<Option<Vec<u8>>, _, _>((&position,), |row| row.get("commitment"))
148            .context("failed to query commitment")?;
149
150        bytes
151            .map(|bytes| {
152                <[u8; 32]>::try_from(bytes)
153                    .map_err(|_| anyhow::anyhow!("commitment was of incorrect length"))
154                    .and_then(|array| StateCommitment::try_from(array).map_err(Into::into))
155            })
156            .transpose()
157    }
158
159    fn commitments(&mut self) -> Self::CommitmentsIter<'_> {
160        // The iterator has to *own* the stmt because the rows borrow from it, so we use the
161        // `genawaiter` crate to shove the entire preparation of the iterator into an (implicit)
162        // async block, which handles the desuguaring to properly own the stmt for us.
163        Box::new(
164            gen!({
165                let mut stmt = match self
166                    .0
167                    .prepare_cached("SELECT position, commitment FROM sct_commitments")
168                    .context("failed to prepare commitments query")
169                {
170                    Ok(stmt) => stmt,
171                    // If an error happens while preparing the statement, shove it inside the first returned
172                    // item of the iterator, because we can't return an outer error:
173                    Err(e) => {
174                        yield_!(Err(e));
175                        return;
176                    }
177                };
178
179                let rows = match stmt
180                    .query_and_then([], |row| {
181                        let position: i64 = row.get("position")?;
182                        let commitment: Vec<u8> = row.get("commitment")?;
183                        let commitment = <[u8; 32]>::try_from(commitment)
184                            .map_err(|_| anyhow::anyhow!("commitment was of incorrect length"))
185                            .and_then(|array| {
186                                StateCommitment::try_from(array).map_err(Into::into)
187                            })?;
188                        anyhow::Ok((Position::from(position as u64), commitment))
189                    })
190                    .context("couldn't query database")
191                {
192                    Ok(rows) => rows,
193                    // If an error happens while querying the database, shove it inside the first
194                    // returned item of the iterator, because we can't return an outer error:
195                    Err(e) => {
196                        yield_!(Err(e));
197                        return;
198                    }
199                };
200
201                // Actually iterate through the rows:
202                for row in rows {
203                    yield_!(row);
204                }
205            })
206            .into_iter(),
207        )
208    }
209}
210
211impl Write for TreeStore<'_, '_> {
212    fn set_position(&mut self, position: StoredPosition) -> Result<(), Self::Error> {
213        let position = Option::from(position).map(|p: Position| u64::from(p) as i64);
214
215        self.0
216            .prepare_cached("UPDATE sct_position SET position = ?1")
217            .context("failed to prepare position update")?
218            .execute([&position])?;
219
220        Ok(())
221    }
222
223    fn set_forgotten(&mut self, forgotten: Forgotten) -> Result<(), Self::Error> {
224        let forgotten = u64::from(forgotten) as i64;
225
226        self.0
227            .prepare_cached("UPDATE sct_forgotten SET forgotten = ?1")
228            .context("failed to prepare forgotten update")?
229            .execute([&forgotten])?;
230
231        Ok(())
232    }
233
234    fn add_hash(
235        &mut self,
236        position: Position,
237        height: u8,
238        hash: Hash,
239        _essential: bool,
240    ) -> Result<(), Self::Error> {
241        let position = u64::from(position) as i64;
242        let hash = hash.to_bytes().to_vec();
243
244        self.0.prepare_cached(
245            "INSERT INTO sct_hashes (position, height, hash) VALUES (?1, ?2, ?3) ON CONFLICT DO NOTHING"
246        ).context("failed to prepare hash insert")?
247            .execute((&position, &height, &hash))
248            .context("failed to insert hash")?;
249
250        Ok(())
251    }
252
253    fn add_commitment(
254        &mut self,
255        position: Position,
256        commitment: StateCommitment,
257    ) -> Result<(), Self::Error> {
258        let position = u64::from(position) as i64;
259        let commitment = <[u8; 32]>::from(commitment).to_vec();
260
261        self.0.prepare_cached(
262            "INSERT INTO sct_commitments (position, commitment) VALUES (?1, ?2) ON CONFLICT DO NOTHING"
263        ).context("failed to prepare commitment insert")?
264            .execute((&position, &commitment))
265            .context("failed to insert commitment")?;
266
267        Ok(())
268    }
269
270    fn delete_range(
271        &mut self,
272        below_height: u8,
273        positions: Range<Position>,
274    ) -> Result<(), Self::Error> {
275        let start = u64::from(positions.start) as i64;
276        let end = u64::from(positions.end) as i64;
277
278        self.0
279            .prepare_cached(
280                "DELETE FROM sct_hashes WHERE position >= ?1 AND position < ?2 AND height < ?3",
281            )
282            .context("failed to prepare hash delete")?
283            .execute((&start, &end, &below_height))
284            .context("failed to delete hashes")?;
285
286        Ok(())
287    }
288}
289
290#[cfg(test)]
291mod test {
292    use super::*;
293
294    use penumbra_sdk_tct::{StateCommitment, Witness};
295
296    #[test]
297    fn tree_store_spot_check() {
298        // Set up the database:
299        let mut db = r2d2_sqlite::rusqlite::Connection::open_in_memory().unwrap();
300        let mut tx = db.transaction().unwrap();
301        tx.execute_batch(include_str!("schema.sql")).unwrap();
302
303        // Now we're exclusively going to talk to the db through the TreeStore:
304        let mut store = TreeStore(&mut tx);
305
306        // Check that the currently stored tree is the empty tree:
307        let deserialized = penumbra_sdk_tct::Tree::from_reader(&mut store).unwrap();
308        assert_eq!(deserialized, penumbra_sdk_tct::Tree::new());
309
310        // Make some kind of tree:
311        let mut tree = penumbra_sdk_tct::Tree::new();
312        tree.insert(Witness::Keep, StateCommitment::try_from([0; 32]).unwrap())
313            .unwrap();
314        tree.end_block().unwrap();
315        tree.insert(Witness::Forget, StateCommitment::try_from([1; 32]).unwrap())
316            .unwrap();
317        tree.end_epoch().unwrap();
318        tree.insert(Witness::Keep, StateCommitment::try_from([2; 32]).unwrap())
319            .unwrap();
320
321        // Write the tree to the database:
322        tree.to_writer(&mut store).unwrap();
323
324        // Read the tree back from the database:
325        let deserialized = penumbra_sdk_tct::Tree::from_reader(&mut store).unwrap();
326
327        assert_eq!(tree, deserialized);
328    }
329}
penumbra_sdk_view/storage/sct.rs

penumbra_sdk_view/storage/
sct.rs
