addresspool.go

/*
 * Copyright (c) 2013-2015 The btcsuite developers
 * Copyright (c) 2015 The Decred developers
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package wallet

import (
	"fmt"
	"sync"

	"github.com/decred/dcrutil"
	"github.com/decred/dcrwallet/waddrmgr"
)

// addressPoolBuffer is the number of addresses to fetch when the address pool
// runs out of new addresses to use.
const addressPoolBuffer = 20

// addressPool is a cache of addresses to use that are generated by the
// address manager. It is safe than directly calling the address manager
// because doing that will increment the cursor of the extended key even
// if the created transaction errors out in some way.
type addressPool struct {
	// Represent addresses as strings because the address interface
	// doesn't have any good way to make comparisons.
	addresses []string
	cursor    int
	account   uint32
	branch    uint32
	index     uint32
	started   bool
	mutex     sync.Mutex
	wallet    *Wallet
}

// newAddressPool creates a new address pool for the wallet default account.
func newAddressPool() *addressPool {
	return &addressPool{
		started: false,
	}
}

// addressPools is a pair of address pools for the two used branches for a
// single account.
type addressPools struct {
	internal *addressPool
	external *addressPool
}

// newAddressPools creates a pair of address pools as an addressPools struct. It
// also initializes the address pools to the passed indexes.
func newAddressPools(account uint32, intIdx, extIdx uint32,
	w *Wallet) (*addressPools, error) {
	a := &addressPools{
		internal: newAddressPool(),
		external: newAddressPool(),
	}
	err := a.external.initialize(account, waddrmgr.ExternalBranch, extIdx, w)
	if err != nil {
		return nil, err
	}
	err = a.internal.initialize(account, waddrmgr.InternalBranch, intIdx, w)
	if err != nil {
		return nil, err
	}

	return a, nil
}

// initialize initializes an address pool for the passed account and branch
// to the address index given. It will automatically load a buffer of addresses
// from the address manager to use for upcoming calls.
func (a *addressPool) initialize(account uint32, branch uint32, index uint32,
	w *Wallet) error {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	// Do not reinitialize an address pool that was already started.
	// This can happen if the RPC client dies due to a disconnect
	// from the daemon.
	if a.started {
		return nil
	}

	// 0 and 1 refer to the external and internal branches of the wallet.
	// Other branches are so far unused.
	if branch > waddrmgr.InternalBranch {
		return fmt.Errorf("unknown branch %v given when attempting to "+
			"initialize address pool for account %v", branch, account)
	}

	// Access the manager and get the synced to index, then insert all
	// the unused addresses into the address pool.
	lastAddrFunc := w.Manager.LastExternalAddress
	if branch == waddrmgr.InternalBranch {
		lastAddrFunc = w.Manager.LastInternalAddress
	}
	_, mgrIdx, err := lastAddrFunc(account)
	if err != nil {
		return fmt.Errorf("failed to retrieve the last used addr index "+
			"from the address manager for branch %v, acct %v: %s", branch,
			account, err.Error())
	}

	if mgrIdx < index {
		return fmt.Errorf("manager is out of sync with the passed index "+
			"(index %v, mgr index %v)", index, mgrIdx)
	}

	if mgrIdx == index {
		a.addresses = make([]string, 0)
	} else {
		fetchNum := mgrIdx - index + 1
		a.addresses = make([]string, fetchNum)
		for i := uint32(0); i < fetchNum; i++ {
			addr, err := w.Manager.AddressDerivedFromDbAcct(index+i, account,
				branch)
			if err != nil {
				return fmt.Errorf("failed to get the address at index %v "+
					"for account %v, branch %v: %s", index+i, account, branch,
					err.Error())
			}
			a.addresses[i] = addr.EncodeAddress()
		}
	}

	a.wallet = w
	a.account = account
	a.branch = branch
	a.index = index

	log.Debugf("Address pool for account %v initialized to next "+
		"address index %v on branch %v", account, a.index, branch)
	log.Debugf("The address manager buffered address space is %v "+
		"many addresses (manager index: %v) for account %v, branch %v",
		len(a.addresses), mgrIdx, account, branch)

	a.cursor = 0
	a.started = true

	return nil
}

// GetNewAddress must be run as many times as necessary with the address pool
// mutex locked. Each time, it returns a single new address while adding that
// address to the toDelete map. If the address pool runs out of addresses, it
// generates more from the address manager.
//
// This function MUST be called with the address pool mutex held and batch
// finish or rollback must be called after.
func (a *addressPool) getNewAddress() (dcrutil.Address, error) {
	if !a.started {
		return nil, fmt.Errorf("failed to getNewAddress; pool not started")
	}

	chainClient, err := a.wallet.requireChainClient()
	if err != nil {
		return nil, err
	}

	// Replenish the pool if we're at the last address.
	if a.cursor == len(a.addresses)-1 || len(a.addresses) == 0 {
		var nextAddrFunc func(uint32, uint32) ([]waddrmgr.ManagedAddress, error)
		switch a.branch {
		case waddrmgr.InternalBranch:
			nextAddrFunc = a.wallet.Manager.NextInternalAddresses
		case waddrmgr.ExternalBranch:
			nextAddrFunc = a.wallet.Manager.NextExternalAddresses
		default:
			return nil, fmt.Errorf("unknown default account branch %v", a.branch)
		}

		addrs, err :=
			nextAddrFunc(a.account, addressPoolBuffer)
		if err != nil {
			return nil, err
		}

		for _, addr := range addrs {
			a.addresses = append(a.addresses, addr.Address().EncodeAddress())
		}
	}

	// As these are all encoded addresses, we should never throw an error
	// converting back.
	curAddressStr := a.addresses[a.cursor]
	curAddress, err := dcrutil.DecodeAddress(curAddressStr, a.wallet.chainParams)
	if err != nil {
		return nil, fmt.Errorf("unexpected error decoding address %s: %s",
			curAddressStr, err.Error())
	}

	log.Debugf("Get new address for branch %v returned %s (idx %v) from "+
		"the address pool", a.branch, curAddressStr, a.index)

	a.cursor++
	a.index++

	// Add the address to the notifications watcher.
	addrs := make([]dcrutil.Address, 1)
	addrs[0] = curAddress
	if err := chainClient.NotifyReceived(addrs); err != nil {
		return nil, err
	}

	return curAddress, nil
}

// GetNewAddress is the exported function that gets a new address
// from the memory pool and then updates the index and writes it to
// disk. This differs from getNewAddress is that it can *not* be
// rolled back after in the event of failure. It should mainly be
// used in calls that provide a single new address to the user for
// them to use externally.
func (a *addressPool) GetNewAddress() (dcrutil.Address, error) {
	defer func() {
		errNotify := a.wallet.notifyAccountAddrIdxs(a.account)
		if errNotify != nil {
			log.Errorf("Failed to push account update notification "+
				"for account %v", a.account)
		}
	}()
	a.mutex.Lock()
	defer a.mutex.Unlock()

	address, err := a.getNewAddress()
	if err == nil {
		a.BatchFinish()
	} else {
		a.BatchRollback()
	}
	return address, err
}

// notifyAccountAddrIdxs sends out an account notification when the address index
// for some account branch has changed.
func (w *Wallet) notifyAccountAddrIdxs(account uint32) error {
	name, err := w.Manager.AccountName(account)
	if err != nil {
		return err
	}

	idxExt, err := w.AddressPoolIndex(account, waddrmgr.ExternalBranch)
	if err != nil {
		return err
	}
	idxInt, err := w.AddressPoolIndex(account, waddrmgr.InternalBranch)
	if err != nil {
		return err
	}

	ap := &waddrmgr.AccountProperties{
		AccountNumber:    account,
		AccountName:      name,
		ExternalKeyCount: idxExt,
		InternalKeyCount: idxInt,
		ImportedKeyCount: 0, // ???
	}

	w.NtfnServer.notifyAccountProperties(ap)

	return nil
}

// BatchFinish must be run after every successful series of usages of
// GetNewAddress to purge the addresses from the unused map.
func (a *addressPool) BatchFinish() {
	log.Debugf("Closing address batch for pool branch %v, next index %v",
		a.branch, a.index)

	isInternal := a.branch == waddrmgr.InternalBranch
	err := a.wallet.Manager.StoreNextToUseAddress(isInternal, a.account,
		a.index)
	if err != nil {
		log.Errorf("Failed to store next to use address idx for "+
			"pool branch %v, account %v in the manager on batch "+
			"finish: %v", a.branch, a.account, err.Error())
	}

	// We used all the addresses, so we need to pull new addresses
	// on the next call of this function.
	if a.cursor >= len(a.addresses) {
		a.addresses = nil
		a.cursor = 0
		return
	}

	a.addresses = a.addresses[a.cursor:len(a.addresses)]
	a.cursor = 0
}

// BatchRollback must be run after every unsuccessful series of usages
// of getNewAddress to restore the cursor to the original position in
// the slice, thus marking all addresses unused again.
func (a *addressPool) BatchRollback() {
	a.index -= uint32(a.cursor)
	a.cursor = 0
}

// Close writes the next to use index for the address pool to disk, then sets
// the address pool as closed.
func (a *addressPool) Close() error {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	if !a.started {
		return fmt.Errorf("attempted to close uninitialized address pool")
	}

	isInternal := a.branch == waddrmgr.InternalBranch
	err := a.wallet.Manager.StoreNextToUseAddress(isInternal, a.account,
		a.index)
	if err != nil {
		return fmt.Errorf("Failed to store next to use address idx for "+
			"pool branch %v, account %v in the manager on address "+
			"pool close: %v", a.branch, a.account, err.Error())
	}

	a.started = false
	return nil
}

// CloseAddressPools grabs one last new address for both internal and external
// acounts. Then it inserts them into the address manager database, so that
// the address manager can be used upon startup to restore the cursor position
// in the address pool.
func (w *Wallet) CloseAddressPools() {
	for _, addressPools := range w.addrPools {
		if addressPools.internal == nil {
			return
		}
		if addressPools.external == nil {
			return
		}
		if !addressPools.internal.started || !addressPools.external.started {
			return
		}

		err := addressPools.internal.Close()
		if err != nil {
			log.Errorf("failed to close default acct internal addr pool: %v",
				err)
		}
		err = addressPools.external.Close()
		if err != nil {
			log.Errorf("failed to close default acct external addr pool: %v",
				err)
		}

		delete(w.addrPools, addressPools.internal.account)
	}

	return
}

// CheckAddressPoolsInitialized checks to make sure an address pool exists
// that that one can safely access functions and internal memory such as
// mutexes.
func (w *Wallet) CheckAddressPoolsInitialized(account uint32) error {
	if w.addrPools[account] == nil {
		return fmt.Errorf("Address pools for account %v "+
			"are undeclared", account)
	}
	if w.addrPools[account].external == nil {
		return fmt.Errorf("External address pool for "+
			"account %v is undeclared", account)
	}
	if w.addrPools[account].internal == nil {
		return fmt.Errorf("Internal address pool for "+
			"account %v is undeclared", account)
	}
	if !w.addrPools[account].external.started {
		return fmt.Errorf("External address pool for the "+
			"account %v is uninitialized", account)
	}
	if !w.addrPools[account].internal.started {
		return fmt.Errorf("Internal address pool for the "+
			"account %v is uninitialized", account)
	}

	return nil
}

// AddressPoolIndex returns the next to use address index for the passed
// branch of the passed account.
func (w *Wallet) AddressPoolIndex(account uint32, branch uint32) (uint32, error) {
	err := w.CheckAddressPoolsInitialized(account)
	if err != nil {
		log.Tracef("Error on fetch of address pool account %v, branch %v from "+
			"the local map: %s", account, branch, err.Error())

		// The address pools are uninitialized. This is usually due to
		// the wallet being in the process of synchronizing. Instead,
		// try to load the last saved address index from the meta bucket
		// of the database. If that fails, give up and return an error.
		isInternal := branch == waddrmgr.InternalBranch
		idx, err := w.Manager.NextToUseAddrPoolIndex(isInternal, account)
		if err != nil {
			return 0, err
		}

		return idx, nil
	}

	switch branch {
	case waddrmgr.ExternalBranch:
		w.addrPools[account].external.mutex.Lock()
		defer w.addrPools[account].external.mutex.Unlock()
		return w.addrPools[account].external.index, nil
	case waddrmgr.InternalBranch:
		w.addrPools[account].internal.mutex.Lock()
		defer w.addrPools[account].internal.mutex.Unlock()
		return w.addrPools[account].internal.index, nil
	}

	return 0, fmt.Errorf("unknown branch number %v", branch)
}

// SyncAddressPoolIndex synchronizes an account's branch to the given address
// by iteratively calling getNewAddress on the respective address pool.
func (w *Wallet) SyncAddressPoolIndex(account uint32, branch uint32,
	index uint32) error {
	// Sanity checks.
	err := w.CheckAddressPoolsInitialized(account)
	if err != nil {
		return err
	}

	defer func() {
		errNotify := w.notifyAccountAddrIdxs(account)
		if errNotify != nil {
			log.Errorf("Failed to push account update notification "+
				"for account %v", account)
		}
	}()
	var addrPool *addressPool
	switch branch {
	case waddrmgr.ExternalBranch:
		addrPool = w.addrPools[account].external
		addrPool.mutex.Lock()
		defer addrPool.mutex.Unlock()
	case waddrmgr.InternalBranch:
		addrPool = w.addrPools[account].internal
		addrPool.mutex.Lock()
		defer addrPool.mutex.Unlock()
	default:
		return fmt.Errorf("unknown branch number %v", branch)
	}
	if index < addrPool.index {
		return fmt.Errorf("the passed index, %v, is before the "+
			"currently synced to address index %v", index,
			addrPool.index)
	}
	if index == addrPool.index {
		return nil
	}

	// Synchronize our address pool by calling getNewAddress
	// iteratively until the next to use index is synced to
	// where we need it.
	toFetch := index - addrPool.index
	for i := uint32(0); i < toFetch; i++ {
		_, err := addrPool.getNewAddress()
		if err != nil {
			addrPool.BatchRollback()
			return err
		}
	}
	addrPool.BatchFinish()

	return nil
}

// NewAddress checks the address pools and then attempts to return a new
// address for the account and branch requested.
func (w *Wallet) NewAddress(account uint32, branch uint32) (dcrutil.Address,
	error) {
	err := w.CheckAddressPoolsInitialized(account)
	if err != nil {
		return nil, err
	}

	var addrPool *addressPool
	switch branch {
	case waddrmgr.ExternalBranch:
		addrPool = w.addrPools[account].external
	case waddrmgr.InternalBranch:
		addrPool = w.addrPools[account].internal
	default:
		return nil, fmt.Errorf("new address failed; unknown branch number %v",
			branch)
	}

	return addrPool.GetNewAddress()
}

// ReusedAddress returns an address that is reused from the external
// branch of the wallet, to cut down on new address usage for wallets.
// Should be used judiciously.
func (w *Wallet) ReusedAddress() (dcrutil.Address, error) {
	addr, err := w.Manager.AddressDerivedFromDbAcct(0,
		waddrmgr.DefaultAccountNum, waddrmgr.ExternalBranch)
	if err != nil {
		return nil, err
	}

	return addr, err
}