From e539efabf160b396cbfd500f38a25be82180fd62 Mon Sep 17 00:00:00 2001
From: lunar-mining <ihateface@autistici.org>
Date: Wed, 22 Dec 2021 14:08:56 +0100
Subject: [PATCH] cashier: add save/load/create merkle tree to cashierdb and
 cashierd

---
 src/bin/cashierd.rs     |  3 ++-
 src/wallet/cashierdb.rs | 45 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/bin/cashierd.rs b/src/bin/cashierd.rs
index a04e86a39..993b1a5dc 100644
--- a/src/bin/cashierd.rs
+++ b/src/bin/cashierd.rs
@@ -692,6 +692,7 @@ async fn start(
     )
     .await?;
 
+    let tree = client.get_tree().await?;
     let merkle_roots = RocksColumn::<columns::MerkleRoots>::new(rocks.clone());
     let nullifiers = RocksColumn::<columns::Nullifiers>::new(rocks);
 
@@ -702,7 +703,7 @@ async fn start(
     let cashier_public_keys = vec![cashier_public];
 
     let state = Arc::new(Mutex::new(State {
-        tree: BridgeTree::<MerkleNode, 32>::new(100),
+        tree,
         merkle_roots,
         nullifiers,
         public_keys: cashier_public_keys,
diff --git a/src/wallet/cashierdb.rs b/src/wallet/cashierdb.rs
index 68f6607a3..575940d70 100644
--- a/src/wallet/cashierdb.rs
+++ b/src/wallet/cashierdb.rs
@@ -1,3 +1,4 @@
+use incrementalmerkletree::bridgetree::BridgeTree;
 use std::{fs::create_dir_all, path::Path, str::FromStr};
 
 use async_std::sync::Arc;
@@ -10,7 +11,10 @@ use sqlx::{
 use super::wallet_api::WalletApi;
 use crate::{
     client::ClientFailed,
-    crypto::keypair::{Keypair, PublicKey, SecretKey},
+    crypto::{
+        keypair::{Keypair, PublicKey, SecretKey},
+        merkle_node::MerkleNode,
+    },
     types::DrkTokenId,
     util::NetworkName,
     Error, Result,
@@ -89,6 +93,45 @@ impl CashierDb {
         Ok(())
     }
 
+    pub async fn tree_gen(&self) -> Result<()> {
+        trace!("Attempting to generate merkle tree");
+        let mut conn = self.conn.acquire().await?;
+
+        match sqlx::query("SELECT * FROM tree").fetch_one(&mut conn).await {
+            Ok(_) => {
+                error!("Tree already exist");
+                Err(Error::from(ClientFailed::TreeExists))
+            }
+            Err(_) => {
+                let tree = BridgeTree::<MerkleNode, 32>::new(100);
+                self.put_tree(&tree).await?;
+                Ok(())
+            }
+        }
+    }
+
+    pub async fn get_tree(&self) -> Result<BridgeTree<MerkleNode, 32>> {
+        trace!("Getting merkle tree");
+        let mut conn = self.conn.acquire().await?;
+
+        let row = sqlx::query("SELECT tree FROM tree").fetch_one(&mut conn).await?;
+        let tree: BridgeTree<MerkleNode, 32> = bincode::deserialize(row.get("tree"))?;
+        Ok(tree)
+    }
+
+    pub async fn put_tree(&self, tree: &BridgeTree<MerkleNode, 32>) -> Result<()> {
+        trace!("Attempting to write merkle tree");
+        let mut conn = self.conn.acquire().await?;
+
+        let tree_bytes = bincode::serialize(tree)?;
+        sqlx::query("INSERT INTO tree(tree) VALUES (?1)")
+            .bind(tree_bytes)
+            .execute(&mut conn)
+            .await?;
+
+        Ok(())
+    }
+
     pub async fn put_main_keys(&self, token_key: &TokenKey, network: &NetworkName) -> Result<()> {
         trace!("Writing main keys into the database");
         let network = self.get_value_serialized(network)?;