Skip to content

Commit

Permalink
refactor, simplify, specify error type
Browse files Browse the repository at this point in the history
  • Loading branch information
Vladislav Melnik committed Apr 15, 2019
1 parent f4ec299 commit 1edae61
Show file tree
Hide file tree
Showing 2 changed files with 70 additions and 126 deletions.
91 changes: 40 additions & 51 deletions wallet/src/account.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
//!
use bitcoin::{
util::{
bip32::{ExtendedPubKey, ExtendedPrivKey, ChildNumber},
bip32::{ExtendedPubKey, ExtendedPrivKey, ChildNumber, Error as Bip32Error},
address::Address,
},
blockdata::{
Expand All @@ -37,7 +37,6 @@ use super::DB;

use std::{
sync::{Arc, RwLock},
error::Error,
collections::HashMap,
};

Expand Down Expand Up @@ -80,7 +79,7 @@ impl Into<AccountAddressType> for usize {
0 => AccountAddressType::P2PKH,
1 => AccountAddressType::P2SHWH,
2 => AccountAddressType::P2WKH,
_ => unreachable!(),
_ => panic!("unknown address code: {}", self),
}
}
}
Expand All @@ -91,15 +90,6 @@ pub enum AddressChain {
Internal,
}

impl From<AddressChain> for usize {
fn from(val: AddressChain) -> usize {
match val {
AddressChain::External => 0,
AddressChain::Internal => 1,
}
}
}

