From 6ab29fcec5bf318d4d57f54b6f89cd791ed1c92e Mon Sep 17 00:00:00 2001
From: Alekos Filini <alekos.filini@gmail.com>
Date: Wed, 5 Feb 2020 11:59:02 +0100
Subject: [PATCH] Add a generalized database trait and a Sled-based
 implementation

---
 core/lib/Cargo.toml               |   7 +
 core/lib/src/database/keyvalue.rs | 390 ++++++++++++++++++++++++++++++
 core/lib/src/database/mod.rs      |  71 ++++++
 core/lib/src/error.rs             |  12 +
 core/lib/src/lib.rs               |   6 +
 core/lib/src/types.rs             |  47 ++++
 6 files changed, 533 insertions(+)
 create mode 100644 core/lib/src/database/keyvalue.rs
 create mode 100644 core/lib/src/database/mod.rs
 create mode 100644 core/lib/src/types.rs

diff --git a/core/lib/Cargo.toml b/core/lib/Cargo.toml
index 6894d5e9..1f0b2c2d 100644
--- a/core/lib/Cargo.toml
+++ b/core/lib/Cargo.toml
@@ -10,3 +10,10 @@ miniscript = { version = "0.12" }
 serde = { version = "^1.0", features = ["derive"] }
 serde_json = { version = "^1.0" }
 base64 = "^0.11"
