jitash_bdk/database/
memory.rs

1// Bitcoin Dev Kit
2// Written in 2020 by Alekos Filini <alekos.filini@gmail.com>
3//
4// Copyright (c) 2020-2021 Bitcoin Dev Kit Developers
5//
6// This file is licensed under the Apache License, Version 2.0 <LICENSE-APACHE
7// or http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
8// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your option.
9// You may not use this file except in accordance with one or both of these
10// licenses.
11
12//! In-memory ephemeral database
13//!
14//! This module defines an in-memory database type called [`MemoryDatabase`] that is based on a
15//! [`BTreeMap`].
16
17use std::any::Any;
18use std::collections::BTreeMap;
19use std::ops::Bound::{Excluded, Included};
20
21use bitcoin::consensus::encode::{deserialize, serialize};
22use bitcoin::hash_types::Txid;
23use bitcoin::{OutPoint, Script, Transaction};
24
25use crate::database::{BatchDatabase, BatchOperations, ConfigurableDatabase, Database, SyncTime};
26use crate::error::Error;
27use crate::types::*;
28
29// path -> script       p{i,e}<path> -> script
30// script -> path       s<script> -> {i,e}<path>
31// outpoint             u<outpoint> -> txout
32// rawtx                r<txid> -> tx
33// transactions         t<txid> -> tx details
34// deriv indexes        c{i,e} -> u32
35// descriptor checksum  d{i,e} -> vec<u8>
36// last sync time       l -> { height, timestamp }
37
38pub(crate) enum MapKey<'a> {
39    Path((Option<KeychainKind>, Option<u32>)),
40    Script(Option<&'a Script>),
41    Utxo(Option<&'a OutPoint>),
42    RawTx(Option<&'a Txid>),
43    Transaction(Option<&'a Txid>),
44    LastIndex(KeychainKind),
45    SyncTime,
46    DescriptorChecksum(KeychainKind),
47}
48
49impl MapKey<'_> {
50    fn as_prefix(&self) -> Vec<u8> {
51        match self {
52            MapKey::Path((st, _)) => {
53                let mut v = b"p".to_vec();
54                if let Some(st) = st {
55                    v.push(st.as_byte());
56                }
57                v
58            }
59            MapKey::Script(_) => b"s".to_vec(),
60            MapKey::Utxo(_) => b"u".to_vec(),
61            MapKey::RawTx(_) => b"r".to_vec(),
62            MapKey::Transaction(_) => b"t".to_vec(),
63            MapKey::LastIndex(st) => [b"c", st.as_ref()].concat(),
64            MapKey::SyncTime => b"l".to_vec(),
65            MapKey::DescriptorChecksum(st) => [b"d", st.as_ref()].concat(),
66        }
67    }
68
69    fn serialize_content(&self) -> Vec<u8> {
70        match self {
71            MapKey::Path((_, Some(child))) => child.to_be_bytes().to_vec(),
72            MapKey::Script(Some(s)) => serialize(*s),
73            MapKey::Utxo(Some(s)) => serialize(*s),
74            MapKey::RawTx(Some(s)) => serialize(*s),
75            MapKey::Transaction(Some(s)) => serialize(*s),
76            _ => vec![],
77        }
78    }
79
80    pub fn as_map_key(&self) -> Vec<u8> {
81        let mut v = self.as_prefix();
82        v.extend_from_slice(&self.serialize_content());
83
84        v
85    }
86}
87
88fn after(key: &[u8]) -> Vec<u8> {
89    let mut key = key.to_owned();
90    let mut idx = key.len();
91    while idx > 0 {
92        if key[idx - 1] == 0xFF {
93            idx -= 1;
94            continue;
95        } else {
96            key[idx - 1] += 1;
97            break;
98        }
99    }
100
101    key
102}
103
104/// In-memory ephemeral database
105///
106/// This database can be used as a temporary storage for wallets that are not kept permanently on
107/// a device, or on platforms that don't provide a filesystem, like `wasm32`.
108///
109/// Once it's dropped its content will be lost.
110///
111/// If you are looking for a permanent storage solution, you can try with the default key-value
112/// database called [`sled`]. See the [`database`] module documentation for more details.
113///
114/// [`database`]: crate::database
115#[derive(Debug, Default)]
116pub struct MemoryDatabase {
117    map: BTreeMap<Vec<u8>, Box<dyn Any + Send + Sync>>,
118    deleted_keys: Vec<Vec<u8>>,
119}
120
121impl MemoryDatabase {
122    /// Create a new empty database
123    pub fn new() -> Self {
124        MemoryDatabase {
125            map: BTreeMap::new(),
126            deleted_keys: Vec::new(),
127        }
128    }
129}
130
131impl BatchOperations for MemoryDatabase {
132    fn set_script_pubkey(
133        &mut self,
134        script: &Script,
135        keychain: KeychainKind,
136        path: u32,
137    ) -> Result<(), Error> {
138        let key = MapKey::Path((Some(keychain), Some(path))).as_map_key();
139        self.map.insert(key, Box::new(script.clone()));
140
141        let key = MapKey::Script(Some(script)).as_map_key();
142        let value = json!({
143            "t": keychain,
144            "p": path,
145        });
146        self.map.insert(key, Box::new(value));
147
148        Ok(())
149    }
150
151    fn set_utxo(&mut self, utxo: &LocalUtxo) -> Result<(), Error> {
152        let key = MapKey::Utxo(Some(&utxo.outpoint)).as_map_key();
153        self.map.insert(
154            key,
155            Box::new((utxo.txout.clone(), utxo.keychain, utxo.is_spent)),
156        );
157
158        Ok(())
159    }
160    fn set_raw_tx(&mut self, transaction: &Transaction) -> Result<(), Error> {
161        let key = MapKey::RawTx(Some(&transaction.txid())).as_map_key();
162        self.map.insert(key, Box::new(transaction.clone()));
163
164        Ok(())
165    }
166    fn set_tx(&mut self, transaction: &TransactionDetails) -> Result<(), Error> {
167        let key = MapKey::Transaction(Some(&transaction.txid)).as_map_key();
168
169        // insert the raw_tx if present
170        if let Some(ref tx) = transaction.transaction {
171            self.set_raw_tx(tx)?;
172        }
173
174        // remove the raw tx from the serialized version
175        let mut transaction = transaction.clone();
176        transaction.transaction = None;
177
178        self.map.insert(key, Box::new(transaction));
179
180        Ok(())
181    }
182    fn set_last_index(&mut self, keychain: KeychainKind, value: u32) -> Result<(), Error> {
183        let key = MapKey::LastIndex(keychain).as_map_key();
184        self.map.insert(key, Box::new(value));
185
186        Ok(())
187    }
188    fn set_sync_time(&mut self, data: SyncTime) -> Result<(), Error> {
189        let key = MapKey::SyncTime.as_map_key();
190        self.map.insert(key, Box::new(data));
191
192        Ok(())
193    }
194
195    fn del_script_pubkey_from_path(
196        &mut self,
197        keychain: KeychainKind,
198        path: u32,
199    ) -> Result<Option<Script>, Error> {
200        let key = MapKey::Path((Some(keychain), Some(path))).as_map_key();
201        let res = self.map.remove(&key);
202        self.deleted_keys.push(key);
203
204        Ok(res.map(|x| x.downcast_ref().cloned().unwrap()))
205    }
206    fn del_path_from_script_pubkey(
207        &mut self,
208        script: &Script,
209    ) -> Result<Option<(KeychainKind, u32)>, Error> {
210        let key = MapKey::Script(Some(script)).as_map_key();
211        let res = self.map.remove(&key);
212        self.deleted_keys.push(key);
213
214        match res {
215            None => Ok(None),
216            Some(b) => {
217                let mut val: serde_json::Value = b.downcast_ref().cloned().unwrap();
218                let st = serde_json::from_value(val["t"].take())?;
219                let path = serde_json::from_value(val["p"].take())?;
220
221                Ok(Some((st, path)))
222            }
223        }
224    }
225    fn del_utxo(&mut self, outpoint: &OutPoint) -> Result<Option<LocalUtxo>, Error> {
226        let key = MapKey::Utxo(Some(outpoint)).as_map_key();
227        let res = self.map.remove(&key);
228        self.deleted_keys.push(key);
229
230        match res {
231            None => Ok(None),
232            Some(b) => {
233                let (txout, keychain, is_spent) = b.downcast_ref().cloned().unwrap();
234                Ok(Some(LocalUtxo {
235                    outpoint: *outpoint,
236                    txout,
237                    keychain,
238                    is_spent,
239                }))
240            }
241        }
242    }
243    fn del_raw_tx(&mut self, txid: &Txid) -> Result<Option<Transaction>, Error> {
244        let key = MapKey::RawTx(Some(txid)).as_map_key();
245        let res = self.map.remove(&key);
246        self.deleted_keys.push(key);
247
248        Ok(res.map(|x| x.downcast_ref().cloned().unwrap()))
249    }
250    fn del_tx(
251        &mut self,
252        txid: &Txid,
253        include_raw: bool,
254    ) -> Result<Option<TransactionDetails>, Error> {
255        let raw_tx = if include_raw {
256            self.del_raw_tx(txid)?
257        } else {
258            None
259        };
260
261        let key = MapKey::Transaction(Some(txid)).as_map_key();
262        let res = self.map.remove(&key);
263        self.deleted_keys.push(key);
264
265        match res {
266            None => Ok(None),
267            Some(b) => {
268                let mut val: TransactionDetails = b.downcast_ref().cloned().unwrap();
269                val.transaction = raw_tx;
270
271                Ok(Some(val))
272            }
273        }
274    }
275    fn del_last_index(&mut self, keychain: KeychainKind) -> Result<Option<u32>, Error> {
276        let key = MapKey::LastIndex(keychain).as_map_key();
277        let res = self.map.remove(&key);
278        self.deleted_keys.push(key);
279
280        match res {
281            None => Ok(None),
282            Some(b) => Ok(Some(*b.downcast_ref().unwrap())),
283        }
284    }
285    fn del_sync_time(&mut self) -> Result<Option<SyncTime>, Error> {
286        let key = MapKey::SyncTime.as_map_key();
287        let res = self.map.remove(&key);
288        self.deleted_keys.push(key);
289
290        Ok(res.map(|b| b.downcast_ref().cloned().unwrap()))
291    }
292}
293
294impl Database for MemoryDatabase {
295    fn check_descriptor_checksum<B: AsRef<[u8]>>(
296        &mut self,
297        keychain: KeychainKind,
298        bytes: B,
299    ) -> Result<(), Error> {
300        let key = MapKey::DescriptorChecksum(keychain).as_map_key();
301
302        let prev = self
303            .map
304            .get(&key)
305            .map(|x| x.downcast_ref::<Vec<u8>>().unwrap());
306        if let Some(val) = prev {
307            if val == &bytes.as_ref().to_vec() {
308                Ok(())
309            } else {
310                Err(Error::ChecksumMismatch)
311            }
312        } else {
313            self.map.insert(key, Box::new(bytes.as_ref().to_vec()));
314            Ok(())
315        }
316    }
317
318    fn iter_script_pubkeys(&self, keychain: Option<KeychainKind>) -> Result<Vec<Script>, Error> {
319        let key = MapKey::Path((keychain, None)).as_map_key();
320        self.map
321            .range::<Vec<u8>, _>((Included(&key), Excluded(&after(&key))))
322            .map(|(_, v)| Ok(v.downcast_ref().cloned().unwrap()))
323            .collect()
324    }
325
326    fn iter_utxos(&self) -> Result<Vec<LocalUtxo>, Error> {
327        let key = MapKey::Utxo(None).as_map_key();
328        self.map
329            .range::<Vec<u8>, _>((Included(&key), Excluded(&after(&key))))
330            .map(|(k, v)| {
331                let outpoint = deserialize(&k[1..]).unwrap();
332                let (txout, keychain, is_spent) = v.downcast_ref().cloned().unwrap();
333                Ok(LocalUtxo {
334                    outpoint,
335                    txout,
336                    keychain,
337                    is_spent,
338                })
339            })
340            .collect()
341    }
342
343    fn iter_raw_txs(&self) -> Result<Vec<Transaction>, Error> {
344        let key = MapKey::RawTx(None).as_map_key();
345        self.map
346            .range::<Vec<u8>, _>((Included(&key), Excluded(&after(&key))))
347            .map(|(_, v)| Ok(v.downcast_ref().cloned().unwrap()))
348            .collect()
349    }
350
351    fn iter_txs(&self, include_raw: bool) -> Result<Vec<TransactionDetails>, Error> {
352        let key = MapKey::Transaction(None).as_map_key();
353        self.map
354            .range::<Vec<u8>, _>((Included(&key), Excluded(&after(&key))))
355            .map(|(k, v)| {
356                let mut txdetails: TransactionDetails = v.downcast_ref().cloned().unwrap();
357                if include_raw {
358                    let txid = deserialize(&k[1..])?;
359                    txdetails.transaction = self.get_raw_tx(&txid)?;
360                }
361
362                Ok(txdetails)
363            })
364            .collect()
365    }
366
367    fn get_script_pubkey_from_path(
368        &self,
369        keychain: KeychainKind,
370        path: u32,
371    ) -> Result<Option<Script>, Error> {
372        let key = MapKey::Path((Some(keychain), Some(path))).as_map_key();
373        Ok(self
374            .map
375            .get(&key)
376            .map(|b| b.downcast_ref().cloned().unwrap()))
377    }
378
379    fn get_path_from_script_pubkey(
380        &self,
381        script: &Script,
382    ) -> Result<Option<(KeychainKind, u32)>, Error> {
383        let key = MapKey::Script(Some(script)).as_map_key();
384        Ok(self.map.get(&key).map(|b| {
385            let mut val: serde_json::Value = b.downcast_ref().cloned().unwrap();
386            let st = serde_json::from_value(val["t"].take()).unwrap();
387            let path = serde_json::from_value(val["p"].take()).unwrap();
388
389            (st, path)
390        }))
391    }
392
393    fn get_utxo(&self, outpoint: &OutPoint) -> Result<Option<LocalUtxo>, Error> {
394        let key = MapKey::Utxo(Some(outpoint)).as_map_key();
395        Ok(self.map.get(&key).map(|b| {
396            let (txout, keychain, is_spent) = b.downcast_ref().cloned().unwrap();
397            LocalUtxo {
398                outpoint: *outpoint,
399                txout,
400                keychain,
401                is_spent,
402            }
403        }))
404    }
405
406    fn get_raw_tx(&self, txid: &Txid) -> Result<Option<Transaction>, Error> {
407        let key = MapKey::RawTx(Some(txid)).as_map_key();
408        Ok(self
409            .map
410            .get(&key)
411            .map(|b| b.downcast_ref().cloned().unwrap()))
412    }
413
414    fn get_tx(&self, txid: &Txid, include_raw: bool) -> Result<Option<TransactionDetails>, Error> {
415        let key = MapKey::Transaction(Some(txid)).as_map_key();
416        Ok(self.map.get(&key).map(|b| {
417            let mut txdetails: TransactionDetails = b.downcast_ref().cloned().unwrap();
418            if include_raw {
419                txdetails.transaction = self.get_raw_tx(txid).unwrap();
420            }
421
422            txdetails
423        }))
424    }
425
426    fn get_last_index(&self, keychain: KeychainKind) -> Result<Option<u32>, Error> {
427        let key = MapKey::LastIndex(keychain).as_map_key();
428        Ok(self.map.get(&key).map(|b| *b.downcast_ref().unwrap()))
429    }
430
431    fn get_sync_time(&self) -> Result<Option<SyncTime>, Error> {
432        let key = MapKey::SyncTime.as_map_key();
433        Ok(self
434            .map
435            .get(&key)
436            .map(|b| b.downcast_ref().cloned().unwrap()))
437    }
438
439    // inserts 0 if not present
440    fn increment_last_index(&mut self, keychain: KeychainKind) -> Result<u32, Error> {
441        let key = MapKey::LastIndex(keychain).as_map_key();
442        let value = self
443            .map
444            .entry(key)
445            .and_modify(|x| *x.downcast_mut::<u32>().unwrap() += 1)
446            .or_insert_with(|| Box::<u32>::new(0))
447            .downcast_mut()
448            .unwrap();
449
450        Ok(*value)
451    }
452}
453
454impl BatchDatabase for MemoryDatabase {
455    type Batch = Self;
456
457    fn begin_batch(&self) -> Self::Batch {
458        MemoryDatabase::new()
459    }
460
461    fn commit_batch(&mut self, mut batch: Self::Batch) -> Result<(), Error> {
462        for key in batch.deleted_keys.iter() {
463            self.map.remove(key);
464        }
465        self.map.append(&mut batch.map);
466        Ok(())
467    }
468}
469
470impl ConfigurableDatabase for MemoryDatabase {
471    type Config = ();
472
473    fn from_config(_config: &Self::Config) -> Result<Self, Error> {
474        Ok(MemoryDatabase::default())
475    }
476}
477
478#[macro_export]
479#[doc(hidden)]
480/// Artificially insert a tx in the database, as if we had found it with a `sync`. This is a hidden
481/// macro and not a `[cfg(test)]` function so it can be called within the context of doctests which
482/// don't have `test` set.
483macro_rules! populate_test_db {
484    ($db:expr, $tx_meta:expr, $current_height:expr$(,)?) => {{
485        $crate::populate_test_db!($db, $tx_meta, $current_height, (@coinbase false))
486    }};
487    ($db:expr, $tx_meta:expr, $current_height:expr, (@coinbase $is_coinbase:expr)$(,)?) => {{
488        use std::str::FromStr;
489        use $crate::database::SyncTime;
490        use $crate::database::{BatchOperations, Database};
491        let mut db = $db;
492        let tx_meta = $tx_meta;
493        let current_height: Option<u32> = $current_height;
494        let mut input = vec![$crate::bitcoin::TxIn::default()];
495        if !$is_coinbase {
496            input[0].previous_output.vout = 0;
497        }
498        let tx = $crate::bitcoin::Transaction {
499            version: 1,
500            lock_time: bitcoin::PackedLockTime(0),
501            input,
502            output: tx_meta
503                .output
504                .iter()
505                .map(|out_meta| $crate::bitcoin::TxOut {
506                    value: out_meta.value,
507                    script_pubkey: $crate::bitcoin::Address::from_str(&out_meta.to_address)
508                        .unwrap()
509                        .script_pubkey(),
510                })
511                .collect(),
512        };
513
514        let txid = tx.txid();
515        // Set Confirmation time only if current height is provided.
516        // panics if `tx_meta.min_confirmation` is Some, and current_height is None.
517        let confirmation_time = tx_meta
518            .min_confirmations
519            .and_then(|v| if v == 0 { None } else { Some(v) })
520            .map(|conf| $crate::BlockTime {
521                height: current_height.expect("Current height is needed for testing transaction with min-confirmation values").checked_sub(conf as u32).unwrap() + 1,
522                timestamp: 0,
523            });
524
525        // Set the database sync_time.
526        // Check if the current_height is less than already known sync height, apply the max
527        // If any of them is None, the other will be applied instead.
528        // If both are None, this will not be set.
529        if let Some(height) = db.get_sync_time().unwrap()
530                                .map(|sync_time| sync_time.block_time.height)
531                                .max(current_height) {
532            let sync_time = SyncTime {
533                block_time: BlockTime {
534                    height,
535                    timestamp: 0
536                }
537            };
538            db.set_sync_time(sync_time).unwrap();
539        }
540
541        let tx_details = $crate::TransactionDetails {
542            transaction: Some(tx.clone()),
543            txid,
544            fee: Some(0),
545            received: 0,
546            sent: 0,
547            confirmation_time,
548        };
549
550        db.set_tx(&tx_details).unwrap();
551        for (vout, out) in tx.output.iter().enumerate() {
552            db.set_utxo(&$crate::LocalUtxo {
553                txout: out.clone(),
554                outpoint: $crate::bitcoin::OutPoint {
555                    txid,
556                    vout: vout as u32,
557                },
558                keychain: $crate::KeychainKind::External,
559                is_spent: false,
560            })
561            .unwrap();
562        }
563
564        txid
565    }};
566}
567
568#[macro_export]
569#[doc(hidden)]
570/// Macro for getting a wallet for use in a doctest
571macro_rules! doctest_wallet {
572    () => {{
573        use $crate::bitcoin::Network;
574        use $crate::database::MemoryDatabase;
575        use $crate::testutils;
576        let descriptor = "wpkh(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW)";
577        let descriptors = testutils!(@descriptors (descriptor) (descriptor));
578
579        let mut db = MemoryDatabase::new();
580        let txid = populate_test_db!(
581            &mut db,
582            testutils! {
583                @tx ( (@external descriptors, 0) => 500_000 ) (@confirmations 1)
584            },
585            Some(100),
586        );
587
588        $crate::Wallet::new(
589            &descriptors.0,
590            descriptors.1.as_ref(),
591            Network::Regtest,
592            db
593        )
594        .unwrap()
595    }}
596}
597
598#[cfg(test)]
599mod test {
600    use super::MemoryDatabase;
601
602    fn get_tree() -> MemoryDatabase {
603        MemoryDatabase::new()
604    }
605
606    #[test]
607    fn test_script_pubkey() {
608        crate::database::test::test_script_pubkey(get_tree());
609    }
610
611    #[test]
612    fn test_batch_script_pubkey() {
613        crate::database::test::test_batch_script_pubkey(get_tree());
614    }
615
616    #[test]
617    fn test_iter_script_pubkey() {
618        crate::database::test::test_iter_script_pubkey(get_tree());
619    }
620
621    #[test]
622    fn test_del_script_pubkey() {
623        crate::database::test::test_del_script_pubkey(get_tree());
624    }
625
626    #[test]
627    fn test_utxo() {
628        crate::database::test::test_utxo(get_tree());
629    }
630
631    #[test]
632    fn test_raw_tx() {
633        crate::database::test::test_raw_tx(get_tree());
634    }
635
636    #[test]
637    fn test_tx() {
638        crate::database::test::test_tx(get_tree());
639    }
640
641    #[test]
642    fn test_last_index() {
643        crate::database::test::test_last_index(get_tree());
644    }
645
646    #[test]
647    fn test_sync_time() {
648        crate::database::test::test_sync_time(get_tree());
649    }
650
651    #[test]
652    fn test_iter_raw_txs() {
653        crate::database::test::test_iter_raw_txs(get_tree());
654    }
655
656    #[test]
657    fn test_del_path_from_script_pubkey() {
658        crate::database::test::test_del_path_from_script_pubkey(get_tree());
659    }
660
661    #[test]
662    fn test_iter_script_pubkeys() {
663        crate::database::test::test_iter_script_pubkeys(get_tree());
664    }
665
666    #[test]
667    fn test_del_utxo() {
668        crate::database::test::test_del_utxo(get_tree());
669    }
670
671    #[test]
672    fn test_del_raw_tx() {
673        crate::database::test::test_del_raw_tx(get_tree());
674    }
675
676    #[test]
677    fn test_del_tx() {
678        crate::database::test::test_del_tx(get_tree());
679    }
680
681    #[test]
682    fn test_del_last_index() {
683        crate::database::test::test_del_last_index(get_tree());
684    }
685
686    #[test]
687    fn test_check_descriptor_checksum() {
688        crate::database::test::test_check_descriptor_checksum(get_tree());
689    }
690}
jitash_bdk/database/memory.rs

jitash_bdk/database/
memory.rs
