milksad
/
rust-bitcoin-unsafe-fast
14
0
Fork 0
Code Issues Pull Requests Activity

Add wallet balance-checking support

This commit is contained in:
Andrew Poelstra 2014-09-10 07:19:12 -05:00
parent ef11e8273b
commit 184721db61
2 changed files with 118 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
src/wallet/address_index.rs
View File

@ -19,34 +19,67 @@
//!
use std::collections::HashMap;
use collections::hash::sip::hash_with_keys;
use blockdata::transaction::{TxOut, PayToPubkeyHash};
use blockdata::utxoset::UtxoSet;
use blockdata::script::Script;
use network::constants::Network;
use wallet::address::Address;
use wallet::wallet::Wallet;
use util::hash::Sha256dHash;
/// An address index
#[deriving(Clone, PartialEq, Eq, Show)]
pub struct AddressIndex {
index: HashMap<Script, Vec<(Sha256dHash, uint)>>
index: HashMap<Script, Vec<(Sha256dHash, uint, TxOut)>>,
network: Network,
k1: u64,
k2: u64
}
impl AddressIndex {
/// Creates a new address index from a wallet (which provides an authenticated
/// hash function for prefix filtering) and UTXO set (which is what gets filtered).
pub fn new(utxo_set: &UtxoSet, wallet: &Wallet) -> AddressIndex {
let (k1, k2) = wallet.siphash_key();
let mut ret = AddressIndex {
index: HashMap::with_capacity(utxo_set.n_utxos() / 256)
index: HashMap::with_capacity(utxo_set.n_utxos() / 256),
network: wallet.network(),
k1: k1,
k2: k2
};
for (key, idx, txo) in utxo_set.iter() {
if wallet.might_be_mine(txo) {
if ret.admissible_txo(txo) {
ret.index.insert_or_update_with(txo.script_pubkey.clone(),
vec![(key, idx)],
|_, v| v.push((key, idx)));
vec![(key, idx, txo.clone())],
|_, v| v.push((key, idx, txo.clone())));
}
}
ret
}
/// A filtering function used for creating a small address index.
#[inline]
pub fn admissible_address(&self, addr: &Address) -> bool {
hash_with_keys(self.k1, self.k2, &addr.as_slice()) & 0xFF == 0
}
/// A filtering function used for creating a small address index.
#[inline]
pub fn admissible_txo(&self, out: &TxOut) -> bool {
match out.classify(self.network) {
PayToPubkeyHash(addr) => self.admissible_address(&addr),
_ => false
}
}
/// Lookup a txout by its scriptpubkey. Returns a slice because there
/// may be more than one for any given scriptpubkey.
#[inline]
pub fn find_by_script<'a>(&'a self, pubkey: &Script) -> &'a [(Sha256dHash, uint, TxOut)] {
self.index.find(pubkey).map(|v| v.as_slice()).unwrap_or(&[])
}
}

100
src/wallet/wallet.rs
View File