+
+# Optional dependencies
+sled = { version = "0.31.0", optional = true }
+
+[features]
+default = ["sled"]
+key-value-db = ["sled"]
diff --git a/core/lib/src/database/keyvalue.rs b/core/lib/src/database/keyvalue.rs
new file mode 100644
index 00000000..f04e9fbd
--- /dev/null
+++ b/core/lib/src/database/keyvalue.rs
@@ -0,0 +1,390 @@
+use std::borrow::Borrow;
+use std::convert::{From, TryInto};
+
+use sled::{Batch, Db, IVec, Tree};
+
+use bitcoin::consensus::encode::{deserialize, serialize};
+use bitcoin::hash_types::Txid;
+use bitcoin::util::bip32::{ChildNumber, DerivationPath};
+use bitcoin::{OutPoint, Script, Transaction, TxOut};
+
+use crate::database::{BatchDatabase, BatchOperations, Database};
+use crate::error::Error;
+use crate::types::*;
+
+// TODO: rename mod to Sled?
+
+// path -> script       p{i,e}<path> -> script
+// script -> path       s<script> -> {i,e}<path>
+// outpoint ->          u -> utxo
+// txhash ->            r -> tx
+// txhash ->            t -> tx details
+// change indexes       c{i,e} -> u32
+
+enum SledKey<'a> {
+    Path((Option<ScriptType>, Option<&'a DerivationPath>)),
+    Script(Option<&'a Script>),
+    UTXO(Option<&'a OutPoint>),
+    RawTx(Option<&'a Txid>),
+    Transaction(Option<&'a Txid>),
+    LastIndex(ScriptType),
+}
+
+impl SledKey<'_> {
+    pub fn as_prefix(&self) -> Vec<u8> {
+        match self {
+            SledKey::Path((st, _)) => {
+                let mut v = b"p".to_vec();
+                if let Some(st) = st {
+                    v.push(st.as_byte());
+                }
+                v
+            }
+            SledKey::Script(_) => b"s".to_vec(),
+            SledKey::UTXO(_) => b"u".to_vec(),
+            SledKey::RawTx(_) => b"r".to_vec(),
+            SledKey::Transaction(_) => b"t".to_vec(),
+            SledKey::LastIndex(st) => [b"c", st.as_ref()].concat(),
+        }
+    }
+
+    fn serialize_content(&self) -> Vec<u8> {
+        match self {
+            SledKey::Path((_, Some(path))) => {
+                let mut res = vec![];
+                for val in *path {
+                    res.extend(&u32::from(*val).to_be_bytes());
+                }
+                res
+            }
+            SledKey::Script(Some(s)) => serialize(*s),
+            SledKey::UTXO(Some(s)) => serialize(*s),
+            SledKey::RawTx(Some(s)) => serialize(*s),
+            SledKey::Transaction(Some(s)) => serialize(*s),
+            _ => vec![],
+        }
+    }
+
+    pub fn as_sled_key(&self) -> Vec<u8> {
+        let mut v = self.as_prefix();
+        v.extend_from_slice(&self.serialize_content());
+
+        v
+    }
+}
+
+macro_rules! impl_batch_operations {
+    ( { $($after_insert:tt)* }, $process_delete:ident ) => {
+        fn set_script_pubkey<P: AsRef<[ChildNumber]>>(&mut self, script: Script, script_type: ScriptType, path: P) -> Result<(), Error> {
+            let deriv_path = DerivationPath::from(path.as_ref());
+            let key = SledKey::Path((Some(script_type), Some(&deriv_path))).as_sled_key();
+            self.insert(key, serialize(&script))$($after_insert)*;
+
+            let key = SledKey::Script(Some(&script)).as_sled_key();
+            let value = json!({
+                "t": script_type,
+                "p": deriv_path,
+            });
+            self.insert(key, serde_json::to_vec(&value)?)$($after_insert)*;
+
+            Ok(())
+        }
+
+        fn set_utxo(&mut self, utxo: UTXO) -> Result<(), Error> {
+            let key = SledKey::UTXO(Some(&utxo.outpoint)).as_sled_key();
+            let value = serialize(&utxo.txout);
+            self.insert(key, value)$($after_insert)*;
+
+            Ok(())
+        }
+
+        fn set_raw_tx(&mut self, transaction: Transaction) -> Result<(), Error> {
+            let key = SledKey::RawTx(Some(&transaction.txid())).as_sled_key();
+            let value = serialize(&transaction);
+            self.insert(key, value)$($after_insert)*;
+
+            Ok(())
+        }
+
+        fn set_tx(&mut self, mut transaction: TransactionDetails) -> Result<(), Error> {
+            let key = SledKey::Transaction(Some(&transaction.txid)).as_sled_key();
+            // remove the raw tx
+            transaction.transaction = None;
+            let value = serde_json::to_vec(&transaction)?;
+            self.insert(key, value)$($after_insert)*;
+
+            Ok(())
+        }
+
+        fn set_last_index(&mut self, script_type: ScriptType, value: u32) -> Result<(), Error> {
+            let key = SledKey::LastIndex(script_type).as_sled_key();
+            self.insert(key, &value.to_be_bytes())$($after_insert)*;
+
+            Ok(())
+        }
+
+        fn del_script_pubkey_from_path<P: AsRef<[ChildNumber]>>(&mut self, script_type: ScriptType, path: &P) -> Result<Option<Script>, Error> {
+            let deriv_path = DerivationPath::from(path.as_ref());
+            let key = SledKey::Path((Some(script_type), Some(&deriv_path))).as_sled_key();
+            let res = self.remove(key);
+            let res = $process_delete!(res);
+
+            Ok(res.map_or(Ok(None), |x| Some(deserialize(&x)).transpose())?)
+        }
+
+        fn del_path_from_script_pubkey(&mut self, script: &Script) -> Result<Option<(ScriptType, DerivationPath)>, Error> {
+            let key = SledKey::Script(Some(script)).as_sled_key();
+            let res = self.remove(key);
+            let res = $process_delete!(res);
+
+            match res {
+                None => Ok(None),
+                Some(b) => {
+                    let mut val: serde_json::Value = serde_json::from_slice(&b)?;
+                    let st = serde_json::from_value(val["t"].take())?;
+                    let path = serde_json::from_value(val["p"].take())?;
+
+                    Ok(Some((st, path)))
+                }
+            }
+        }
+
+        fn del_utxo(&mut self, outpoint: &OutPoint) -> Result<Option<UTXO>, Error> {
+            let key = SledKey::UTXO(Some(outpoint)).as_sled_key();
+            let res = self.remove(key);
+            let res = $process_delete!(res);
+
+            match res {
+                None => Ok(None),
+                Some(b) => {
+                    let txout = deserialize(&b)?;
+                    Ok(Some(UTXO { outpoint: outpoint.clone(), txout }))
+                }
+            }
+        }
+
+        fn del_raw_tx(&mut self, txid: &Txid) -> Result<Option<Transaction>, Error> {
+            let key = SledKey::RawTx(Some(txid)).as_sled_key();
+            let res = self.remove(key);
+            let res = $process_delete!(res);
+
+            Ok(res.map_or(Ok(None), |x| Some(deserialize(&x)).transpose())?)
+        }
+
+        fn del_tx(&mut self, txid: &Txid, include_raw: bool) -> Result<Option<TransactionDetails>, Error> {
+            let raw_tx = if include_raw {
+                self.del_raw_tx(txid)?
+            } else {
+                None
+            };
+
+            let key = SledKey::Transaction(Some(txid)).as_sled_key();
+            let res = self.remove(key);
+            let res = $process_delete!(res);
+
+            match res {
+                None => Ok(None),
+                Some(b) => {
+                    let mut val: TransactionDetails = serde_json::from_slice(&b)?;
+                    val.transaction = raw_tx;
+
+                    Ok(Some(val))
+                }
+            }
+        }
+
+        fn del_last_index(&mut self, script_type: ScriptType) -> Result<Option<u32>, Error> {
+            let key = SledKey::LastIndex(script_type).as_sled_key();
+            let res = self.remove(key);
+            let res = $process_delete!(res);
+
+            match res {
+                None => Ok(None),
+                Some(b) => {
+                    let array: [u8; 4] = b.as_ref().try_into().map_err(|_| Error::InvalidU32Bytes(b.to_vec()))?;
+                    let val = u32::from_be_bytes(array);
+                    Ok(Some(val))
+                }
+            }
+        }
+    }
+}
+
+macro_rules! process_delete_tree {
+    ($res:expr) => {
+        $res?
+    };
+}
+impl BatchOperations for Tree {
+    impl_batch_operations!({?}, process_delete_tree);
+}
+
+macro_rules! process_delete_batch {
+    ($res:expr) => {
+        None as Option<sled::IVec>
+    };
+}
+impl BatchOperations for Batch {
+    impl_batch_operations!({}, process_delete_batch);
+}
+
+impl Database for Tree {
+    fn iter_script_pubkeys(&self, script_type: Option<ScriptType>) -> Vec<Result<Script, Error>> {
+        let key = SledKey::Path((script_type, None)).as_sled_key();
+        self.scan_prefix(key)
+            .map(|x| -> Result<_, Error> {
+                let (_, v) = x?;
+                Ok(deserialize(&v)?)
+            })
+            .collect()
+    }
+
+    fn iter_utxos(&self) -> Vec<Result<UTXO, Error>> {
+        let key = SledKey::UTXO(None).as_sled_key();
+        self.scan_prefix(key)
+            .map(|x| -> Result<_, Error> {
+                let (k, v) = x?;
+                let outpoint = deserialize(&k[1..])?;
+                let txout = deserialize(&v)?;
+                Ok(UTXO { outpoint, txout })
+            })
+            .collect()
+    }
+
+    fn iter_raw_txs(&self) -> Vec<Result<Transaction, Error>> {
+        let key = SledKey::RawTx(None).as_sled_key();
+        self.scan_prefix(key)
+            .map(|x| -> Result<_, Error> {
+                let (_, v) = x?;
+                Ok(deserialize(&v)?)
+            })
+            .collect()
+    }
+
+    fn iter_txs(&self, include_raw: bool) -> Vec<Result<TransactionDetails, Error>> {
+        let key = SledKey::RawTx(None).as_sled_key();
+        self.scan_prefix(key)
+            .map(|x| -> Result<_, Error> {
+                let (k, v) = x?;
+                let mut txdetails: TransactionDetails = serde_json::from_slice(&v)?;
+                if include_raw {
+                    let txid = deserialize(&k[1..])?;
+                    txdetails.transaction = self.get_raw_tx(&txid)?;
+                }
+
+                Ok(txdetails)
+            })
+            .collect()
+    }
+
+    fn get_script_pubkey_from_path<P: AsRef<[ChildNumber]>>(
+        &self,
+        script_type: ScriptType,
+        path: P,
+    ) -> Result<Option<Script>, Error> {
+        let deriv_path = DerivationPath::from(path.as_ref());
+        let key = SledKey::Path((Some(script_type), Some(&deriv_path))).as_sled_key();
+        Ok(self.get(key)?.map(|b| deserialize(&b)).transpose()?)
+    }
+
+    fn get_path_from_script_pubkey(
+        &self,
+        script: &Script,
+    ) -> Result<Option<(ScriptType, DerivationPath)>, Error> {
+        let key = SledKey::Script(Some(script)).as_sled_key();
+        self.get(key)?
+            .map(|b| -> Result<_, Error> {
+                let mut val: serde_json::Value = serde_json::from_slice(&b)?;
+                let st = serde_json::from_value(val["t"].take())?;
+                let path = serde_json::from_value(val["p"].take())?;
+
+                Ok((st, path))
+            })
+            .transpose()
+    }
+
+    fn get_utxo(&self, outpoint: &OutPoint) -> Result<Option<UTXO>, Error> {
+        let key = SledKey::UTXO(Some(outpoint)).as_sled_key();
+        self.get(key)?
+            .map(|b| -> Result<_, Error> {
+                let txout = deserialize(&b)?;
+                Ok(UTXO {
+                    outpoint: outpoint.clone(),
+                    txout,
+                })
+            })
+            .transpose()
+    }
+
+    fn get_raw_tx(&self, txid: &Txid) -> Result<Option<Transaction>, Error> {
+        let key = SledKey::RawTx(Some(txid)).as_sled_key();
+        Ok(self.get(key)?.map(|b| deserialize(&b)).transpose()?)
+    }
+
+    fn get_tx(&self, txid: &Txid, include_raw: bool) -> Result<Option<TransactionDetails>, Error> {
+        let key = SledKey::Transaction(Some(txid)).as_sled_key();
+        self.get(key)?
+            .map(|b| -> Result<_, Error> {
+                let mut txdetails: TransactionDetails = serde_json::from_slice(&b)?;
+                if include_raw {
+                    txdetails.transaction = self.get_raw_tx(&txid)?;
+                }
+
+                Ok(txdetails)
+            })
+            .transpose()
+    }
+
+    fn get_last_index(&self, script_type: ScriptType) -> Result<Option<u32>, Error> {
+        let key = SledKey::LastIndex(script_type).as_sled_key();
+        self.get(key)?
+            .map(|b| -> Result<_, Error> {
+                let array: [u8; 4] = b
+                    .as_ref()
+                    .try_into()
+                    .map_err(|_| Error::InvalidU32Bytes(b.to_vec()))?;
+                let val = u32::from_be_bytes(array);
+                Ok(val)
+            })
+            .transpose()
+    }
+
+    // inserts 0 if not present
+    fn increment_last_index(&mut self, script_type: ScriptType) -> Result<Option<u32>, Error> {
+        let key = SledKey::LastIndex(script_type).as_sled_key();
+        self.fetch_and_update(key, |prev| {
+            let new = match prev {
+                Some(b) => {
+                    let array: [u8; 4] = b.try_into().unwrap_or([0; 4]);
+                    let val = u32::from_be_bytes(array);
+
+                    val + 1
+                }
+                None => 0,
+            };
+
+            Some(new.to_be_bytes().to_vec())
+        })?
+        .map(|b| -> Result<_, Error> {
+            let array: [u8; 4] = b
+                .as_ref()
+                .try_into()
+                .map_err(|_| Error::InvalidU32Bytes(b.to_vec()))?;
+            let val = u32::from_be_bytes(array);
+            Ok(val)
+        })
+        .transpose()
+    }
+}
+
+impl BatchDatabase for Tree {
+    type Batch = sled::Batch;
+
+    fn start_batch(&self) -> Self::Batch {
+        sled::Batch::default()
+    }
+
+    fn apply_batch(&mut self, batch: Self::Batch) -> Result<(), Error> {
+        Ok(self.apply_batch(batch)?)
+    }
+}
diff --git a/core/lib/src/database/mod.rs b/core/lib/src/database/mod.rs
new file mode 100644
index 00000000..872ddf70
--- /dev/null
+++ b/core/lib/src/database/mod.rs
@@ -0,0 +1,71 @@
+use bitcoin::hash_types::Txid;
+use bitcoin::util::bip32::{ChildNumber, DerivationPath};
+use bitcoin::{OutPoint, Script, Transaction, TxOut};
+
+use crate::error::Error;
+use crate::types::*;
+
+#[cfg(any(feature = "key-value-db", feature = "default"))]
+pub mod keyvalue;
+
+pub trait BatchOperations {
+    fn set_script_pubkey<P: AsRef<[ChildNumber]>>(
+        &mut self,
+        script: Script,
+        script_type: ScriptType,
+        path: P,
+    ) -> Result<(), Error>;
+    fn set_utxo(&mut self, utxo: UTXO) -> Result<(), Error>;
+    fn set_raw_tx(&mut self, transaction: Transaction) -> Result<(), Error>;
+    fn set_tx(&mut self, transaction: TransactionDetails) -> Result<(), Error>;
+    fn set_last_index(&mut self, script_type: ScriptType, value: u32) -> Result<(), Error>;
+
+    fn del_script_pubkey_from_path<P: AsRef<[ChildNumber]>>(
+        &mut self,
+        script_type: ScriptType,
+        path: P,
+    ) -> Result<Option<Script>, Error>;
+    fn del_path_from_script_pubkey(
+        &mut self,
+        script: &Script,
+    ) -> Result<Option<(ScriptType, DerivationPath)>, Error>;
+    fn del_utxo(&mut self, outpoint: &OutPoint) -> Result<Option<UTXO>, Error>;
+    fn del_raw_tx(&mut self, txid: &Txid) -> Result<Option<Transaction>, Error>;
+    fn del_tx(
+        &mut self,
+        txid: &Txid,
+        include_raw: bool,
+    ) -> Result<Option<TransactionDetails>, Error>;
+    fn del_last_index(&mut self, script_type: ScriptType) -> Result<Option<u32>, Error>;
+}
+
+pub trait Database: BatchOperations {
+    fn iter_script_pubkeys(&self, script_type: Option<ScriptType>) -> Vec<Result<Script, Error>>;
+    fn iter_utxos(&self) -> Vec<Result<UTXO, Error>>;
+    fn iter_raw_txs(&self) -> Vec<Result<Transaction, Error>>;
+    fn iter_txs(&self, include_raw: bool) -> Vec<Result<TransactionDetails, Error>>;
+
+    fn get_script_pubkey_from_path<P: AsRef<[ChildNumber]>>(
+        &self,
+        script_type: ScriptType,
+        path: P,
+    ) -> Result<Option<Script>, Error>;
+    fn get_path_from_script_pubkey(
+        &self,
+        script: &Script,
+    ) -> Result<Option<(ScriptType, DerivationPath)>, Error>;
+    fn get_utxo(&self, outpoint: &OutPoint) -> Result<Option<UTXO>, Error>;
+    fn get_raw_tx(&self, txid: &Txid) -> Result<Option<Transaction>, Error>;
+    fn get_tx(&self, txid: &Txid, include_raw: bool) -> Result<Option<TransactionDetails>, Error>;
+    fn get_last_index(&self, script_type: ScriptType) -> Result<Option<u32>, Error>;
+
+    // inserts 0 if not present
+    fn increment_last_index(&mut self, script_type: ScriptType) -> Result<Option<u32>, Error>;
+}
+
+pub trait BatchDatabase: Database {
+    type Batch: BatchOperations;
+
+    fn start_batch(&self) -> Self::Batch;
+    fn apply_batch(&mut self, batch: Self::Batch) -> Result<(), Error>;
+}
diff --git a/core/lib/src/error.rs b/core/lib/src/error.rs
index 78c6c571..80b72c15 100644
--- a/core/lib/src/error.rs
+++ b/core/lib/src/error.rs
@@ -2,9 +2,16 @@
 pub enum Error {
     KeyMismatch(bitcoin::secp256k1::PublicKey, bitcoin::secp256k1::PublicKey),
     MissingInputUTXO(usize),
+    InvalidU32Bytes(Vec<u8>),
+    Generic(String),
 
+    Encode(bitcoin::consensus::encode::Error),
     BIP32(bitcoin::util::bip32::Error),
     Secp256k1(bitcoin::secp256k1::Error),
+    JSON(serde_json::Error),
+
+    #[cfg(any(feature = "key-value-db", feature = "default"))]
+    Sled(sled::Error),
 }
 
 macro_rules! impl_error {
@@ -17,5 +24,10 @@ macro_rules! impl_error {
     };
 }
 
+impl_error!(bitcoin::consensus::encode::Error, Encode);
 impl_error!(bitcoin::util::bip32::Error, BIP32);
 impl_error!(bitcoin::secp256k1::Error, Secp256k1);
+impl_error!(serde_json::Error, JSON);
+
+#[cfg(any(feature = "key-value-db", feature = "default"))]
+impl_error!(sled::Error, Sled);
diff --git a/core/lib/src/lib.rs b/core/lib/src/lib.rs
index 242393c9..b2bc2221 100644
--- a/core/lib/src/lib.rs
+++ b/core/lib/src/lib.rs
@@ -2,10 +2,16 @@ pub extern crate bitcoin;
 extern crate log;
 pub extern crate miniscript;
 extern crate serde;
+#[macro_use]
 extern crate serde_json;
 
+#[cfg(any(feature = "key-value-db", feature = "default"))]
+extern crate sled;
+
 #[macro_use]
 pub mod error;
+pub mod database;
 pub mod descriptor;
 pub mod psbt;
 pub mod signer;
+pub mod types;
diff --git a/core/lib/src/types.rs b/core/lib/src/types.rs
new file mode 100644
index 00000000..1ed23aa5
--- /dev/null
+++ b/core/lib/src/types.rs
@@ -0,0 +1,47 @@
+use std::convert::AsRef;
+
+use bitcoin::blockdata::transaction::{OutPoint, Transaction, TxOut};
+use bitcoin::hash_types::Txid;
+
+use serde::{Deserialize, Serialize};
+
+// TODO serde flatten?
+#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
+pub enum ScriptType {
+    External = 0,
+    Internal = 1,
+}
+
+impl ScriptType {
+    pub fn as_byte(&self) -> u8 {
+        match self {
+            ScriptType::External => 'e' as u8,
+            ScriptType::Internal => 'i' as u8,
+        }
+    }
+}
+
+impl AsRef<[u8]> for ScriptType {
+    fn as_ref(&self) -> &[u8] {
+        match self {
+            ScriptType::External => b"e",
+            ScriptType::Internal => b"i",
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct UTXO {
+    pub outpoint: OutPoint,
+    pub txout: TxOut,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct TransactionDetails {
+    pub transaction: Option<Transaction>,
+    pub txid: Txid,
+    pub timestamp: u64,
+    pub received: u64,
+    pub sent: u64,
+    pub height: Option<u32>,
+}