diff --git a/lib/src/blaze/fetch_taddr_txns.rs b/lib/src/blaze/fetch_taddr_txns.rs index 90fc80e6..bcb5e1c4 100644 --- a/lib/src/blaze/fetch_taddr_txns.rs +++ b/lib/src/blaze/fetch_taddr_txns.rs @@ -150,6 +150,7 @@ mod test { use zcash_primitives::transaction::{Transaction, TransactionData}; use crate::lightwallet::keys::Keys; + use crate::lightwallet::wallettkey::WalletTKey; use super::FetchTaddrTxns; @@ -158,7 +159,7 @@ mod test { // 5 t addresses let mut keys = Keys::new_empty(); let gened_taddrs: Vec<_> = (0..5).into_iter().map(|n| format!("taddr{}", n)).collect(); - keys.taddresses = gened_taddrs.clone(); + keys.tkeys = gened_taddrs.iter().map(|ta| WalletTKey::empty(ta)).collect::>(); let ftt = FetchTaddrTxns::new(Arc::new(RwLock::new(keys))); diff --git a/lib/src/lightclient.rs b/lib/src/lightclient.rs index 58c175a6..145d82b9 100644 --- a/lib/src/lightclient.rs +++ b/lib/src/lightclient.rs @@ -1,9 +1,15 @@ use self::lightclient_config::LightClientConfig; -use crate::{blaze::{ +use crate::{ + blaze::{ block_witness_data::BlockAndWitnessData, fetch_compact_blocks::FetchCompactBlocks, fetch_full_tx::FetchFullTxns, fetch_taddr_txns::FetchTaddrTxns, sync_status::SyncStatus, syncdata::BlazeSyncData, trial_decryptions::TrialDecryptions, update_notes::UpdateNotes, - }, compact_formats::RawTransaction, grpc_connector::GrpcConnector, lightclient::lightclient_config::MAX_REORG, lightwallet::{self, LightWallet, data::WalletTx, message::Message, now}}; + }, + compact_formats::RawTransaction, + grpc_connector::GrpcConnector, + lightclient::lightclient_config::MAX_REORG, + lightwallet::{self, data::WalletTx, message::Message, now, LightWallet}, +}; use futures::future::join_all; use json::{array, object, JsonValue}; use log::{error, info, warn}; @@ -928,14 +934,34 @@ impl LightClient { } else if key.starts_with(self.config.hrp_sapling_viewing_key()) { self.do_import_vk(key, birthday).await } else if key.starts_with("K") || key.starts_with("L") { - Err(format!("Can't import t-address keys yet!")) + self.do_import_tk(key).await } else { - Err(format!("'{}' was not recognized as either a spending key or a viewing key because it didn't start with either '{}' or '{}'", - key, self.config.hrp_sapling_private_key(), self.config.hrp_sapling_viewing_key())) + Err(format!( + "'{}' was not recognized as either a spending key or a viewing key", + key, + )) } } - /// Import a new private key + /// Import a new transparent private key + pub async fn do_import_tk(&self, sk: String) -> Result { + if !self.wallet.is_unlocked_for_spending().await { + error!("Wallet is locked"); + return Err("Wallet is locked".to_string()); + } + + let address = self.wallet.add_imported_tk(sk).await; + if address.starts_with("Error") { + let e = address; + error!("{}", e); + return Err(e); + } + + self.do_save().await?; + Ok(array![address]) + } + + /// Import a new z-address private key pub async fn do_import_sk(&self, sk: String, birthday: u64) -> Result { if !self.wallet.is_unlocked_for_spending().await { error!("Wallet is locked"); diff --git a/lib/src/lightclient/tests.rs b/lib/src/lightclient/tests.rs index b47fba49..e51965db 100644 --- a/lib/src/lightclient/tests.rs +++ b/lib/src/lightclient/tests.rs @@ -3,7 +3,6 @@ use group::GroupEncoding; use json::JsonValue; use jubjub::ExtendedPoint; use rand::rngs::OsRng; -use secp256k1::{PublicKey, Secp256k1}; use tempdir::TempDir; use tokio::runtime::Runtime; use tonic::transport::Channel; @@ -663,7 +662,6 @@ async fn z_incoming_viewkey() { #[tokio::test] async fn t_incoming_t_outgoing() { - let secp = Secp256k1::new(); let (data, config, ready_rx, stop_tx, h1) = create_test_server().await; ready_rx.await.unwrap(); @@ -675,9 +673,9 @@ async fn t_incoming_t_outgoing() { mine_random_blocks(&mut fcbl, &data, &lc, 10).await; // 2. Get an incoming tx to a t address - let sk = lc.wallet.keys().read().await.tkeys[0]; - let pk = PublicKey::from_secret_key(&secp, &sk); - let taddr = lc.wallet.keys().read().await.address_from_sk(&sk); + let sk = lc.wallet.keys().read().await.tkeys[0].clone(); + let pk = sk.pubkey().unwrap(); + let taddr = sk.address; let value = 100_000; let mut ftx = FakeTransaction::new(); @@ -762,7 +760,6 @@ async fn t_incoming_t_outgoing() { #[tokio::test] async fn mixed_txn() { - let secp = Secp256k1::new(); let (data, config, ready_rx, stop_tx, h1) = create_test_server().await; ready_rx.await.unwrap(); @@ -782,9 +779,9 @@ async fn mixed_txn() { mine_random_blocks(&mut fcbl, &data, &lc, 5).await; // 3. Send an incoming t-address txn - let sk = lc.wallet.keys().read().await.tkeys[0]; - let pk = PublicKey::from_secret_key(&secp, &sk); - let taddr = lc.wallet.keys().read().await.address_from_sk(&sk); + let sk = lc.wallet.keys().read().await.tkeys[0].clone(); + let pk = sk.pubkey().unwrap(); + let taddr = sk.address; let tvalue = 200_000; let mut ftx = FakeTransaction::new(); @@ -865,7 +862,6 @@ async fn mixed_txn() { #[tokio::test] async fn aborted_resync() { - let secp = Secp256k1::new(); let (data, config, ready_rx, stop_tx, h1) = create_test_server().await; ready_rx.await.unwrap(); @@ -885,9 +881,9 @@ async fn aborted_resync() { mine_random_blocks(&mut fcbl, &data, &lc, 5).await; // 3. Send an incoming t-address txn - let sk = lc.wallet.keys().read().await.tkeys[0]; - let pk = PublicKey::from_secret_key(&secp, &sk); - let taddr = lc.wallet.keys().read().await.address_from_sk(&sk); + let sk = lc.wallet.keys().read().await.tkeys[0].clone(); + let pk = sk.pubkey().unwrap(); + let taddr = sk.address; let tvalue = 200_000; let mut ftx = FakeTransaction::new(); @@ -976,7 +972,6 @@ async fn aborted_resync() { #[tokio::test] async fn no_change() { - let secp = Secp256k1::new(); let (data, config, ready_rx, stop_tx, h1) = create_test_server().await; ready_rx.await.unwrap(); @@ -996,9 +991,9 @@ async fn no_change() { mine_random_blocks(&mut fcbl, &data, &lc, 5).await; // 3. Send an incoming t-address txn - let sk = lc.wallet.keys().read().await.tkeys[0]; - let pk = PublicKey::from_secret_key(&secp, &sk); - let taddr = lc.wallet.keys().read().await.address_from_sk(&sk); + let sk = lc.wallet.keys().read().await.tkeys[0].clone(); + let pk = sk.pubkey().unwrap(); + let taddr = sk.address; let tvalue = 200_000; let mut ftx = FakeTransaction::new(); diff --git a/lib/src/lightwallet.rs b/lib/src/lightwallet.rs index f78336b0..a6d92c09 100644 --- a/lib/src/lightwallet.rs +++ b/lib/src/lightwallet.rs @@ -1,5 +1,6 @@ use crate::compact_formats::TreeState; use crate::lightwallet::data::WalletTx; +use crate::lightwallet::wallettkey::WalletTKey; use crate::{ blaze::fetch_full_tx::FetchFullTxns, lightclient::lightclient_config::LightClientConfig, @@ -51,6 +52,7 @@ pub(crate) mod keys; pub(crate) mod message; pub(crate) mod utils; pub(crate) mod wallet_txns; +pub(crate) mod wallettkey; mod walletzkey; pub fn now() -> u64 { @@ -104,7 +106,7 @@ pub struct WalletOptions { impl Default for WalletOptions { fn default() -> Self { WalletOptions { - download_memos: MemoDownloadOption::WalletMemos + download_memos: MemoDownloadOption::WalletMemos, } } } @@ -114,7 +116,6 @@ impl WalletOptions { return 1; } - pub fn read(mut reader: R) -> io::Result { let _version = reader.read_u64::()?; @@ -123,13 +124,14 @@ impl WalletOptions { 1 => MemoDownloadOption::WalletMemos, 2 => MemoDownloadOption::AllMemos, v => { - return Err(io::Error::new(io::ErrorKind::InvalidData, format!("Bad download option {}", v))); + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Bad download option {}", v), + )); } }; - Ok(Self { - download_memos - }) + Ok(Self { download_memos }) } pub fn write(&self, mut writer: W) -> io::Result<()> { @@ -471,6 +473,34 @@ impl LightWallet { } } + pub async fn add_imported_tk(&self, sk: String) -> String { + if self.keys.read().await.encrypted { + return "Error: Can't import transparent address key while wallet is encrypted".to_string(); + } + + let sk = match WalletTKey::from_sk_string(&self.config, sk) { + Err(e) => return format!("Error: {}", e), + Ok(k) => k, + }; + + let address = sk.address.clone(); + + if self + .keys + .read() + .await + .tkeys + .iter() + .find(|&tk| tk.address == address) + .is_some() + { + return "Error: Key already exists".to_string(); + } + + self.keys.write().await.tkeys.push(sk); + return address; + } + // Add a new imported spending key to the wallet /// NOTE: This will not rescan the wallet pub async fn add_imported_sk(&self, sk: String, birthday: u64) -> String { @@ -832,7 +862,6 @@ impl LightWallet { if highest_account.unwrap() == 0 { // Remove unused addresses self.keys.write().await.tkeys.truncate(1); - self.keys.write().await.taddresses.truncate(1); } } @@ -1326,7 +1355,6 @@ impl LightWallet { #[cfg(test)] mod test { - use secp256k1::{PublicKey, Secp256k1}; use zcash_primitives::transaction::components::Amount; use crate::{ @@ -1426,10 +1454,9 @@ mod test { assert_eq!(utxos.len(), 0); // 4. Get an incoming tx to a t address - let secp = Secp256k1::new(); - let sk = lc.wallet.keys().read().await.tkeys[0]; - let pk = PublicKey::from_secret_key(&secp, &sk); - let taddr = lc.wallet.keys().read().await.address_from_sk(&sk); + let sk = lc.wallet.keys().read().await.tkeys[0].clone(); + let pk = sk.pubkey().unwrap(); + let taddr = sk.address; let tvalue = 100_000; let mut ftx = FakeTransaction::new(); diff --git a/lib/src/lightwallet/keys.rs b/lib/src/lightwallet/keys.rs index cfbab5e2..a6e76c99 100644 --- a/lib/src/lightwallet/keys.rs +++ b/lib/src/lightwallet/keys.rs @@ -3,12 +3,11 @@ use std::{ io::{self, Error, ErrorKind, Read, Write}, }; -use base58::ToBase58; +use base58::{FromBase58, ToBase58}; use bip39::{Language, Mnemonic, Seed}; use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt}; use rand::{rngs::OsRng, Rng}; use ripemd160::Digest; -use secp256k1::SecretKey; use sha2::Sha256; use sodiumoxide::crypto::secretbox; use zcash_client_backend::{ @@ -24,13 +23,13 @@ use zcash_primitives::{ use crate::{ lightclient::lightclient_config::{LightClientConfig, GAP_RULE_UNUSED_ADDRESSES}, - lightwallet::{ - extended_key::{ExtendedPrivKey, KeyIndex}, - utils, - }, + lightwallet::utils, }; -use super::walletzkey::{WalletZKey, WalletZKeyType}; +use super::{ + wallettkey::{WalletTKey, WalletTKeyType}, + walletzkey::{WalletZKey, WalletZKeyType}, +}; /// Sha256(Sha256(value)) pub fn double_sha256(payload: &[u8]) -> Vec { @@ -61,6 +60,35 @@ impl ToBase58Check for [u8] { } } +/// A trait for converting base58check encoded values. +pub trait FromBase58Check { + /// Convert a value of `self`, interpreted as base58check encoded data, into the tuple with version and payload as bytes vector. + fn from_base58check(&self) -> io::Result<(u8, Vec)>; +} + +impl FromBase58Check for str { + fn from_base58check(&self) -> io::Result<(u8, Vec)> { + let mut payload: Vec = match self.from_base58() { + Ok(payload) => payload, + Err(error) => return Err(io::Error::new(ErrorKind::InvalidData, format!("{:?}", error))), + }; + if payload.len() < 5 { + return Err(io::Error::new( + ErrorKind::InvalidData, + format!("Invalid Checksum length"), + )); + } + + let checksum_index = payload.len() - 4; + let provided_checksum = payload.split_off(checksum_index); + let checksum = double_sha256(&payload)[..4].to_vec(); + if checksum != provided_checksum { + return Err(io::Error::new(ErrorKind::InvalidData, format!("Invalid Checksum"))); + } + Ok((payload[0], payload[1..].to_vec())) + } +} + // Manages all the keys in the wallet. Note that the RwLock for this is present in `lightwallet.rs`, so we'll // assume that this is already gone through a RwLock, so we don't lock any of the individual fields. pub struct Keys { @@ -85,14 +113,13 @@ pub struct Keys { pub(crate) zkeys: Vec, // Transparent keys. If the wallet is locked, then the secret keys will be encrypted, - // but the addresses will be present. - pub(crate) tkeys: Vec, - pub(crate) taddresses: Vec, + // but the addresses will be present. This Vec contains both wallet and imported tkeys + pub(crate) tkeys: Vec, } impl Keys { pub fn serialized_version() -> u64 { - return 20; + return 21; } #[cfg(test)] @@ -107,7 +134,6 @@ impl Keys { seed: [0u8; 32], zkeys: vec![], tkeys: vec![], - taddresses: vec![], } } @@ -136,8 +162,7 @@ impl Keys { let bip39_seed = Seed::new(&Mnemonic::from_entropy(&seed_bytes, Language::English).unwrap(), ""); // Derive only the first sk and address - let tpk = Self::get_taddr_from_bip39seed(&config, &bip39_seed.as_bytes(), 0); - let taddr = Self::address_from_prefix_sk(&config.base58_pubkey_address(), &tpk); + let tpk = WalletTKey::new_hdkey(config, 0, &bip39_seed.as_bytes()); let mut zkeys = vec![]; for hdkey_num in 0..num_zaddrs { @@ -154,7 +179,6 @@ impl Keys { seed: seed_bytes, zkeys, tkeys: vec![tpk], - taddresses: vec![taddr], }) } @@ -242,21 +266,33 @@ impl Keys { Vector::read(&mut reader, |r| WalletZKey::read(r))? }; - let tkeys = Vector::read(&mut reader, |r| { - let mut tpk_bytes = [0u8; 32]; - r.read_exact(&mut tpk_bytes)?; - secp256k1::SecretKey::from_slice(&tpk_bytes).map_err(|e| io::Error::new(ErrorKind::InvalidData, e)) - })?; + let tkeys = if version <= 20 { + let tkeys = Vector::read(&mut reader, |r| { + let mut tpk_bytes = [0u8; 32]; + r.read_exact(&mut tpk_bytes)?; + secp256k1::SecretKey::from_slice(&tpk_bytes).map_err(|e| io::Error::new(ErrorKind::InvalidData, e)) + })?; + + let taddresses = if version >= 4 { + // Read the addresses + Vector::read(&mut reader, |r| utils::read_string(r))? + } else { + // Calculate the addresses + tkeys + .iter() + .map(|sk| WalletTKey::address_from_prefix_sk(&config.base58_pubkey_address(), sk)) + .collect() + }; - let taddresses = if version >= 4 { - // Read the addresses - Vector::read(&mut reader, |r| utils::read_string(r))? - } else { - // Calculate the addresses tkeys .iter() - .map(|sk| Self::address_from_prefix_sk(&config.base58_pubkey_address(), sk)) - .collect() + .zip(taddresses.iter()) + .enumerate() + .map(|(i, (sk, taddr))| WalletTKey::from_raw(sk, taddr, i as u32)) + .collect::>() + } else { + // Read the TKeys + Vector::read(&mut reader, |r| WalletTKey::read(r))? }; Ok(Self { @@ -268,7 +304,6 @@ impl Keys { seed: seed_bytes, zkeys, tkeys, - taddresses, }) } @@ -295,13 +330,25 @@ impl Keys { let zkeys = Vector::read(&mut reader, |r| WalletZKey::read(r))?; - let tkeys = Vector::read(&mut reader, |r| { - let mut tpk_bytes = [0u8; 32]; - r.read_exact(&mut tpk_bytes)?; - secp256k1::SecretKey::from_slice(&tpk_bytes).map_err(|e| io::Error::new(ErrorKind::InvalidData, e)) - })?; + let tkeys = if version <= 20 { + let tkeys = Vector::read(&mut reader, |r| { + let mut tpk_bytes = [0u8; 32]; + r.read_exact(&mut tpk_bytes)?; + secp256k1::SecretKey::from_slice(&tpk_bytes).map_err(|e| io::Error::new(ErrorKind::InvalidData, e)) + })?; + + let taddresses = Vector::read(&mut reader, |r| utils::read_string(r))?; - let taddresses = Vector::read(&mut reader, |r| utils::read_string(r))?; + tkeys + .iter() + .zip(taddresses.iter()) + .enumerate() + .map(|(i, (sk, taddr))| WalletTKey::from_raw(sk, taddr, i as u32)) + .collect::>() + } else { + // Read the TKeys + Vector::read(&mut reader, |r| WalletTKey::read(r))? + }; Ok(Self { config: config.clone(), @@ -312,7 +359,6 @@ impl Keys { seed: seed_bytes, zkeys, tkeys, - taddresses, }) } @@ -339,10 +385,7 @@ impl Keys { Vector::write(&mut writer, &self.zkeys, |w, zk| zk.write(w))?; // Write the transparent private keys - Vector::write(&mut writer, &self.tkeys, |w, pk| w.write_all(&pk[..]))?; - - // Write the transparent addresses - Vector::write(&mut writer, &self.taddresses, |w, a| utils::write_string(w, a))?; + Vector::write(&mut writer, &self.tkeys, |w, sk| sk.write(w))?; Ok(()) } @@ -382,7 +425,7 @@ impl Keys { } pub fn get_all_taddrs(&self) -> Vec { - self.taddresses.iter().map(|t| t.clone()).collect() + self.tkeys.iter().map(|tk| tk.address.clone()).collect::>() } pub fn have_spending_key(&self, extfvk: &ExtendedFullViewingKey) -> bool { @@ -402,10 +445,9 @@ impl Keys { } pub fn get_taddr_to_sk_map(&self) -> HashMap { - self.taddresses + self.tkeys .iter() - .zip(self.tkeys.iter()) - .map(|(addr, sk)| (addr.clone(), sk.clone())) + .map(|tk| (tk.address.clone(), tk.key.unwrap().clone())) .collect() } @@ -416,11 +458,12 @@ impl Keys { } let last_addresses = { - self.taddresses + self.tkeys .iter() + .filter(|tk| tk.keytype == WalletTKeyType::HdKey) .rev() .take(GAP_RULE_UNUSED_ADDRESSES) - .map(|s| s.clone()) + .map(|s| s.address.clone()) .collect::>() }; @@ -509,14 +552,19 @@ impl Keys { return "Error: Can't add key while wallet is locked".to_string(); } - let pos = self.tkeys.len() as u32; - let bip39_seed = bip39::Seed::new(&Mnemonic::from_entropy(&self.seed, Language::English).unwrap(), ""); + // Find the highest pos we have + let pos = self + .tkeys + .iter() + .filter(|sk| sk.hdkey_num.is_some()) + .max_by(|sk1, sk2| sk1.hdkey_num.unwrap().cmp(&sk2.hdkey_num.unwrap())) + .map_or(0, |sk| sk.hdkey_num.unwrap() + 1); - let sk = Self::get_taddr_from_bip39seed(&self.config, &bip39_seed.as_bytes(), pos); - let address = self.address_from_sk(&sk); + let bip39_seed = bip39::Seed::new(&Mnemonic::from_entropy(&self.seed, Language::English).unwrap(), ""); - self.tkeys.push(sk); - self.taddresses.push(address.clone()); + let key = WalletTKey::new_hdkey(&self.config, pos, &bip39_seed.as_bytes()); + let address = key.address.clone(); + self.tkeys.push(key); address } @@ -553,12 +601,7 @@ impl Keys { pub fn get_t_secret_keys(&self) -> Vec<(String, String)> { self.tkeys .iter() - .map(|sk| { - ( - self.address_from_sk(sk), - sk[..].to_base58check(&self.config.base58_secretkey_prefix(), &[0x01]), - ) - }) + .map(|sk| (sk.address.clone(), sk.sk_as_string(&self.config).unwrap_or_default())) .collect::>() } @@ -582,6 +625,11 @@ impl Keys { .map(|k| k.encrypt(&key)) .collect::>>()?; + self.tkeys + .iter_mut() + .map(|k| k.encrypt(&key)) + .collect::>>()?; + self.encrypted = true; self.lock()?; @@ -599,9 +647,13 @@ impl Keys { // Empty the seed and the secret keys self.seed.copy_from_slice(&[0u8; 32]); - self.tkeys = vec![]; - // Remove all the private key from the zkeys + // Remove all the private key from the zkeys and tkeys + self.tkeys + .iter_mut() + .map(|tk| tk.lock()) + .collect::>>()?; + self.zkeys .iter_mut() .map(|zk| zk.lock()) @@ -641,32 +693,20 @@ impl Keys { // The seed bytes is the raw entropy. To pass it to HD wallet generation, // we need to get the 64 byte bip39 entropy let bip39_seed = bip39::Seed::new(&Mnemonic::from_entropy(&seed, Language::English).unwrap(), ""); + let config = self.config.clone(); // Transparent keys - let mut tkeys = vec![]; - for pos in 0..self.taddresses.len() { - let sk = Self::get_taddr_from_bip39seed(&self.config, &bip39_seed.as_bytes(), pos as u32); - let address = self.address_from_sk(&sk); - - if address != self.taddresses[pos] { - return Err(io::Error::new( - ErrorKind::InvalidData, - format!("taddress mismatch at {}. {} vs {}", pos, address, self.taddresses[pos]), - )); - } - - tkeys.push(sk); - } + self.tkeys + .iter_mut() + .map(|tk| tk.unlock(&config, bip39_seed.as_bytes(), &key)) + .collect::>>()?; - let config = self.config.clone(); // Go over the zkeys, and add the spending keys again self.zkeys .iter_mut() .map(|zk| zk.unlock(&config, bip39_seed.as_bytes(), &key)) .collect::>>()?; - // Everything checks out, so we'll update our wallet with the decrypted values - self.tkeys = tkeys; self.seed.copy_from_slice(&seed); self.encrypted = true; @@ -687,7 +727,12 @@ impl Keys { self.unlock(passwd)?; } - // Remove encryption from individual zkeys + // Remove encryption from individual zkeys and tkeys + self.tkeys + .iter_mut() + .map(|tk| tk.remove_encryption()) + .collect::>>()?; + self.zkeys .iter_mut() .map(|zk| zk.remove_encryption()) @@ -722,39 +767,6 @@ impl Keys { } } - pub fn address_from_sk(&self, sk: &secp256k1::SecretKey) -> String { - Self::address_from_prefix_sk(&self.config.base58_pubkey_address(), sk) - } - - pub fn address_from_prefix_sk(prefix: &[u8; 2], sk: &secp256k1::SecretKey) -> String { - let secp = secp256k1::Secp256k1::new(); - let pk = secp256k1::PublicKey::from_secret_key(&secp, &sk); - - // Encode into t address - let mut hash160 = ripemd160::Ripemd160::new(); - hash160.update(Sha256::digest(&pk.serialize()[..].to_vec())); - - hash160.finalize().to_base58check(prefix, &[]) - } - - pub fn get_taddr_from_bip39seed(config: &LightClientConfig, bip39_seed: &[u8], pos: u32) -> SecretKey { - assert_eq!(bip39_seed.len(), 64); - - let ext_t_key = ExtendedPrivKey::with_seed(bip39_seed).unwrap(); - ext_t_key - .derive_private_key(KeyIndex::hardened_from_normalize_index(44).unwrap()) - .unwrap() - .derive_private_key(KeyIndex::hardened_from_normalize_index(config.get_coin_type()).unwrap()) - .unwrap() - .derive_private_key(KeyIndex::hardened_from_normalize_index(0).unwrap()) - .unwrap() - .derive_private_key(KeyIndex::Normal(0)) - .unwrap() - .derive_private_key(KeyIndex::Normal(pos)) - .unwrap() - .private_key - } - pub fn get_zaddr_from_bip39seed( config: &LightClientConfig, bip39_seed: &[u8], diff --git a/lib/src/lightwallet/wallettkey.rs b/lib/src/lightwallet/wallettkey.rs new file mode 100644 index 00000000..e210190b --- /dev/null +++ b/lib/src/lightwallet/wallettkey.rs @@ -0,0 +1,364 @@ +use std::io::{self, Error, ErrorKind, Read, Write}; + +use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt}; +use ripemd160::Digest; +use secp256k1::SecretKey; +use sha2::Sha256; +use sodiumoxide::crypto::secretbox; +use zcash_primitives::serialize::{Optional, Vector}; + +use crate::{ + lightclient::lightclient_config::LightClientConfig, + lightwallet::extended_key::{ExtendedPrivKey, KeyIndex}, +}; + +use super::{ + keys::{FromBase58Check, ToBase58Check}, + utils, +}; + +#[derive(Debug, PartialEq, Clone, Copy)] +pub enum WalletTKeyType { + HdKey = 0, + ImportedKey = 1, +} + +#[derive(Debug, Clone)] +pub struct WalletTKey { + pub(super) keytype: WalletTKeyType, + locked: bool, + pub(super) key: Option, + pub(crate) address: String, + + // If this is a HD key, what is the key number + pub(super) hdkey_num: Option, + + // If locked, the encrypted private key is stored here + enc_key: Option>, + nonce: Option>, +} + +impl WalletTKey { + pub fn get_taddr_from_bip39seed(config: &LightClientConfig, bip39_seed: &[u8], pos: u32) -> secp256k1::SecretKey { + assert_eq!(bip39_seed.len(), 64); + + let ext_t_key = ExtendedPrivKey::with_seed(bip39_seed).unwrap(); + ext_t_key + .derive_private_key(KeyIndex::hardened_from_normalize_index(44).unwrap()) + .unwrap() + .derive_private_key(KeyIndex::hardened_from_normalize_index(config.get_coin_type()).unwrap()) + .unwrap() + .derive_private_key(KeyIndex::hardened_from_normalize_index(0).unwrap()) + .unwrap() + .derive_private_key(KeyIndex::Normal(0)) + .unwrap() + .derive_private_key(KeyIndex::Normal(pos)) + .unwrap() + .private_key + } + + pub fn address_from_prefix_sk(prefix: &[u8; 2], sk: &secp256k1::SecretKey) -> String { + let secp = secp256k1::Secp256k1::new(); + let pk = secp256k1::PublicKey::from_secret_key(&secp, &sk); + + // Encode into t address + let mut hash160 = ripemd160::Ripemd160::new(); + hash160.update(Sha256::digest(&pk.serialize()[..].to_vec())); + + hash160.finalize().to_base58check(prefix, &[]) + } + + pub fn from_raw(sk: &secp256k1::SecretKey, taddr: &String, num: u32) -> Self { + WalletTKey { + keytype: WalletTKeyType::HdKey, + key: Some(sk.clone()), + address: taddr.clone(), + hdkey_num: Some(num), + locked: false, + enc_key: None, + nonce: None, + } + } + + pub fn from_sk_string(config: &LightClientConfig, sks: String) -> io::Result { + let (_v, mut bytes) = sks.as_str().from_base58check()?; + let suffix = bytes.split_off(32); + + // Assert the suffix + if suffix.len() !=1 || suffix[0] != 0x01 { + return Err(io::Error::new(ErrorKind::InvalidData, format!("Invalid Suffix: {:?}", suffix))); + } + + let key = SecretKey::from_slice(&bytes).map_err(|e| io::Error::new(ErrorKind::InvalidData, e))?; + let address = Self::address_from_prefix_sk(&config.base58_pubkey_address(), &key); + + Ok(WalletTKey { + keytype: WalletTKeyType::ImportedKey, + key: Some(key), + address, + hdkey_num: None, + locked: false, + enc_key: None, + nonce: None, + }) + } + + pub fn new_hdkey(config: &LightClientConfig, hdkey_num: u32, bip39_seed: &[u8]) -> Self { + let pos = hdkey_num; + + let sk = Self::get_taddr_from_bip39seed(&config, bip39_seed, pos); + let address = Self::address_from_prefix_sk(&config.base58_pubkey_address(), &sk); + + WalletTKey { + keytype: WalletTKeyType::HdKey, + key: Some(sk), + address, + hdkey_num: Some(hdkey_num), + locked: false, + enc_key: None, + nonce: None, + } + } + + // Return the wallet string representation of a secret key + pub fn sk_as_string(&self, config: &LightClientConfig) -> io::Result { + if self.key.is_none() { + return Err(io::Error::new(ErrorKind::NotFound, "Wallet locked")); + } + + Ok(self.key.as_ref().unwrap()[..].to_base58check(&config.base58_secretkey_prefix(), &[0x01])) + } + + #[cfg(test)] + pub fn empty(ta: &String) -> Self { + WalletTKey { + keytype: WalletTKeyType::HdKey, + key: None, + address: ta.clone(), + hdkey_num: None, + locked: false, + enc_key: None, + nonce: None, + } + } + + #[cfg(test)] + pub fn pubkey(&self) -> io::Result { + if self.key.is_none() { + return Err(io::Error::new(ErrorKind::NotFound, "Wallet locked")); + } + + Ok(secp256k1::PublicKey::from_secret_key( + &secp256k1::Secp256k1::new(), + &self.key.unwrap(), + )) + } + + fn serialized_version() -> u8 { + return 1; + } + + pub fn read(mut inp: R) -> io::Result { + let version = inp.read_u8()?; + assert!(version <= Self::serialized_version()); + + let keytype: WalletTKeyType = match inp.read_u32::()? { + 0 => Ok(WalletTKeyType::HdKey), + 1 => Ok(WalletTKeyType::ImportedKey), + n => Err(io::Error::new( + ErrorKind::InvalidInput, + format!("Unknown zkey type {}", n), + )), + }?; + + let locked = inp.read_u8()? > 0; + + let key = Optional::read(&mut inp, |r| { + let mut tpk_bytes = [0u8; 32]; + r.read_exact(&mut tpk_bytes)?; + secp256k1::SecretKey::from_slice(&tpk_bytes).map_err(|e| io::Error::new(ErrorKind::InvalidData, e)) + })?; + let address = utils::read_string(&mut inp)?; + + let hdkey_num = Optional::read(&mut inp, |r| r.read_u32::())?; + + let enc_key = Optional::read(&mut inp, |r| Vector::read(r, |r| r.read_u8()))?; + let nonce = Optional::read(&mut inp, |r| Vector::read(r, |r| r.read_u8()))?; + + Ok(WalletTKey { + keytype, + locked, + key, + address, + hdkey_num, + enc_key, + nonce, + }) + } + + pub fn write(&self, mut out: W) -> io::Result<()> { + out.write_u8(Self::serialized_version())?; + + out.write_u32::(self.keytype.clone() as u32)?; + + out.write_u8(self.locked as u8)?; + + Optional::write(&mut out, &self.key, |w, sk| w.write_all(&sk[..]))?; + utils::write_string(&mut out, &self.address)?; + + Optional::write(&mut out, &self.hdkey_num, |o, n| o.write_u32::(*n))?; + + // Write enc_key + Optional::write(&mut out, &self.enc_key, |o, v| { + Vector::write(o, v, |o, n| o.write_u8(*n)) + })?; + + // Write nonce + Optional::write(&mut out, &self.nonce, |o, v| Vector::write(o, v, |o, n| o.write_u8(*n))) + } + + pub fn lock(&mut self) -> io::Result<()> { + match self.keytype { + WalletTKeyType::HdKey => { + // For HD keys, just empty out the keys, since they will be reconstructed from the hdkey_num + self.key = None; + self.locked = true; + } + WalletTKeyType::ImportedKey => { + // For imported keys, encrypt the key into enckey + // assert that we have the encrypted key. + if self.enc_key.is_none() { + return Err(Error::new( + ErrorKind::InvalidInput, + "Can't lock when imported key is not encrypted", + )); + } + self.key = None; + self.locked = true; + } + } + + Ok(()) + } + + pub fn unlock(&mut self, config: &LightClientConfig, bip39_seed: &[u8], key: &secretbox::Key) -> io::Result<()> { + match self.keytype { + WalletTKeyType::HdKey => { + let sk = Self::get_taddr_from_bip39seed(&config, &bip39_seed, self.hdkey_num.unwrap()); + let address = Self::address_from_prefix_sk(&config.base58_pubkey_address(), &sk); + + if address != self.address { + return Err(io::Error::new( + ErrorKind::InvalidData, + format!( + "address mismatch at {}. {:?} vs {:?}", + self.hdkey_num.unwrap(), + address, + self.address + ), + )); + } + + self.key = Some(sk) + } + WalletTKeyType::ImportedKey => { + // For imported keys, we need to decrypt from the encrypted key + let nonce = secretbox::Nonce::from_slice(&self.nonce.as_ref().unwrap()).unwrap(); + let sk_bytes = match secretbox::open(&self.enc_key.as_ref().unwrap(), &nonce, &key) { + Ok(s) => s, + Err(_) => { + return Err(io::Error::new( + ErrorKind::InvalidData, + "Decryption failed. Is your password correct?", + )); + } + }; + + let key = secp256k1::SecretKey::from_slice(&sk_bytes[..]) + .map_err(|e| io::Error::new(ErrorKind::InvalidData, e))?; + self.key = Some(key); + } + }; + + self.locked = false; + Ok(()) + } + + pub fn encrypt(&mut self, key: &secretbox::Key) -> io::Result<()> { + match self.keytype { + WalletTKeyType::HdKey => { + // For HD keys, we don't need to do anything, since the hdnum has all the info to recreate this key + } + WalletTKeyType::ImportedKey => { + // For imported keys, encrypt the key into enckey + let nonce = secretbox::gen_nonce(); + + let sk_bytes = &self.key.as_ref().unwrap()[..]; + + self.enc_key = Some(secretbox::seal(&sk_bytes, &nonce, &key)); + self.nonce = Some(nonce.as_ref().to_vec()); + } + } + + // Also lock after encrypt + self.lock() + } + + pub fn remove_encryption(&mut self) -> io::Result<()> { + if self.locked { + return Err(Error::new( + ErrorKind::InvalidInput, + "Can't remove encryption while locked", + )); + } + + match self.keytype { + WalletTKeyType::HdKey => { + // For HD keys, we don't need to do anything, since the hdnum has all the info to recreate this key + Ok(()) + } + WalletTKeyType::ImportedKey => { + self.enc_key = None; + self.nonce = None; + Ok(()) + } + } + } +} + +#[cfg(test)] +mod test { + + use rand::{rngs::OsRng, Rng}; + use secp256k1::SecretKey; + + use crate::lightclient::lightclient_config::LightClientConfig; + + use super::WalletTKey; + + #[test] + fn tkey_encode_decode() { + let config = LightClientConfig::create_unconnected("main".to_string(), None); + + for _i in 0..10 { + let mut b = [0u8; 32]; + + // Gen random key + OsRng.fill(&mut b); + let sk = SecretKey::from_slice(&b).unwrap(); + let address = WalletTKey::address_from_prefix_sk(&config.base58_pubkey_address(), &sk); + let wtk = WalletTKey::from_raw(&sk, &address, 0); + + // export private key + let sks = wtk.sk_as_string(&config).unwrap(); + // println!("key:{}", sks); + + // Import it back + let wtk2 = WalletTKey::from_sk_string(&config, sks).unwrap(); + + // Make sure they're the same + assert_eq!(wtk.address, wtk2.address); + assert_eq!(wtk.key.unwrap(), wtk2.key.unwrap()); + } + } +}