@ -19,15 +19,15 @@
use std::collections::HashMap;
use std::default::Default;
use std::io::extensions::u64_from_be_bytes;
use collections::hash::sip::hash_with_keys;
use serialize::{Decoder, Decodable, Encoder, Encodable};
use secp256k1::key::PublicKey;
use blockdata::transaction::{PayToPubkeyHash, TxOut};
use blockdata::utxoset::UtxoSet;
use network::constants::Network;
use wallet::bip32::{mod, ChildNumber, ExtendedPrivKey, Normal, Hardened};
use wallet::bip32::{mod, ChildNumber, ExtendedPrivKey, ExtendedPubKey, Normal, Hardened};
use wallet::address::Address;
use wallet::address_index::AddressIndex;
/// A Wallet error
#[deriving(Clone, PartialEq, Eq, Show)]
@ -37,7 +37,9 @@ pub enum Error {
/// Tried to add an account when one already exists with that name
DuplicateAccount,
/// An error occured in a BIP32 derivation
Bip32Error(bip32::Error)
Bip32Error(bip32::Error),
/// Tried to use a wallet without an address index
NoAddressIndex
}
/// Each account has two chains, as specified in BIP32
@ -80,7 +82,8 @@ impl Default for Account {
#[deriving(Clone, PartialEq, Eq, Show)]
pub struct Wallet {
master: ExtendedPrivKey,
accounts: HashMap<String, Account>
accounts: HashMap<String, Account>,
index: Option<AddressIndex>
}
impl<S: Encoder<E>, E> Encodable<S, E> for Wallet {
@ -112,7 +115,8 @@ impl<D: Decoder<E>, E> Decodable<D, E> for Wallet {
}
Ok(ret)
})
}))
})),
index: None
})
})
}
@ -127,10 +131,30 @@ impl Wallet {
Ok(Wallet {
master: try!(ExtendedPrivKey::new_master(network, seed)),
accounts: accounts
accounts: accounts,
index: None
})
}
/// Creates the address index
#[inline]
pub fn build_index(&mut self, utxo_set: &UtxoSet) {
let new = AddressIndex::new(utxo_set, self);
self.index = Some(new);
}
/// Accessor for the wallet's address index
#[inline]
pub fn index<'a>(&'a self) -> Option<&'a AddressIndex> {
self.index.as_ref()
}
/// Mutable accessor for the wallet's address index
#[inline]
pub fn index_mut<'a>(&'a mut self) -> Option<&'a mut AddressIndex> {
self.index.as_mut()
}
/// Adds an account to a wallet
pub fn account_insert(&mut self, name: String)
-> Result<(), Error> {
@ -152,6 +176,7 @@ impl Wallet {
}
/// Locates an account in a wallet
#[inline]
pub fn account_find<'a>(&'a self, name: &str)
-> Option<&'a Account> {
self.accounts.find_equiv(&name)
@ -162,10 +187,10 @@ impl Wallet {
account: &str,
chain: AccountChain)
-> Result<Address, Error> {
let (k1, k2) = self.siphash_key();
// TODO: unnecessary allocation, waiting on *_equiv in stdlib
let account = self.accounts.find_mut(&account.to_string());
let account = match account { Some(a) => a, None => return Err(AccountNotFound) };
let index = match self.index { Some(ref i) => i, None => return Err(NoAddressIndex) };
let (mut i, master) = match chain {
Internal => (account.internal_next,
@ -183,7 +208,7 @@ impl Wallet {
let mut address = Address::from_key(
master.network,
&PublicKey::from_secret_key(&sk.secret_key, true));
while !admissible_address(k1, k2, &address) {
while !index.admissible_address(&address) {
i += 1;
sk = try!(master.ckd_priv(Normal(i)).map_err(Bip32Error));
address = Address::from_key(
@ -219,22 +244,53 @@ impl Wallet {
u64_from_be_bytes(ck_slice, 8, 8))
}
/// A filter used for creating a small address index
#[inline]
pub fn might_be_mine(&self, out: &TxOut) -> bool {
let (k1, k2) = self.siphash_key();
match out.classify(self.network()) {
PayToPubkeyHash(addr) => admissible_address(k1, k2, &addr),
_ => false
/// Total balance
pub fn total_balance(&self) -> Result<u64, Error> {
let mut ret = 0;
for (_, account) in self.accounts.iter() {
ret += try!(self.account_balance(account));
}
Ok(ret)
}
/// Account balance
pub fn balance(&self, account: &str) -> Result<u64, Error> {
let account = self.accounts.find_equiv(&account);
let account = match account { Some(a) => a, None => return Err(AccountNotFound) };
self.account_balance(account)
}
fn account_balance(&self, account: &Account) -> Result<u64, Error> {
let index = match self.index { Some(ref i) => i, None => return Err(NoAddressIndex) };
let mut ret = 0;
// Sum internal balance
let master = try!(ExtendedPrivKey::from_path(
&self.master,
account.internal_path.as_slice()).map_err(Bip32Error));
for &cnum in account.internal_used.iter() {
let sk = try!(master.ckd_priv(cnum).map_err(Bip32Error));
let pk = ExtendedPubKey::from_private(&sk);
let addr = Address::from_key(pk.network, &pk.public_key);
for &(_, _, ref out) in index.find_by_script(&addr.script_pubkey()).iter() {
ret += out.value;
}
}
// Sum external balance
let master = try!(ExtendedPrivKey::from_path(
&self.master,
account.external_path.as_slice()).map_err(Bip32Error));
for &cnum in account.external_used.iter() {
let sk = try!(master.ckd_priv(cnum).map_err(Bip32Error));
let pk = ExtendedPubKey::from_private(&sk);
let addr = Address::from_key(pk.network, &pk.public_key);
for &(_, _, ref out) in index.find_by_script(&addr.script_pubkey()).iter() {
ret += out.value;
}
}
/// A filter used for creating a small address index. Note that this
/// function, `might_be_mine` and `siphash_key` are used by wizards-wallet
/// to create a cheap UTXO index
#[inline]
pub fn admissible_address(k1: u64, k2: u64, addr: &Address) -> bool {
hash_with_keys(k1, k2, &addr.as_slice()) & 0xFF == 0
Ok(ret)
}
}