From d4e756bc231266b4b589bf24db50444305c7c8c4 Mon Sep 17 00:00:00 2001 From: Josh Rickmar Date: Tue, 10 Dec 2013 16:15:25 -0500 Subject: [PATCH] Add getaddressbalance websocket extension request. --- account.go | 70 +++++++++++++++++++++++++++++++++++++++++++++++++ accountstore.go | 6 +++++ cmdmgr.go | 44 +++++++++++++++++++++++++++++++ createtx.go | 3 +++ 4 files changed, 123 insertions(+) diff --git a/account.go b/account.go index f387bd3..e21fc20 100644 --- a/account.go +++ b/account.go @@ -31,6 +31,39 @@ import ( "time" ) +// ErrNotFound describes an error where a map lookup failed due to a +// key not being in the map. +var ErrNotFound = errors.New("not found") + +// addressAccountMap holds a map of addresses to names of the +// accounts that hold each address. +var addressAccountMap = struct { + sync.RWMutex + m map[string]string +}{ + m: make(map[string]string), +} + +// MarkAddressForAccount marks an address as belonging to an account. +func MarkAddressForAccount(address, account string) { + addressAccountMap.Lock() + addressAccountMap.m[address] = account + addressAccountMap.Unlock() +} + +// LookupAccountByAddress returns the account name for address. error +// will be set to ErrNotFound if the address has not been marked as +// associated with any account. +func LookupAccountByAddress(address string) (string, error) { + addressAccountMap.RLock() + defer addressAccountMap.RUnlock() + account, ok := addressAccountMap.m[address] + if !ok { + return "", ErrNotFound + } + return account, nil +} + // Account is a structure containing all the components for a // complete wallet. It contains the Armory-style wallet (to store // addresses and keys), and tx and utxo data stores, along with locks @@ -118,6 +151,37 @@ func (a *Account) CalculateBalance(confirms int) float64 { return float64(bal) / float64(btcutil.SatoshiPerBitcoin) } +// CalculateAddressBalance sums the amounts of all unspent transaction +// outputs to a single address's pubkey hash and returns the balance +// as a float64. +// +// If confirmations is 0, all UTXOs, even those not present in a +// block (height -1), will be used to get the balance. Otherwise, +// a UTXO must be in a block. If confirmations is 1 or greater, +// the balance will be calculated based on how many how many blocks +// include a UTXO. +func (a *Account) CalculateAddressBalance(pubkeyHash []byte, confirms int) float64 { + var bal uint64 // Measured in satoshi + + bs, err := GetCurBlock() + if bs.Height == int32(btcutil.BlockHeightUnknown) || err != nil { + return 0. + } + + a.UtxoStore.RLock() + for _, u := range a.UtxoStore.s { + // Utxos not yet in blocks (height -1) should only be + // added if confirmations is 0. + if confirms == 0 || (u.Height != -1 && int(bs.Height-u.Height+1) >= confirms) { + if bytes.Equal(pubkeyHash, u.AddrHash[:]) { + bal += u.Amt + } + } + } + a.UtxoStore.RUnlock() + return float64(bal) / float64(btcutil.SatoshiPerBitcoin) +} + // ListTransactions returns a slice of maps with details about a recorded // transaction. This is intended to be used for listtransactions RPC // replies. @@ -287,6 +351,9 @@ func (a *Account) ImportWIFPrivateKey(wif string, bs *wallet.BlockStamp) (string log.Errorf("cannot write dirty wallet: %v", err) } + // Associate the imported address with this account. + MarkAddressForAccount(addr, a.Name()) + log.Infof("Imported payment address %v", addr) // Return the payment address string of the imported private key. @@ -463,6 +530,9 @@ func (a *Account) NewAddress() (string, error) { log.Errorf("cannot sync dirty wallet: %v", err) } + // Mark this new address as belonging to this account. + MarkAddressForAccount(addr, a.Name()) + // Request updates from btcd for new transactions sent to this address. a.ReqNewTxsForAddress(addr) diff --git a/accountstore.go b/accountstore.go index 1ceae92..348457c 100644 --- a/accountstore.go +++ b/accountstore.go @@ -418,5 +418,11 @@ func (store *AccountStore) OpenAccount(name string, cfg *config) error { default: log.Warnf("cannot open wallet: %v", err) } + + // Mark all active payment addresses as belonging to this account. + for addr := range a.ActivePaymentAddresses() { + MarkAddressForAccount(addr, name) + } + return nil } diff --git a/cmdmgr.go b/cmdmgr.go index 360826b..1d4d834 100644 --- a/cmdmgr.go +++ b/cmdmgr.go @@ -22,6 +22,7 @@ import ( "errors" "fmt" "github.com/conformal/btcjson" + "github.com/conformal/btcutil" "github.com/conformal/btcwallet/tx" "github.com/conformal/btcwallet/wallet" "github.com/conformal/btcwire" @@ -59,6 +60,7 @@ var rpcHandlers = map[string]cmdHandler{ // Extensions exclusive to websocket connections. var wsHandlers = map[string]cmdHandler{ + "getaddressbalance": GetAddressBalance, "getbalances": GetBalances, "listalltransactions": ListAllTransactions, "walletislocked": WalletIsLocked, @@ -283,6 +285,48 @@ func GetBalances(frontend chan []byte, cmd btcjson.Cmd) { NotifyBalances(frontend) } +// GetAddressBalance replies to a getaddressbalance extension request +// by replying with the current balance (sum of unspent transaction +// output amounts) for a single address. +func GetAddressBalance(frontend chan []byte, icmd btcjson.Cmd) { + // Type assert icmd to access parameters. + cmd, ok := icmd.(*btcws.GetAddressBalanceCmd) + if !ok { + ReplyError(frontend, icmd.Id(), &btcjson.ErrInternal) + return + } + + // Is address valid? + pkhash, net, err := btcutil.DecodeAddress(cmd.Address) + if err != nil || net != cfg.Net() { + ReplyError(frontend, cmd.Id(), &btcjson.ErrInvalidAddressOrKey) + return + } + + // Look up account which holds this address. + aname, err := LookupAccountByAddress(cmd.Address) + if err == ErrNotFound { + e := &btcjson.Error{ + Code: btcjson.ErrInvalidAddressOrKey.Code, + Message: "Address not found in wallet", + } + ReplyError(frontend, cmd.Id(), e) + return + } + + // Get the account which holds the address in the request. + // This should not fail, so if it does, return an internal + // error to the frontend. + a, err := accountstore.Account(aname) + if err != nil { + ReplyError(frontend, cmd.Id(), &btcjson.ErrInternal) + return + } + + bal := a.CalculateAddressBalance(pkhash, int(cmd.Minconf)) + ReplySuccess(frontend, cmd.Id(), bal) +} + // ImportPrivKey replies to an importprivkey request by parsing // a WIF-encoded private key and adding it to an account. func ImportPrivKey(frontend chan []byte, icmd btcjson.Cmd) { diff --git a/createtx.go b/createtx.go index 621e079..0efc849 100644 --- a/createtx.go +++ b/createtx.go @@ -250,6 +250,9 @@ func (a *Account) txToPairs(pairs map[string]int64, minconf int) (*CreatedTx, er return nil, fmt.Errorf("failed to get next address: %s", err) } + // Mark change address as belonging to this account. + MarkAddressForAccount(changeAddr, a.Name()) + changeAddrHash, _, err = btcutil.DecodeAddress(changeAddr) if err != nil { return nil, fmt.Errorf("cannot decode new address: %s", err)