impl Into<u32> for AddressChain {
fn into(self) -> u32 {
match self {
Expand Down Expand Up @@ -143,7 +133,7 @@ impl Utxo {
pk_script: Script,
addr_type: AccountAddressType,
) -> Self {
Self {
Utxo {
value,
key_path,
out_point,
Expand Down Expand Up @@ -180,7 +170,7 @@ pub struct SecretKeyHelper {

impl SecretKeyHelper {
fn new(addr_type: AccountAddressType, addr_chain: AddressChain, index: u32) -> Self {
Self {
SecretKeyHelper {
addr_type,
addr_chain,
index,
Expand Down Expand Up @@ -237,7 +227,7 @@ impl Account {
&self.utxo_list
}

pub fn next_external_pk(&mut self) -> Result<PublicKey, Box<dyn Error>> {
pub fn next_external_pk(&mut self) -> Result<PublicKey, Bip32Error> {
let path = &[
ChildNumber::Normal { index: 0 }, // TODO(evg): use addr chain enum instead?
ChildNumber::Normal {
Expand Down Expand Up @@ -265,7 +255,7 @@ impl Account {
Ok(extended_pub_key.public_key)
}

pub fn next_internal_pk(&mut self) -> Result<PublicKey, Box<dyn Error>> {
pub fn next_internal_pk(&mut self) -> Result<PublicKey, Bip32Error> {
let path = &[
ChildNumber::Normal { index: 1 },
ChildNumber::Normal {
Expand Down Expand Up @@ -294,6 +284,21 @@ impl Account {
}

pub fn addr_from_pk(&self, pk: &PublicKey) -> String {
fn p2pkh_addr_from_public_key(pk: &PublicKey, network: Network) -> String {
let addr = Address::p2pkh(pk, network);
addr.to_string()
}

fn p2shwh_addr_from_public_key(pk: &PublicKey, network: Network) -> String {
let addr = Address::p2shwpkh(pk, network);
addr.to_string()
}

fn p2wkh_addr_from_public_key_bip0173(pk: &PublicKey, network: Network) -> String {
let addr = Address::p2wpkh(pk, network);
addr.to_string()
}

match self.address_type {
AccountAddressType::P2PKH => p2pkh_addr_from_public_key(pk, self.network),
AccountAddressType::P2SHWH => p2shwh_addr_from_public_key(pk, self.network),
Expand All @@ -302,14 +307,29 @@ impl Account {
}

pub fn script_from_pk(&self, pk: &PublicKey) -> Script {
fn p2pkh_script_from_public_key(pk: &PublicKey, network: Network) -> Script {
let p2pkh_addr = Address::p2pkh(pk, network);
p2pkh_addr.script_pubkey()
}

fn p2shwh_script_from_public_key(pk: &PublicKey, network: Network) -> Script {
let addr = Address::p2shwpkh(pk, network);
addr.script_pubkey()
}

fn p2wkh_script_from_public_key(pk: &PublicKey, network: Network) -> Script {
let p2wkh_addr = Address::p2wpkh(pk, network);
p2wkh_addr.script_pubkey()
}

match self.address_type {
AccountAddressType::P2PKH => p2pkh_script_from_public_key(pk, self.network),
AccountAddressType::P2SHWH => p2shwh_script_from_public_key(pk, self.network),
AccountAddressType::P2WKH => p2wkh_script_from_public_key(pk, self.network),
}
}

pub fn new_address(&mut self) -> Result<String, Box<dyn Error>> {
pub fn new_address(&mut self) -> Result<String, Bip32Error> {
let pk = self.next_external_pk()?;
let addr = self.addr_from_pk(&pk);
self.btc_address_list.push(addr.clone());
Expand All @@ -320,7 +340,7 @@ impl Account {
Ok(addr)
}

pub fn new_change_address(&mut self) -> Result<String, Box<dyn Error>> {
pub fn new_change_address(&mut self) -> Result<String, Bip32Error> {
let pk = self.next_internal_pk()?;
let addr = self.addr_from_pk(&pk);
self.btc_address_list.push(addr.clone());
Expand All @@ -332,42 +352,12 @@ impl Account {
}
}

pub fn p2pkh_addr_from_public_key(pk: &PublicKey, network: Network) -> String {
let addr = Address::p2pkh(pk, network);
addr.to_string()
}

pub fn p2shwh_addr_from_public_key(pk: &PublicKey, network: Network) -> String {
let addr = Address::p2shwpkh(pk, network);
addr.to_string()
}

pub fn p2wkh_addr_from_public_key_bip0173(pk: &PublicKey, network: Network) -> String {
let addr = Address::p2wpkh(pk, network);
addr.to_string()
}

pub fn p2pkh_script_from_public_key(pk: &PublicKey, network: Network) -> Script {
let p2pkh_addr = Address::p2pkh(pk, network);
p2pkh_addr.script_pubkey()
}

pub fn p2shwh_script_from_public_key(pk: &PublicKey, network: Network) -> Script {
let addr = Address::p2shwpkh(pk, network);
addr.script_pubkey()
}

pub fn p2wkh_script_from_public_key(pk: &PublicKey, network: Network) -> Script {
let p2wkh_addr = Address::p2wpkh(pk, network);
p2wkh_addr.script_pubkey()
}

#[cfg(test)]
mod test {
use bitcoin::{
network::constants::Network,
Block, Transaction,
};
};
use bitcoin_hashes::sha256d::Hash as Sha256dHash;
use std::{fmt, error::Error};

Expand All @@ -387,8 +377,7 @@ mod test {
}
}

impl Error for FakeBlockChainIoError {
}
impl Error for FakeBlockChainIoError {}

impl BlockChainIO for FakeBlockChainIO {
type Error = FakeBlockChainIoError;
Expand Down
105 changes: 30 additions & 75 deletions wallet/src/walletlibrary.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -257,14 +257,19 @@ pub struct WalletLibrary {

impl WalletLibraryInterface for WalletLibrary {
fn new_address(&mut self, address_type: AccountAddressType) -> Result<String, Box<dyn Error>> {
self.get_account_mut(address_type).new_address()
self.get_account_mut(address_type)
.new_address()
// converts Bip32Error into `Box<dyn Error>`
.map_err(Into::into)
}

fn new_change_address(
&mut self,
address_type: AccountAddressType,
) -> Result<String, Box<dyn Error>> {
self.get_account_mut(address_type).new_change_address()
self.get_account_mut(address_type)
.new_change_address()
.map_err(Into::into)
}

fn get_utxo_list(&self) -> Vec<Utxo> {
Expand Down Expand Up @@ -418,10 +423,7 @@ impl WalletLibraryInterface for WalletLibrary {

// TODO(evg): use SigHashType enum
let hash = tx.signature_hash(i, &pk_script, 0x1);
let signature = ctx.sign(
&Message::from_slice(&hash[..]).unwrap(),
&sk.key,
);
let signature = ctx.sign(&Message::from_slice(&hash[..]).unwrap(), &sk.key);

let mut serialized_sig = signature.serialize_der();
serialized_sig.push(0x1);
Expand All @@ -442,7 +444,8 @@ impl WalletLibraryInterface for WalletLibrary {
utxo.value,
);

let signature = ctx.sign(&Message::from_slice(&tx_sig_hash[..]).unwrap(), &sk.key);
let signature =
ctx.sign(&Message::from_slice(&tx_sig_hash[..]).unwrap(), &sk.key);

let mut serialized_sig = signature.serialize_der();
serialized_sig.push(0x1);
Expand All @@ -463,7 +466,8 @@ impl WalletLibraryInterface for WalletLibrary {
utxo.value,
);

let signature = ctx.sign(&Message::from_slice(&tx_sig_hash[..]).unwrap(), &sk.key);
let signature =
ctx.sign(&Message::from_slice(&tx_sig_hash[..]).unwrap(), &sk.key);

let mut serialized_sig = signature.serialize_der();
serialized_sig.push(0x1);
Expand Down Expand Up @@ -534,99 +538,49 @@ impl WalletLibraryInterface for WalletLibrary {
&mut self.p2wkh_account,
];
for (account_index, account) in account_list.iter_mut().enumerate() {
for output_index in 0..tx.output.len() {
let output = &tx.output[output_index];
for (output_index, output) in tx.output.iter().enumerate() {
let actual = &output.script_pubkey.to_bytes();
let mut joined = account.external_pk_list.clone();
joined.extend_from_slice(&account.internal_pk_list);

// TODO(evg): something better?
let external_pk_list_len = account.external_pk_list.len();
let get_pk_index = |mut raw: usize| -> (usize, AddressChain) {
let mut addr_chain = AddressChain::External;
if raw >= external_pk_list_len {
raw -= external_pk_list_len;
addr_chain = AddressChain::Internal;
}
(raw, addr_chain)
let get_pk_index = |raw: usize| -> KeyPath {
let cache = if raw >= external_pk_list_len {
(raw - external_pk_list_len, AddressChain::Internal)
} else {
(raw, AddressChain::External)
};
KeyPath::new(cache.1, cache.0 as u32)
};

let op = OutPoint {
txid: tx.txid(),
vout: output_index as u32,
};

if output.script_pubkey.is_p2pkh()
&& account.address_type == AccountAddressType::P2PKH
if (output.script_pubkey.is_p2pkh()
&& account.address_type == AccountAddressType::P2PKH)
|| (output.script_pubkey.is_p2sh()
&& account.address_type == AccountAddressType::P2SHWH)
|| (output.script_pubkey.is_v0_p2wpkh()
&& account.address_type == AccountAddressType::P2WKH)
{
// TODO(evg): use correct index
for pk_index in 0..joined.len() {
let pk = &joined[pk_index];
let script = account.script_from_pk(pk);
let expected = &script.to_bytes();
if actual == expected {
let cache = get_pk_index(pk_index);
let key_path = KeyPath::new(cache.1, cache.0 as u32);

let utxo = Utxo::new(
output.value,
key_path,
op,
account_index as u32,
script,
AccountAddressType::P2PKH,
);

account.grab_utxo(utxo.clone());
self.op_to_utxo.insert(op, utxo);
}
}
}

if output.script_pubkey.is_p2sh()
&& account.address_type == AccountAddressType::P2SHWH
{
for pk_index in 0..joined.len() {
let pk = &joined[pk_index];
let script = account.script_from_pk(pk);
let expected = &script.to_bytes();
if actual == expected {
let cache = get_pk_index(pk_index);
let key_path = KeyPath::new(cache.1, cache.0 as u32);

let utxo = Utxo::new(
output.value,
key_path,
op,
account_index as u32,
script,
AccountAddressType::P2SHWH,
);

account.grab_utxo(utxo.clone());
self.op_to_utxo.insert(op, utxo);
}
}
}

if output.script_pubkey.is_v0_p2wpkh()
&& account.address_type == AccountAddressType::P2WKH
{
for pk_index in 0..joined.len() {
let pk = &joined[pk_index];
let script = account.script_from_pk(pk);
let expected = &script.to_bytes();
if actual == expected {
let cache = get_pk_index(pk_index);
let key_path = KeyPath::new(cache.1, cache.0 as u32);
let key_path = get_pk_index(pk_index);

let utxo = Utxo::new(
output.value,
key_path,
op,
account_index as u32,
script,
AccountAddressType::P2WKH,
account.address_type.clone(),
);

account.grab_utxo(utxo.clone());
Expand Down Expand Up @@ -666,7 +620,8 @@ impl WalletLibrary {
(master_key, mnemonic)
}
WalletLibraryMode::Decrypt => {
let randomness = db.get_bip39_randomness()
let randomness = db
.get_bip39_randomness()
.ok_or(WalletError::HasNoWalletInDatabase)?;
let (master_key, mnemonic) =
KeyFactory::decrypt(&randomness, wc.network, &wc.passphrase, &wc.salt)?;
Expand Down

0 comments on commit 1edae61

Please sign in to comment.