wallet: update HD recovery logic to map outpoints to addresses

This commit is contained in:
Olaoluwa Osuntokun 2018-06-14 22:05:39 -07:00
parent b29e917a24
commit d6ee9fd0e4

View file

@ -3,7 +3,9 @@ package wallet
import ( import (
"time" "time"
"github.com/btcsuite/btcd/chaincfg"
"github.com/btcsuite/btcd/chaincfg/chainhash" "github.com/btcsuite/btcd/chaincfg/chainhash"
"github.com/btcsuite/btcd/txscript"
"github.com/btcsuite/btcd/wire" "github.com/btcsuite/btcd/wire"
"github.com/btcsuite/btcutil" "github.com/btcsuite/btcutil"
"github.com/btcsuite/btcutil/hdkeychain" "github.com/btcsuite/btcutil/hdkeychain"
@ -29,15 +31,22 @@ type RecoveryManager struct {
// state encapsulates and allocates the necessary recovery state for all // state encapsulates and allocates the necessary recovery state for all
// key scopes and subsidiary derivation paths. // key scopes and subsidiary derivation paths.
state *RecoveryState state *RecoveryState
// chainParams are the parameters that describe the chain we're trying
// to recover funds on.
chainParams *chaincfg.Params
} }
// NewRecoveryManager initializes a new RecoveryManager with a derivation // NewRecoveryManager initializes a new RecoveryManager with a derivation
// look-ahead of `recoveryWindow` child indexes, and pre-allocates a backing // look-ahead of `recoveryWindow` child indexes, and pre-allocates a backing
// array for `batchSize` blocks to scan at once. // array for `batchSize` blocks to scan at once.
func NewRecoveryManager(recoveryWindow, batchSize uint32) *RecoveryManager { func NewRecoveryManager(recoveryWindow, batchSize uint32,
chainParams *chaincfg.Params) *RecoveryManager {
return &RecoveryManager{ return &RecoveryManager{
recoveryWindow: recoveryWindow, recoveryWindow: recoveryWindow,
blockBatch: make([]wtxmgr.BlockMeta, 0, batchSize), blockBatch: make([]wtxmgr.BlockMeta, 0, batchSize),
chainParams: chainParams,
state: NewRecoveryState(recoveryWindow), state: NewRecoveryState(recoveryWindow),
} }
} }
@ -121,7 +130,14 @@ func (rm *RecoveryManager) Resurrect(ns walletdb.ReadBucket,
// to our global set of watched outpoints, so that we can watch them for // to our global set of watched outpoints, so that we can watch them for
// spends. // spends.
for _, credit := range credits { for _, credit := range credits {
rm.state.AddWatchedOutPoint(&credit.OutPoint) _, addrs, _, err := txscript.ExtractPkScriptAddrs(
credit.PkScript, rm.chainParams,
)
if err != nil {
return err
}
rm.state.AddWatchedOutPoint(&credit.OutPoint, addrs[0])
} }
return nil return nil
@ -191,7 +207,7 @@ type RecoveryState struct {
// watchedOutPoints contains the set of all outpoints known to the // watchedOutPoints contains the set of all outpoints known to the
// wallet. This is updated iteratively as new outpoints are found during // wallet. This is updated iteratively as new outpoints are found during
// a rescan. // a rescan.
watchedOutPoints map[wire.OutPoint]struct{} watchedOutPoints map[wire.OutPoint]btcutil.Address
} }
// NewRecoveryState creates a new RecoveryState using the provided // NewRecoveryState creates a new RecoveryState using the provided
@ -203,7 +219,7 @@ func NewRecoveryState(recoveryWindow uint32) *RecoveryState {
return &RecoveryState{ return &RecoveryState{
recoveryWindow: recoveryWindow, recoveryWindow: recoveryWindow,
scopes: scopes, scopes: scopes,
watchedOutPoints: make(map[wire.OutPoint]struct{}), watchedOutPoints: make(map[wire.OutPoint]btcutil.Address),
} }
} }
@ -227,14 +243,16 @@ func (rs *RecoveryState) StateForScope(
// WatchedOutPoints returns the global set of outpoints that are known to belong // WatchedOutPoints returns the global set of outpoints that are known to belong
// to the wallet during recovery. // to the wallet during recovery.
func (rs *RecoveryState) WatchedOutPoints() map[wire.OutPoint]struct{} { func (rs *RecoveryState) WatchedOutPoints() map[wire.OutPoint]btcutil.Address {
return rs.watchedOutPoints return rs.watchedOutPoints
} }
// AddWatchedOutPoint updates the recovery state's set of known outpoints that // AddWatchedOutPoint updates the recovery state's set of known outpoints that
// we will monitor for spends during recovery. // we will monitor for spends during recovery.
func (rs *RecoveryState) AddWatchedOutPoint(outPoint *wire.OutPoint) { func (rs *RecoveryState) AddWatchedOutPoint(outPoint *wire.OutPoint,
rs.watchedOutPoints[*outPoint] = struct{}{} addr btcutil.Address) {
rs.watchedOutPoints[*outPoint] = addr
} }
// ScopeRecoveryState is used to manage the recovery of addresses generated // ScopeRecoveryState is used to manage the recovery of addresses generated