1 package wallet
2 3 import (
4 "github.com/p9c/p9/pkg/btcaddr"
5 "github.com/p9c/p9/pkg/chaincfg"
6 "time"
7 8 "github.com/p9c/p9/pkg/chainhash"
9 "github.com/p9c/p9/pkg/txscript"
10 "github.com/p9c/p9/pkg/util/hdkeychain"
11 "github.com/p9c/p9/pkg/waddrmgr"
12 "github.com/p9c/p9/pkg/walletdb"
13 "github.com/p9c/p9/pkg/wire"
14 "github.com/p9c/p9/pkg/wtxmgr"
15 )
16 17 // RecoveryManager maintains the state required to recover previously used addresses, and coordinates batched processing
18 // of the blocks to search.
19 type RecoveryManager struct {
20 // recoveryWindow defines the key-derivation lookahead used when attempting to recover the set of used addresses.
21 recoveryWindow uint32
22 // started is true after the first block has been added to the batch.
23 started bool
24 // blockBatch contains a list of blocks that have not yet been searched for recovered addresses.
25 blockBatch []wtxmgr.BlockMeta
26 // state encapsulates and allocates the necessary recovery state for all key scopes and subsidiary derivation paths.
27 state *RecoveryState
28 // chainParams are the parameters that describe the chain we're trying to recover funds on.
29 chainParams *chaincfg.Params
30 }
31 32 // NewRecoveryManager initializes a new RecoveryManager with a derivation look-ahead of `recoveryWindow` child indexes,
33 // and pre-allocates a backing array for `batchSize` blocks to scan at once.
34 func NewRecoveryManager(
35 recoveryWindow, batchSize uint32,
36 chainParams *chaincfg.Params,
37 ) *RecoveryManager {
38 return &RecoveryManager{
39 recoveryWindow: recoveryWindow,
40 blockBatch: make([]wtxmgr.BlockMeta, 0, batchSize),
41 chainParams: chainParams,
42 state: NewRecoveryState(recoveryWindow),
43 }
44 }
45 46 // Resurrect restores all known addresses for the provided scopes that can be found in the walletdb namespace, in
47 // addition to restoring all outpoints that have been previously found. This method ensures that the recovery state's
48 // horizons properly start from the last found address of a prior recovery attempt.
49 func (rm *RecoveryManager) Resurrect(
50 ns walletdb.ReadBucket,
51 scopedMgrs map[waddrmgr.KeyScope]*waddrmgr.ScopedKeyManager,
52 credits []wtxmgr.Credit,
53 ) (e error) {
54 // First, for each scope that we are recovering, rederive all of the addresses up to the last found address known to
55 // each branch.
56 for keyScope, scopedMgr := range scopedMgrs {
57 // Load the current account properties for this scope, using the the default account number.
58 //
59 // TODO(conner): rescan for all created accounts if we allow users to use non-default address
60 scopeState := rm.state.StateForScope(keyScope)
61 var acctProperties *waddrmgr.AccountProperties
62 acctProperties, e = scopedMgr.AccountProperties(
63 ns, waddrmgr.DefaultAccountNum,
64 )
65 if e != nil {
66 return e
67 }
68 // Fetch the external key count, which bounds the indexes we will need to rederive.
69 externalCount := acctProperties.ExternalKeyCount
70 // Walk through all indexes through the last external key, deriving each address and adding it to the external
71 // branch recovery state's set of addresses to look for.
72 for i := uint32(0); i < externalCount; i++ {
73 keyPath := externalKeyPath(i)
74 var addr waddrmgr.ManagedAddress
75 addr, e = scopedMgr.DeriveFromKeyPath(ns, keyPath)
76 if e != nil && e != hdkeychain.ErrInvalidChild || addr == nil {
77 return e
78 } else if e == hdkeychain.ErrInvalidChild {
79 scopeState.ExternalBranch.MarkInvalidChild(i)
80 continue
81 }
82 scopeState.ExternalBranch.AddAddr(i, addr.Address())
83 }
84 // Fetch the internal key count, which bounds the indexes we will need to rederive.
85 internalCount := acctProperties.InternalKeyCount
86 // Walk through all indexes through the last internal key, deriving each address and adding it to the internal
87 // branch recovery state's set of addresses to look for.
88 for i := uint32(0); i < internalCount; i++ {
89 keyPath := internalKeyPath(i)
90 var addr waddrmgr.ManagedAddress
91 addr, e = scopedMgr.DeriveFromKeyPath(ns, keyPath)
92 if e != nil && e != hdkeychain.ErrInvalidChild || addr == nil {
93 return e
94 } else if e == hdkeychain.ErrInvalidChild {
95 scopeState.InternalBranch.MarkInvalidChild(i)
96 continue
97 }
98 scopeState.InternalBranch.AddAddr(i, addr.Address())
99 }
100 // The key counts will point to the next key that can be derived, so we subtract one to point to last known key.
101 // If the key count is zero, then no addresses have been found.
102 if externalCount > 0 {
103 scopeState.ExternalBranch.ReportFound(externalCount - 1)
104 }
105 if internalCount > 0 {
106 scopeState.InternalBranch.ReportFound(internalCount - 1)
107 }
108 }
109 // In addition, we will re-add any outpoints that are known the wallet to our global set of watched outpoints, so
110 // that we can watch them for spends.
111 for _, credit := range credits {
112 var addrs []btcaddr.Address
113 _, addrs, _, e = txscript.ExtractPkScriptAddrs(
114 credit.PkScript, rm.chainParams,
115 )
116 if e != nil {
117 return e
118 }
119 rm.state.AddWatchedOutPoint(&credit.OutPoint, addrs[0])
120 }
121 return nil
122 }
123 124 // AddToBlockBatch appends the block information, consisting of hash and height, to the batch of blocks to be searched.
125 func (rm *RecoveryManager) AddToBlockBatch(
126 hash *chainhash.Hash, height int32,
127 timestamp time.Time,
128 ) {
129 if !rm.started {
130 T.F(
131 "Seed birthday surpassed, starting recovery of wallet from height=%d hash=%v with recovery-window=%d",
132 height, *hash, rm.recoveryWindow,
133 )
134 rm.started = true
135 }
136 block := wtxmgr.BlockMeta{
137 Block: wtxmgr.Block{
138 Hash: *hash,
139 Height: height,
140 },
141 Time: timestamp,
142 }
143 rm.blockBatch = append(rm.blockBatch, block)
144 }
145 146 // BlockBatch returns a buffer of blocks that have not yet been searched.
147 func (rm *RecoveryManager) BlockBatch() []wtxmgr.BlockMeta {
148 return rm.blockBatch
149 }
150 151 // ResetBlockBatch resets the internal block buffer to conserve memory.
152 func (rm *RecoveryManager) ResetBlockBatch() {
153 rm.blockBatch = rm.blockBatch[:0]
154 }
155 156 // State returns the current RecoveryState.
157 func (rm *RecoveryManager) State() *RecoveryState {
158 return rm.state
159 }
160 161 // RecoveryState manages the initialization and lookup of ScopeRecoveryStates for any actively used key scopes.
162 //
163 // In order to ensure that all addresses are properly recovered, the window should be sized as the sum of maximum
164 // possible inter-block and intra-block gap between used addresses of a particular branch.
165 //
166 // These are defined as:
167 //
168 // - Inter-Block Gap: The maximum difference between the derived child indexes of the last addresses used in any block
169 // and the next address consumed by a later block.
170 //
171 // - Intra-Block Gap: The maximum difference between the derived child indexes of the first address used in any block
172 // and the last address used in the same block.
173 type RecoveryState struct {
174 // recoveryWindow defines the key-derivation lookahead used when attempting to recover the set of used addresses.
175 // This value will be used to instantiate a new RecoveryState for each requested scope.
176 recoveryWindow uint32
177 // scopes maintains a map of each requested key scope to its active RecoveryState.
178 scopes map[waddrmgr.KeyScope]*ScopeRecoveryState
179 // watchedOutPoints contains the set of all outpoints known to the wallet. This is updated iteratively as new
180 // outpoints are found during a rescan.
181 watchedOutPoints map[wire.OutPoint]btcaddr.Address
182 }
183 184 // NewRecoveryState creates a new RecoveryState using the provided recoveryWindow. Each RecoveryState that is
185 // subsequently initialized for a particular key scope will receive the same recoveryWindow.
186 func NewRecoveryState(recoveryWindow uint32) *RecoveryState {
187 scopes := make(map[waddrmgr.KeyScope]*ScopeRecoveryState)
188 return &RecoveryState{
189 recoveryWindow: recoveryWindow,
190 scopes: scopes,
191 watchedOutPoints: make(map[wire.OutPoint]btcaddr.Address),
192 }
193 }
194 195 // StateForScope returns a ScopeRecoveryState for the provided key scope. If one does not already exist, a new one will
196 // be generated with the RecoveryState's recoveryWindow.
197 func (rs *RecoveryState) StateForScope(
198 keyScope waddrmgr.KeyScope,
199 ) *ScopeRecoveryState {
200 // If the account recovery state already exists, return it.
201 if scopeState, ok := rs.scopes[keyScope]; ok {
202 return scopeState
203 }
204 // Otherwise, initialize the recovery state for this scope with the chosen recovery window.
205 rs.scopes[keyScope] = NewScopeRecoveryState(rs.recoveryWindow)
206 return rs.scopes[keyScope]
207 }
208 209 // WatchedOutPoints returns the global set of outpoints that are known to belong to the wallet during recovery.
210 func (rs *RecoveryState) WatchedOutPoints() map[wire.OutPoint]btcaddr.Address {
211 return rs.watchedOutPoints
212 }
213 214 // AddWatchedOutPoint updates the recovery state's set of known outpoints that we will monitor for spends during
215 // recovery.
216 func (rs *RecoveryState) AddWatchedOutPoint(
217 outPoint *wire.OutPoint,
218 addr btcaddr.Address,
219 ) {
220 rs.watchedOutPoints[*outPoint] = addr
221 }
222 223 // ScopeRecoveryState is used to manage the recovery of addresses generated under a particular BIP32 account. Each
224 // account tracks both an external and internal branch recovery state, both of which use the same recovery window.
225 type ScopeRecoveryState struct {
226 // ExternalBranch is the recovery state of addresses generated for external use, i.e. receiving addresses.
227 ExternalBranch *BranchRecoveryState
228 // InternalBranch is the recovery state of addresses generated for internal use, i.e. change addresses.
229 InternalBranch *BranchRecoveryState
230 }
231 232 // NewScopeRecoveryState initializes an ScopeRecoveryState with the chosen recovery window.
233 func NewScopeRecoveryState(recoveryWindow uint32) *ScopeRecoveryState {
234 return &ScopeRecoveryState{
235 ExternalBranch: NewBranchRecoveryState(recoveryWindow),
236 InternalBranch: NewBranchRecoveryState(recoveryWindow),
237 }
238 }
239 240 // BranchRecoveryState maintains the required state in-order to properly recover addresses derived from a particular
241 // account's internal or external derivation branch.
242 //
243 // A branch recovery state supports operations for:
244 //
245 // - Expanding the look-ahead horizon based on which indexes have been found.
246 //
247 // - Registering derived addresses with indexes within the horizon.
248 //
249 // - Reporting an invalid child index that falls into the horizon.
250 //
251 // - Reporting that an address has been found.
252 //
253 // - Retrieving all currently derived addresses for the branch.
254 //
255 // - Looking up a particular address by its child index.
256 type BranchRecoveryState struct {
257 // recoveryWindow defines the key-derivation lookahead used when attempting to recover the set of addresses on this
258 // branch.
259 recoveryWindow uint32
260 // horizion records the highest child index watched by this branch.
261 horizon uint32
262 // nextUnfound maintains the child index of the successor to the highest index that has been found during recovery
263 // of this branch.
264 nextUnfound uint32
265 // addresses is a map of child index to address for all actively watched addresses belonging to this branch.
266 addresses map[uint32]btcaddr.Address
267 // invalidChildren records the set of child indexes that derive to invalid keys.
268 invalidChildren map[uint32]struct{}
269 }
270 271 // NewBranchRecoveryState creates a new BranchRecoveryState that can be used to track either the external or internal
272 // branch of an account's derivation path.
273 func NewBranchRecoveryState(recoveryWindow uint32) *BranchRecoveryState {
274 return &BranchRecoveryState{
275 recoveryWindow: recoveryWindow,
276 addresses: make(map[uint32]btcaddr.Address),
277 invalidChildren: make(map[uint32]struct{}),
278 }
279 }
280 281 // ExtendHorizon returns the current horizon and the number of addresses that must be derived in order to maintain the
282 // desired recovery window.
283 func (brs *BranchRecoveryState) ExtendHorizon() (uint32, uint32) {
284 // Compute the new horizon, which should surpass our last found address by the recovery window.
285 curHorizon := brs.horizon
286 nInvalid := brs.NumInvalidInHorizon()
287 minValidHorizon := brs.nextUnfound + brs.recoveryWindow + nInvalid
288 // If the current horizon is sufficient, we will not have to derive any new keys.
289 if curHorizon >= minValidHorizon {
290 return curHorizon, 0
291 }
292 // Otherwise, the number of addresses we should derive corresponds to the delta of the two horizons, and we update
293 // our new horizon.
294 delta := minValidHorizon - curHorizon
295 brs.horizon = minValidHorizon
296 return curHorizon, delta
297 }
298 299 // AddAddr adds a freshly derived address from our lookahead into the map of known addresses for this branch.
300 func (brs *BranchRecoveryState) AddAddr(index uint32, addr btcaddr.Address) {
301 brs.addresses[index] = addr
302 }
303 304 // GetAddr returns the address derived from a given child index.
305 func (brs *BranchRecoveryState) GetAddr(index uint32) btcaddr.Address {
306 return brs.addresses[index]
307 }
308 309 // ReportFound updates the last found index if the reported index exceeds the current value.
310 func (brs *BranchRecoveryState) ReportFound(index uint32) {
311 if index >= brs.nextUnfound {
312 brs.nextUnfound = index + 1
313 // Prune all invalid child indexes that fall below our last found index. We don't need to keep these entries any
314 // longer, since they will not affect our required look-ahead.
315 for childIndex := range brs.invalidChildren {
316 if childIndex < index {
317 delete(brs.invalidChildren, childIndex)
318 }
319 }
320 }
321 }
322 323 // MarkInvalidChild records that a particular child index results in deriving an invalid address. In addition, the
324 // branch's horizon is increment, as we expect the caller to perform an additional derivation to replace the invalid
325 // child. This is used to ensure that we are always have the proper lookahead when an invalid child is encountered.
326 func (brs *BranchRecoveryState) MarkInvalidChild(index uint32) {
327 brs.invalidChildren[index] = struct{}{}
328 brs.horizon++
329 }
330 331 // NextUnfound returns the child index of the successor to the highest found child index.
332 func (brs *BranchRecoveryState) NextUnfound() uint32 {
333 return brs.nextUnfound
334 }
335 336 // Addrs returns a map of all currently derived child indexes to the their corresponding addresses.
337 func (brs *BranchRecoveryState) Addrs() map[uint32]btcaddr.Address {
338 return brs.addresses
339 }
340 341 // NumInvalidInHorizon computes the number of invalid child indexes that lie between the last found and current horizon.
342 // This informs how many additional indexes to derive in order to maintain the proper number of valid addresses within
343 // our horizon.
344 func (brs *BranchRecoveryState) NumInvalidInHorizon() uint32 {
345 var nInvalid uint32
346 for childIndex := range brs.invalidChildren {
347 if brs.nextUnfound <= childIndex && childIndex < brs.horizon {
348 nInvalid++
349 }
350 }
351 return nInvalid
352 }
353