Add getaddressbalance websocket extension request.

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)
 

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
 }

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) {

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)