Diversifier addresses

hhanh00 · Jul 8, 2021 · 7d40ad2 · 7d40ad2
1 parent 4548e88
commit 7d40ad2
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 3 deletions.
diff --git a/src/db.rs b/src/db.rs
@@ -5,7 +5,7 @@ use std::collections::HashMap;
 use zcash_primitives::consensus::{NetworkUpgrade, Parameters};
 use zcash_primitives::merkle_tree::IncrementalWitness;
 use zcash_primitives::sapling::{Diversifier, Node, Note, Rseed};
-use zcash_primitives::zip32::ExtendedFullViewingKey;
+use zcash_primitives::zip32::{ExtendedFullViewingKey, DiversifierIndex};
 
 #[allow(dead_code)]
 pub const DEFAULT_DB_PATH: &str = "zec.db";
@@ -98,6 +98,13 @@ impl DbAdapter {
             NO_PARAMS,
         )?;
 
+        self.connection.execute(
+            "CREATE TABLE IF NOT EXISTS diversifiers (
+            account INTEGER PRIMARY KEY NOT NULL,
+            diversifier_index BLOB NOT NULL)",
+            NO_PARAMS,
+        )?;
+
         Ok(())
     }
 
@@ -338,7 +345,7 @@ impl DbAdapter {
             nf.clone_from_slice(&nf_vec);
             let nf_ref = NfRef {
                 id_note,
-                account
+                account,
             };
             Ok((nf_ref, nf))
         })?;
@@ -467,6 +474,29 @@ impl DbAdapter {
         log::debug!("-get_ivk");
         Ok(ivk)
     }
+
+    pub fn get_diversifier(&self, account: u32) -> anyhow::Result<DiversifierIndex> {
+        let diversifier_index = self.connection.query_row(
+            "SELECT diversifier_index FROM diversifiers WHERE account = ?1",
+            params![account],
+            |row| {
+                let d: Vec<u8> = row.get(0)?;
+                let mut div = [0u8; 11];
+                div.copy_from_slice(&d);
+                Ok(div)
+            }
+        ).optional()?.unwrap_or_else(|| [0u8; 11]);
+        Ok(DiversifierIndex(diversifier_index))
+    }
+
+    pub fn store_diversifier(&self, account: u32, diversifier_index: &DiversifierIndex) -> anyhow::Result<()> {
+        let diversifier_bytes = diversifier_index.0.to_vec();
+        self.connection.execute(
+            "INSERT INTO diversifiers(account, diversifier_index) VALUES (?1, ?2) ON CONFLICT \
+            (account) DO UPDATE SET diversifier_index = excluded.diversifier_index",
+            params![account, diversifier_bytes])?;
+        Ok(())
+    }
 }
 
 #[cfg(test)]

diff --git a/src/wallet.rs b/src/wallet.rs
@@ -12,7 +12,7 @@ use tokio::sync::Mutex;
 use tonic::Request;
 use zcash_client_backend::address::RecipientAddress;
 use zcash_client_backend::data_api::wallet::ANCHOR_OFFSET;
-use zcash_client_backend::encoding::decode_extended_spending_key;
+use zcash_client_backend::encoding::{decode_extended_spending_key, decode_extended_full_viewing_key, encode_payment_address};
 use zcash_params::{OUTPUT_PARAMS, SPEND_PARAMS};
 use zcash_primitives::consensus::{BlockHeight, BranchId, Parameters};
 use zcash_primitives::transaction::builder::{Builder, Progress};
@@ -265,6 +265,17 @@ impl Wallet {
     pub fn get_ivk(&self, account: u32) -> anyhow::Result<String> {
         self.db.get_ivk(account)
     }
+
+    pub fn new_diversified_address(&self, account: u32) -> anyhow::Result<String> {
+        let ivk = self.get_ivk(account)?;
+        let fvk = decode_extended_full_viewing_key(NETWORK.hrp_sapling_extended_full_viewing_key(), &ivk)?.unwrap();
+        let mut diversifier_index = self.db.get_diversifier(account)?;
+        diversifier_index.increment().unwrap();
+        let (new_diversifier_index, pa) = fvk.address(diversifier_index).map_err(|_| anyhow::anyhow!("Cannot generate new address"))?;
+        self.db.store_diversifier(account, &new_diversifier_index)?;
+        let pa = encode_payment_address(NETWORK.hrp_sapling_payment_address(), &pa);
+        Ok(pa)
+    }
 }
 
 #[cfg(test)]
@@ -297,4 +308,11 @@ mod tests {
         // let tx_id = wallet.send_payment(1, &pa, 1000).await.unwrap();
         // println!("TXID = {}", tx_id);
     }
+
+    #[test]
+    pub fn test_diversified_address() {
+        let wallet = Wallet::new("zec.db");
+        let address = wallet.new_diversified_address(1).unwrap();
+        println!("{}", address);
+    }
 }