create WalletApi trait and implement it for both cashierdb and walletdb

src/bin/darkfid.rs

@@ -1,15 +1,16 @@
 use drk::blockchain::Rocks;
 use drk::cli::{Config, DarkfidCli, DarkfidConfig};
 use drk::util::join_config_path;
-use drk::wallet::WalletDb;
+use drk::wallet::{WalletApi, WalletDb};
+use drk::serial::serialize;
 use drk::Result;
 
 use drk::client::Client;
 
 use async_executor::Executor;
 use easy_parallel::Parallel;
-use ff::Field;
 use rand::rngs::OsRng;
+use ff::Field;
 
 use async_std::sync::Arc;
 use std::net::SocketAddr;
@@ -31,23 +32,25 @@ async fn start(executor: Arc<Executor<'_>>, config: Arc<DarkfidConfig>) -> Resul
 
     let wallet = WalletDb::new(&walletdb_path, config.password.clone())?;
 
+    let mint_params_path = join_config_path(&PathBuf::from("mint.params"))?;
+    let spend_params_path = join_config_path(&PathBuf::from("spend.params"))?;
+
     // wallet secret key
     let secret: jubjub::Fr;
-    if let Some(prv) = wallet.get_private().ok() {
-        secret = prv;
+    if let Ok(prv) = wallet.get_private_keys() {
+        secret = prv[0];
     } else {
         secret = jubjub::Fr::random(&mut OsRng);
+        let public = zcash_primitives::constants::SPENDING_KEY_GENERATOR * secret;
+        wallet.put_keypair(serialize(&public), serialize(&secret))?;
     }
 
-    let mint_params_path = join_config_path(&PathBuf::from("mint.params"))?;
-    let spend_params_path = join_config_path(&PathBuf::from("spend.params"))?;
-
     let mut client = Client::new(
-        secret,
         rocks,
         (connect_addr, sub_addr),
         (mint_params_path, spend_params_path),
         wallet.clone(),
+        secret.clone()
     )?;
 
     client.start().await?;

src/client/client.rs

@@ -15,9 +15,11 @@ use crate::service::{CashierClient, GatewayClient, GatewaySlabsSubscriber};
 use crate::state::{state_transition, ProgramState, StateUpdate};
 use crate::wallet::WalletPtr;
 use crate::{tx, Result};
+use crate::wallet::WalletApi;
 
 use super::ClientFailed;
 
+
 use async_executor::Executor;
 use bellman::groth16;
 use bls12_381::Bls12;
@@ -39,11 +41,11 @@ pub struct Client {
 
 impl Client {
     pub fn new(
-        secret: jubjub::Fr,
         rocks: Arc<Rocks>,
         gateway_addrs: (SocketAddr, SocketAddr),
         params_paths: (PathBuf, PathBuf),
         wallet: WalletPtr,
+        secret: jubjub::Fr,
     ) -> Result<Self> {
         let slabstore = RocksColumn::<columns::Slabs>::new(rocks.clone());
         let merkle_roots = RocksColumn::<columns::MerkleRoots>::new(rocks.clone());
@@ -198,7 +200,8 @@ impl Client {
             let tx = tx::Transaction::decode(&slab.get_payload()[..])?;
             let mut client = client.lock().await;
             let update = state_transition(&client.state, tx)?;
-            client.state.apply(update).await?;
+            let secret_keys = client.state.wallet.get_private_keys()?;
+            client.state.apply(update, secret_keys.clone()).await?;
         }
     }
 }
@@ -249,14 +252,14 @@ impl ProgramState for State {
 }
 
 impl State {
-    pub async fn apply(&mut self, update: StateUpdate) -> Result<()> {
+    pub async fn apply(&mut self, update: StateUpdate, secret_keys: Vec<jubjub::Fr>) -> Result<()> {
         // Extend our list of nullifiers with the ones from the update
         for nullifier in update.nullifiers {
             self.nullifiers.put(nullifier, vec![] as Vec<u8>)?;
         }
 
         // Update merkle tree and witnesses
-        for (coin, enc_note) in update.coins.into_iter().zip(update.enc_notes.into_iter()) {
+        for (coin, enc_note) in update.coins.into_iter().zip(update.enc_notes.iter()) {
             // Add the new coins to the merkle tree
             let node = MerkleNode::from_coin(&coin);
             self.tree.append(node).expect("Append to merkle tree");
@@ -271,32 +274,37 @@ impl State {
                     .update_witness(coin_id.clone(), witness.clone())?;
             }
 
-            if let Some((note, secret)) = self.try_decrypt_note(enc_note).await {
-                // We need to keep track of the witness for this coin.
-                // This allows us to prove inclusion of the coin in the merkle tree with ZK.
-                // Just as we update the merkle tree with every new coin, so we do the same with
-                // the witness.
-
-                // Derive the current witness from the current tree.
-                // This is done right after we add our coin to the tree (but before any other
-                // coins are added)
-
-                // Make a new witness for this coin
-                let witness = IncrementalWitness::from_tree(&self.tree);
-
-                self.wallet
-                    .put_own_coins(coin.clone(), note.clone(), secret, witness.clone())?;
+            for secret in secret_keys.iter() {
+                if let Some(note) = Self::try_decrypt_note(enc_note.clone(), secret.clone()) {
+                    // We need to keep track of the witness for this coin.
+                    // This allows us to prove inclusion of the coin in the merkle tree with ZK.
+                    // Just as we update the merkle tree with every new coin, so we do the same with
+                    // the witness.
+
+                    // Derive the current witness from the current tree.
+                    // This is done right after we add our coin to the tree (but before any other
+                    // coins are added)
+
+                    // Make a new witness for this coin
+                    let witness = IncrementalWitness::from_tree(&self.tree);
+
+                    self.wallet.put_own_coins(
+                        coin.clone(),
+                        note.clone(),
+                        secret.clone(),
+                        witness.clone(),
+                    )?;
+                }
             }
         }
         Ok(())
     }
 
-    async fn try_decrypt_note(&self, ciphertext: EncryptedNote) -> Option<(Note, jubjub::Fr)> {
-        let secret = self.wallet.get_private().ok()?;
+    fn try_decrypt_note(ciphertext: &EncryptedNote, secret: jubjub::Fr) -> Option<Note> {
         match ciphertext.decrypt(&secret) {
             Ok(note) => {
                 // ... and return the decrypted note for this coin.
-                return Some((note, secret.clone()));
+                return Some(note);
             }
             Err(_) => {}
         }

src/rpc/adapters/client_adapter.rs

@@ -2,6 +2,7 @@ use crate::client::{Client, ClientFailed};
 use crate::serial::serialize;
 use crate::service::CashierClient;
 use crate::{Error, Result};
+use crate::wallet::WalletApi;
 
 use jsonrpc_core::BoxFuture;
 use jsonrpc_derive::rpc;
@@ -56,7 +57,7 @@ impl RpcClientAdapter {
     }
 
     async fn get_key_process(client: Arc<Mutex<Client>>) -> Result<String> {
-        let key_public = client.lock().await.state.wallet.get_public()?;
+        let key_public = client.lock().await.state.wallet.get_public_keys()?[0];
         let bs58_address = bs58::encode(serialize(&key_public)).into_string();
         Ok(bs58_address)
     }
@@ -122,7 +123,7 @@ impl RpcClientAdapter {
         client: Arc<Mutex<Client>>,
         cashier_client: Arc<Mutex<CashierClient>>,
     ) -> Result<String> {
-        let deposit_addr = client.lock().await.state.wallet.get_public()?;
+        let deposit_addr = client.lock().await.state.wallet.get_public_keys()?[0];
         let btc_public = cashier_client
             .lock()
             .await

src/service/cashier.rs

@@ -5,6 +5,7 @@ use crate::client::Client;
 use crate::serial::{deserialize, serialize};
 use crate::wallet::{CashierDbPtr, WalletPtr};
 use crate::{Error, Result};
+use crate::wallet::WalletApi; 
 
 use ff::Field;
 use rand::rngs::OsRng;
@@ -51,25 +52,24 @@ impl CashierService {
         // create btc client
         let btc_client = Arc::new(ElectrumClient::new(&client_address)?);
 
-        let cashier_secret: jubjub::Fr;
+        let rocks = Rocks::new(&cashier_database_path)?;
 
-        if let Ok(secret) = wallet.get_private() {
-            cashier_secret = secret;
+        // wallet secret key
+        let secret: jubjub::Fr;
+        if let Ok(prv) = wallet.get_private_keys() {
+            secret = prv[0];
         } else {
-            wallet.init_db()?;
-            let keys = wallet.key_gen();
-            wallet.put_keypair(keys.0, keys.1)?;
-            cashier_secret = wallet.get_private()?;
+            secret = jubjub::Fr::random(&mut OsRng);
+            let public = zcash_primitives::constants::SPENDING_KEY_GENERATOR * secret;
+            wallet.put_keypair(serialize(&public), serialize(&secret))?;
         }
 
-        let rocks = Rocks::new(&cashier_database_path)?;
-
         let client = Client::new(
-            cashier_secret,
             rocks,
             gateway_addrs,
             params_paths,
             client_wallet.clone(),
+            secret.clone(),
         )?;
 
         let client = Arc::new(Mutex::new(client));

src/wallet/cashierdb.rs

@@ -1,3 +1,4 @@
+use super::WalletApi;
 use crate::client::ClientFailed;
 use crate::serial;
 use crate::serial::{deserialize, serialize, Decodable, Encodable};
@@ -19,16 +20,8 @@ pub struct CashierDb {
     pub password: String,
 }
 
-impl CashierDb {
-    pub fn new(path: &PathBuf, password: String) -> Result<CashierDbPtr> {
-        debug!(target: "CASHIERDB", "new() Constructor called");
-        Ok(Arc::new(Self {
-            path: path.to_owned(),
-            password,
-        }))
-    }
-
-    pub fn init_db(&self) -> Result<()> {
+impl WalletApi for CashierDb {
+    fn init_db(&self) -> Result<()> {
         if !self.password.trim().is_empty() {
             let contents = include_str!("../../res/cashier.sql");
             let conn = Connection::open(&self.path)?;
@@ -42,6 +35,64 @@ impl CashierDb {
         Ok(())
     }
 
+    fn key_gen(&self) -> Result<(Vec<u8>, Vec<u8>)> {
+        debug!(target: "CASHIERDB", "Generating cashier keys...");
+        let secret: jubjub::Fr = jubjub::Fr::random(&mut OsRng);
+        let public = zcash_primitives::constants::SPENDING_KEY_GENERATOR * secret;
+        let pubkey = serial::serialize(&public);
+        let privkey = serial::serialize(&secret);
+        Ok((pubkey, privkey))
+    }
+
+    fn put_keypair(&self, key_public: Vec<u8>, key_private: Vec<u8>) -> Result<()> {
+        let conn = Connection::open(&self.path)?;
+        conn.pragma_update(None, "key", &self.password)?;
+        conn.execute(
+            "INSERT INTO keys(key_public, key_private) VALUES (?1, ?2)",
+            params![key_public, key_private],
+        )?;
+        Ok(())
+    }
+
+    fn get_public_keys(&self) -> Result<Vec<jubjub::SubgroupPoint>> {
+        debug!(target: "CASHIERDB", "Returning keys...");
+        let conn = Connection::open(&self.path)?;
+        conn.pragma_update(None, "key", &self.password)?;
+        let mut stmt = conn.prepare("SELECT key_public FROM keys")?;
+        let key_iter = stmt.query_map::<Vec<u8>, _, _>([], |row| row.get(0))?;
+        let mut pub_keys = Vec::new();
+        for key in key_iter {
+            let public: jubjub::SubgroupPoint =
+                self.get_value_deserialized::<jubjub::SubgroupPoint>(key?)?;
+            pub_keys.push(public);
+        }
+        Ok(pub_keys)
+    }
+
+    fn get_private_keys(&self) -> Result<Vec<jubjub::Fr>> {
+        debug!(target: "CASHIERDB", "Returning keys...");
+        let conn = Connection::open(&self.path)?;
+        conn.pragma_update(None, "key", &self.password)?;
+        let mut stmt = conn.prepare("SELECT key_private FROM keys")?;
+        let key_iter = stmt.query_map::<Vec<u8>, _, _>([], |row| row.get(0))?;
+        let mut keys = Vec::new();
+        for key in key_iter {
+            let private: jubjub::Fr = self.get_value_deserialized(key?)?;
+            keys.push(private);
+        }
+        Ok(keys)
+    }
+}
+
+impl CashierDb {
+    pub fn new(path: &PathBuf, password: String) -> Result<CashierDbPtr> {
+        debug!(target: "CASHIERDB", "new() Constructor called");
+        Ok(Arc::new(Self {
+            path: path.to_owned(),
+            password,
+        }))
+    }
+
     pub fn get_keys_by_dkey(&self, dkey_pub: &Vec<u8>) -> Result<()> {
         debug!(target: "CASHIERDB", "Check for existing dkey");
         //let dkey_id = self.get_value_deserialized(dkey_pub)?;
@@ -151,55 +202,6 @@ impl CashierDb {
         Ok(())
     }
 
-    pub fn key_gen(&self) -> (Vec<u8>, Vec<u8>) {
-        debug!(target: "CASHIERDB", "Generating cashier keys...");
-        let secret: jubjub::Fr = jubjub::Fr::random(&mut OsRng);
-        let public = zcash_primitives::constants::SPENDING_KEY_GENERATOR * secret;
-        let pubkey = serial::serialize(&public);
-        let privkey = serial::serialize(&secret);
-        (pubkey, privkey)
-    }
-
-    pub fn put_keypair(&self, key_public: Vec<u8>, key_private: Vec<u8>) -> Result<()> {
-        let conn = Connection::open(&self.path)?;
-        conn.pragma_update(None, "key", &self.password)?;
-        conn.execute(
-            "INSERT INTO keys(key_public, key_private) VALUES (?1, ?2)",
-            params![key_public, key_private],
-        )?;
-        Ok(())
-    }
-
-    pub fn get_public(&self) -> Result<jubjub::SubgroupPoint> {
-        debug!(target: "CASHIERDB", "Returning keys...");
-        let conn = Connection::open(&self.path)?;
-        conn.pragma_update(None, "key", &self.password)?;
-        let mut stmt = conn.prepare("SELECT key_public FROM keys")?;
-        let key_iter = stmt.query_map::<Vec<u8>, _, _>([], |row| row.get(0))?;
-        let mut pub_keys = Vec::new();
-        for key in key_iter {
-            pub_keys.push(key?);
-        }
-        let public: jubjub::SubgroupPoint =
-            self.get_value_deserialized(pub_keys.pop().expect("load public_key from cashierdb"))?;
-        Ok(public)
-    }
-
-    pub fn get_private(&self) -> Result<jubjub::Fr> {
-        debug!(target: "CASHIERDB", "Returning keys...");
-        let conn = Connection::open(&self.path)?;
-        conn.pragma_update(None, "key", &self.password)?;
-        let mut stmt = conn.prepare("SELECT key_private FROM keys")?;
-        let key_iter = stmt.query_map::<Vec<u8>, _, _>([], |row| row.get(0))?;
-        let mut keys = Vec::new();
-        for key in key_iter {
-            keys.push(key?);
-        }
-        let private: jubjub::Fr =
-            self.get_value_deserialized(keys.pop().expect("load private_key from cashierdb"))?;
-        Ok(private)
-    }
-
     pub fn test_wallet(&self) -> Result<()> {
         let conn = Connection::open(&self.path)?;
         conn.pragma_update(None, "key", &self.password)?;
@@ -266,8 +268,8 @@ mod tests {
 
         wallet.put_keypair(key_public, key_private)?;
 
-        let public2 = wallet.get_public()?;
-        let secret2 = wallet.get_private()?;
+        let public2 = wallet.get_public_keys()?[0];
+        let secret2 = wallet.get_private_keys()?[0];
 
         assert_eq!(public, public2);
         assert_eq!(secret, secret2);

src/wallet/mod.rs

@@ -1,5 +1,7 @@
 pub mod cashierdb;
 pub mod walletdb;
+pub mod wallet_api;
 
-pub use cashierdb::{CashierDb, CashierDbPtr};
+pub use wallet_api::WalletApi;
 pub use walletdb::{WalletDb, WalletPtr};
+pub use cashierdb::{CashierDb, CashierDbPtr};

src/wallet/wallet_api.rs

@@ -0,0 +1,9 @@
+use crate::Result;
+
+pub trait WalletApi {
+    fn init_db(&self) -> Result<()>; 
+    fn get_public_keys(&self) -> Result<Vec<jubjub::SubgroupPoint>>;
+    fn key_gen(&self) -> Result<(Vec<u8>, Vec<u8>)>;
+    fn put_keypair(&self, key_public: Vec<u8>, key_private: Vec<u8>) -> Result<()>;
+    fn get_private_keys(&self) -> Result<Vec<jubjub::Fr>>;
+}

src/wallet/walletdb.rs

@@ -1,3 +1,4 @@
+use super::WalletApi;
 use crate::client::ClientFailed;
 use crate::crypto::{
     coin::Coin, merkle::IncrementalWitness, merkle_node::MerkleNode, note::Note, OwnCoins,
@@ -21,16 +22,8 @@ pub struct WalletDb {
     pub password: String,
 }
 
-impl WalletDb {
-    pub fn new(path: &PathBuf, password: String) -> Result<WalletPtr> {
-        debug!(target: "WALLETDB", "new() Constructor called");
-        Ok(Arc::new(Self {
-            path: path.to_owned(),
-            password,
-        }))
-    }
-
-    pub fn init_db(&self) -> Result<()> {
+impl WalletApi for WalletDb {
+    fn init_db(&self) -> Result<()> {
         if !self.password.trim().is_empty() {
             let contents = include_str!("../../res/schema.sql");
             let conn = Connection::open(&self.path)?;
@@ -44,6 +37,65 @@ impl WalletDb {
         Ok(())
     }
 
+    fn key_gen(&self) -> Result<(Vec<u8>, Vec<u8>)> {
+        debug!(target: "WALLETDB", "Attempting to generate keys...");
+        let secret: jubjub::Fr = jubjub::Fr::random(&mut OsRng);
+        let public = zcash_primitives::constants::SPENDING_KEY_GENERATOR * secret;
+        let pubkey = serial::serialize(&public);
+        let privkey = serial::serialize(&secret);
+        self.put_keypair(pubkey.clone(), privkey.clone())?;
+        Ok((pubkey, privkey))
+    }
+
+    fn put_keypair(&self, key_public: Vec<u8>, key_private: Vec<u8>) -> Result<()> {
+        let conn = Connection::open(&self.path)?;
+        conn.pragma_update(None, "key", &self.password)?;
+        conn.execute(
+            "INSERT INTO keys(key_public, key_private) VALUES (?1, ?2)",
+            params![key_public, key_private],
+        )?;
+        Ok(())
+    }
+    fn get_public_keys(&self) -> Result<Vec<jubjub::SubgroupPoint>> {
+        debug!(target: "WALLETDB", "Returning keys...");
+        let conn = Connection::open(&self.path)?;
+        conn.pragma_update(None, "key", &self.password)?;
+        let mut stmt = conn.prepare("SELECT key_public FROM keys")?;
+        // this just gets the first key. maybe we should randomize this
+        let key_iter = stmt.query_map([], |row| row.get(0))?;
+        let mut pub_keys = Vec::new();
+        for key in key_iter {
+            let public: jubjub::SubgroupPoint =
+                self.get_value_deserialized::<jubjub::SubgroupPoint>(key?)?;
+            pub_keys.push(public);
+        }
+        Ok(pub_keys)
+    }
+
+    fn get_private_keys(&self) -> Result<Vec<jubjub::Fr>> {
+        debug!(target: "WALLETDB", "Returning keys...");
+        let conn = Connection::open(&self.path)?;
+        conn.pragma_update(None, "key", &self.password)?;
+        let mut stmt = conn.prepare("SELECT key_private FROM keys")?;
+        let key_iter = stmt.query_map([], |row| row.get(0))?;
+        let mut keys = Vec::new();
+        for key in key_iter {
+            let private: jubjub::Fr = self.get_value_deserialized(key?)?;
+            keys.push(private);
+        }
+        Ok(keys)
+    }
+}
+
+impl WalletDb {
+    pub fn new(path: &PathBuf, password: String) -> Result<WalletPtr> {
+        debug!(target: "WALLETDB", "new() Constructor called");
+        Ok(Arc::new(Self {
+            path: path.to_owned(),
+            password,
+        }))
+    }
+
     pub fn get_own_coins(&self) -> Result<OwnCoins> {
         // open connection
         let conn = Connection::open(&self.path)?;
@@ -186,26 +238,6 @@ impl WalletDb {
         Ok(())
     }
 
-    pub fn key_gen(&self) -> Result<(Vec<u8>, Vec<u8>)> {
-        debug!(target: "WALLETDB", "Attempting to generate keys...");
-        let secret: jubjub::Fr = jubjub::Fr::random(&mut OsRng);
-        let public = zcash_primitives::constants::SPENDING_KEY_GENERATOR * secret;
-        let pubkey = serial::serialize(&public);
-        let privkey = serial::serialize(&secret);
-        self.put_keypair(pubkey.clone(), privkey.clone())?;
-        Ok((pubkey, privkey))
-    }
-
-    pub fn put_keypair(&self, key_public: Vec<u8>, key_private: Vec<u8>) -> Result<()> {
-        let conn = Connection::open(&self.path)?;
-        conn.pragma_update(None, "key", &self.password)?;
-        conn.execute(
-            "INSERT INTO keys(key_public, key_private) VALUES (?1, ?2)",
-            params![key_public, key_private],
-        )?;
-        Ok(())
-    }
-
     pub fn put_cashier_pub(&self, key_public: Vec<u8>) -> Result<()> {
         debug!(target: "WALLETDB", "Save cashier keys...");
         let conn = Connection::open(&self.path)?;
@@ -217,23 +249,6 @@ impl WalletDb {
         Ok(())
     }
 
-    pub fn get_public(&self) -> Result<jubjub::SubgroupPoint> {
-        debug!(target: "WALLETDB", "Returning keys...");
-        let conn = Connection::open(&self.path)?;
-        conn.pragma_update(None, "key", &self.password)?;
-        let mut stmt = conn.prepare("SELECT key_public FROM keys")?;
-        // this just gets the first key. maybe we should randomize this
-        let key_iter = stmt.query_map([], |row| row.get(0))?;
-        let mut pub_keys = Vec::new();
-        for key in key_iter {
-            pub_keys.push(key?);
-        }
-        let public: jubjub::SubgroupPoint =
-            self.get_value_deserialized(pub_keys.pop().expect("Load public_key from walletdb"))?;
-
-        Ok(public)
-    }
-
     pub fn get_cashier_public_keys(&self) -> Result<Vec<jubjub::SubgroupPoint>> {
         debug!(target: "WALLETDB", "Returning keys...");
         let conn = Connection::open(&self.path)?;
@@ -249,21 +264,6 @@ impl WalletDb {
         Ok(pub_keys)
     }
 
-    pub fn get_private(&self) -> Result<jubjub::Fr> {
-        debug!(target: "WALLETDB", "Returning keys...");
-        let conn = Connection::open(&self.path)?;
-        conn.pragma_update(None, "key", &self.password)?;
-        let mut stmt = conn.prepare("SELECT key_private FROM keys")?;
-        let key_iter = stmt.query_map([], |row| row.get(0))?;
-        let mut keys = Vec::new();
-        for key in key_iter {
-            keys.push(key?);
-        }
-        let private: jubjub::Fr =
-            self.get_value_deserialized(keys.pop().expect("Load private key from walletdb"))?;
-        Ok(private)
-    }
-
     pub fn test_wallet(&self) -> Result<()> {
         let conn = Connection::open(&self.path)?;
         conn.pragma_update(None, "key", &self.password)?;
@@ -330,11 +330,11 @@ mod tests {
 
         wallet.put_keypair(key_public, key_private)?;
 
-        let public2 = wallet.get_public()?;
-        let secret2 = wallet.get_private()?;
+        let public2 = wallet.get_public_keys()?;
+        let secret2 = wallet.get_private_keys()?;
 
-        assert_eq!(public, public2);
-        assert_eq!(secret, secret2);
+        assert_eq!(public, public2[0]);
+        assert_eq!(secret, secret2[0]);
 
         wallet.destroy()?;