wallet/recovery.go

   1 package wallet
   2
   3 import (
   4         "encoding/json"
   5         "fmt"
   6         "reflect"
   7         "sync"
   8         "sync/atomic"
   9         "time"
  10
  11         "github.com/bytom/account"
  12         "github.com/bytom/blockchain/signers"
  13         "github.com/bytom/crypto/ed25519/chainkd"
  14         "github.com/bytom/crypto/sha3pool"
  15         "github.com/bytom/errors"
  16         "github.com/bytom/protocol/bc"
  17         "github.com/bytom/protocol/bc/types"
  18         dbm "github.com/bytom/database/leveldb"
  19 )
  20
  21 const (
  22         // acctRecoveryWindow defines the account derivation lookahead used when
  23         // attempting to recover the set of used accounts.
  24         acctRecoveryWindow = uint64(6)
  25
  26         // addrRecoveryWindow defines the address derivation lookahead used when
  27         // attempting to recover the set of used addresses.
  28         addrRecoveryWindow = uint64(128)
  29 )
  30
  31 //recoveryKey key for db store recovery info.
  32 var (
  33         recoveryKey = []byte("RecoveryInfo")
  34
  35         // ErrRecoveryBusy another recovery in progress, can not get recovery manager lock
  36         ErrRecoveryBusy = errors.New("another recovery in progress")
  37
  38         // ErrInvalidAcctID can not find account by account id
  39         ErrInvalidAcctID = errors.New("invalid account id")
  40 )
  41
  42 // branchRecoveryState maintains the required state in-order to properly
  43 // recover addresses derived from a particular account's internal or external
  44 // derivation branch.
  45 //
  46 // A branch recovery state supports operations for:
  47 //  - Expanding the look-ahead horizon based on which indexes have been found.
  48 //  - Registering derived addresses with indexes within the horizon.
  49 //  - Reporting an invalid child index that falls into the horizon.
  50 //  - Reporting that an address has been found.
  51 //  - Retrieving all currently derived addresses for the branch.
  52 //  - Looking up a particular address by its child index.
  53 type branchRecoveryState struct {
  54         // recoveryWindow defines the key-derivation lookahead used when
  55         // attempting to recover the set of addresses on this branch.
  56         RecoveryWindow uint64
  57
  58         // horizion records the highest child index watched by this branch.
  59         Horizon uint64
  60
  61         // nextUnfound maintains the child index of the successor to the highest
  62         // index that has been found during recovery of this branch.
  63         NextUnfound uint64
  64 }
  65
  66 // newBranchRecoveryState creates a new branchRecoveryState that can be used to
  67 // track either the external or internal branch of an account's derivation path.
  68 func newBranchRecoveryState(recoveryWindow uint64) *branchRecoveryState {
  69         return &branchRecoveryState{
  70                 RecoveryWindow: recoveryWindow,
  71                 Horizon:        1,
  72                 NextUnfound:    1,
  73         }
  74 }
  75
  76 // extendHorizon returns the current horizon and the number of addresses that
  77 // must be derived in order to maintain the desired recovery window.
  78 func (brs *branchRecoveryState) extendHorizon() (uint64, uint64) {
  79         // Compute the new horizon, which should surpass our last found address
  80         // by the recovery window.
  81         curHorizon := brs.Horizon
  82
  83         minValidHorizon := brs.NextUnfound + brs.RecoveryWindow
  84
  85         // If the current horizon is sufficient, we will not have to derive any
  86         // new keys.
  87         if curHorizon >= minValidHorizon {
  88                 return curHorizon, 0
  89         }
  90
  91         // Otherwise, the number of addresses we should derive corresponds to
  92         // the delta of the two horizons, and we update our new horizon.
  93         delta := minValidHorizon - curHorizon
  94         brs.Horizon = minValidHorizon
  95
  96         return curHorizon, delta
  97 }
  98
  99 // reportFound updates the last found index if the reported index exceeds the
 100 // current value.
 101 func (brs *branchRecoveryState) reportFound(index uint64) {
 102         if index >= brs.NextUnfound {
 103                 brs.NextUnfound = index + 1
 104         }
 105 }
 106
 107 // addressRecoveryState is used to manage the recovery of addresses generated
 108 // under a particular BIP32/BIP44 account. Each account tracks both an external and
 109 // internal branch recovery state, both of which use the same recovery window.
 110 type addressRecoveryState struct {
 111         // ExternalBranch is the recovery state of addresses generated for
 112         // external use, i.e. receiving addresses.
 113         ExternalBranch *branchRecoveryState
 114
 115         // InternalBranch is the recovery state of addresses generated for
 116         // internal use, i.e. change addresses.
 117         InternalBranch *branchRecoveryState
 118
 119         Account *account.Account
 120 }
 121
 122 func newAddressRecoveryState(recoveryWindow uint64, account *account.Account) *addressRecoveryState {
 123         return &addressRecoveryState{
 124                 ExternalBranch: newBranchRecoveryState(recoveryWindow),
 125                 InternalBranch: newBranchRecoveryState(recoveryWindow),
 126                 Account:        account,
 127         }
 128 }
 129
 130 // recoveryState used to record the status of a recovery process.
 131 type recoveryState struct {
 132         // XPubs recovery account xPubs
 133         XPubs []chainkd.XPub
 134
 135         // The time to start the recovery task, used to detemine whether
 136         // recovery task is completed.
 137         StartTime time.Time
 138
 139         // XPubsStatus maintains a map of each requested XPub to its active
 140         // account recovery state.
 141         XPubsStatus *branchRecoveryState
 142
 143         // AcctStatus maintains a map of each requested key scope to its active
 144         // recovery state.
 145         AccountsStatus map[string]*addressRecoveryState
 146 }
 147
 148 func newRecoveryState() *recoveryState {
 149         return &recoveryState{
 150                 AccountsStatus: make(map[string]*addressRecoveryState),
 151                 StartTime:      time.Now(),
 152         }
 153 }
 154
 155 // stateForScope returns a ScopeRecoveryState for the provided key scope. If one
 156 // does not already exist, a new one will be generated with the RecoveryState's
 157 // recoveryWindow.
 158 func (rs *recoveryState) stateForScope(account *account.Account) {
 159         // If the account recovery state already exists, return it.
 160         if _, ok := rs.AccountsStatus[account.ID]; ok {
 161                 return
 162         }
 163
 164         // Otherwise, initialize the recovery state for this scope with the
 165         // chosen recovery window.
 166         rs.AccountsStatus[account.ID] = newAddressRecoveryState(addrRecoveryWindow, account)
 167 }
 168
 169 // recoveryManager manage recovery wallet from key.
 170 type recoveryManager struct {
 171         mu sync.Mutex
 172
 173         db         dbm.DB
 174         accountMgr *account.Manager
 175
 176         locked int32
 177
 178         started bool
 179
 180         // state encapsulates and allocates the necessary recovery state for all
 181         // key scopes and subsidiary derivation paths.
 182         state *recoveryState
 183
 184         //addresses all addresses derivation lookahead used when
 185         // attempting to recover the set of used addresses.
 186         addresses map[bc.Hash]*account.CtrlProgram
 187 }
 188
 189 // newRecoveryManager create recovery manger.
 190 func newRecoveryManager(db dbm.DB, accountMgr *account.Manager) *recoveryManager {
 191         return &recoveryManager{
 192                 db:         db,
 193                 accountMgr: accountMgr,
 194                 addresses:  make(map[bc.Hash]*account.CtrlProgram),
 195                 state:      newRecoveryState(),
 196         }
 197 }
 198
 199 func (m *recoveryManager) AddrResurrect(accts []*account.Account) error {
 200         m.mu.Lock()
 201         defer m.mu.Unlock()
 202
 203         for _, acct := range accts {
 204                 m.state.stateForScope(acct)
 205                 if err := m.extendScanAddresses(acct.ID, false); err != nil {
 206                         return err
 207                 }
 208
 209                 //Bip32 path no change field, no need to create addresses repeatedly.
 210                 if acct.DeriveRule == signers.BIP0032 {
 211                         continue
 212                 }
 213                 if err := m.extendScanAddresses(acct.ID, true); err != nil {
 214                         return err
 215                 }
 216         }
 217
 218         m.state.StartTime = time.Now()
 219         m.started = true
 220         return nil
 221 }
 222
 223 func (m *recoveryManager) AcctResurrect(xPubs []chainkd.XPub) error {
 224         m.mu.Lock()
 225         defer m.mu.Unlock()
 226
 227         if !m.tryStartXPubsRec() {
 228                 return ErrRecoveryBusy
 229         }
 230
 231         m.state.XPubs = xPubs
 232         m.state.XPubsStatus = newBranchRecoveryState(acctRecoveryWindow)
 233
 234         if err := m.extendScanAccounts(); err != nil {
 235                 m.stopXPubsRec()
 236                 return err
 237         }
 238         m.state.StartTime = time.Now()
 239         m.started = true
 240         return nil
 241 }
 242
 243 func (m *recoveryManager) commitStatusInfo() error {
 244         rawStatus, err := json.Marshal(m.state)
 245         if err != nil {
 246                 return err
 247         }
 248
 249         m.db.Set(recoveryKey, rawStatus)
 250         return nil
 251 }
 252
 253 func genAcctAlias(xPubs []chainkd.XPub, index uint64) string {
 254         var tmp []byte
 255         for _, xPub := range xPubs {
 256                 tmp = append(tmp, xPub[:6]...)
 257         }
 258         return fmt.Sprintf("%x:%x", tmp, index)
 259 }
 260
 261 func (m *recoveryManager) extendScanAccounts() error {
 262         if m.state.XPubsStatus == nil {
 263                 return nil
 264         }
 265
 266         curHorizon, delta := m.state.XPubsStatus.extendHorizon()
 267         for index := curHorizon; index < curHorizon+delta; index++ {
 268                 alias := genAcctAlias(m.state.XPubs, index)
 269                 account, err := account.CreateAccount(m.state.XPubs, len(m.state.XPubs), alias, index, signers.BIP0044)
 270                 if err != nil {
 271                         return err
 272                 }
 273
 274                 m.state.stateForScope(account)
 275                 //generate resurrect address for new account.
 276                 if err := m.extendScanAddresses(account.ID, true); err != nil {
 277                         return err
 278                 }
 279
 280                 if err := m.extendScanAddresses(account.ID, false); err != nil {
 281                         return err
 282                 }
 283         }
 284
 285         return nil
 286 }
 287
 288 func getCPHash(cp []byte) bc.Hash {
 289         var h [32]byte
 290         sha3pool.Sum256(h[:], cp)
 291         return bc.NewHash(h)
 292 }
 293
 294 func (m *recoveryManager) extendAddress(acct *account.Account, index uint64, change bool) error {
 295         cp, err := account.CreateCtrlProgram(acct, index, change)
 296         if err != nil {
 297                 return err
 298         }
 299
 300         m.addresses[getCPHash(cp.ControlProgram)] = cp
 301         return nil
 302 }
 303
 304 func (m *recoveryManager) extendScanAddresses(accountID string, change bool) error {
 305         state, ok := m.state.AccountsStatus[accountID]
 306         if !ok {
 307                 return ErrInvalidAcctID
 308         }
 309
 310         var curHorizon, delta uint64
 311         if change {
 312                 curHorizon, delta = state.InternalBranch.extendHorizon()
 313         } else {
 314                 curHorizon, delta = state.ExternalBranch.extendHorizon()
 315         }
 316         for index := curHorizon; index < curHorizon+delta; index++ {
 317                 if err := m.extendAddress(state.Account, index, change); err != nil {
 318                         return err
 319                 }
 320         }
 321         return nil
 322 }
 323
 324 func (m *recoveryManager) processBlock(b *types.Block) error {
 325         for _, tx := range b.Transactions {
 326                 for _, output := range tx.Outputs {
 327                         if cp, ok := m.addresses[getCPHash(output.ControlProgram)]; ok {
 328                                 status, ok := m.state.AccountsStatus[cp.AccountID]
 329                                 if !ok {
 330                                         return ErrInvalidAcctID
 331                                 }
 332
 333                                 if err := m.reportFound(status.Account, cp); err != nil {
 334                                         return err
 335                                 }
 336                         }
 337                 }
 338         }
 339
 340         return nil
 341 }
 342
 343 // FilterRecoveryTxs Filter transactions that meet the recovery address
 344 func (m *recoveryManager) FilterRecoveryTxs(b *types.Block) error {
 345         m.mu.Lock()
 346         defer m.mu.Unlock()
 347
 348         if !m.started {
 349                 return nil
 350         }
 351         if b.Time().After(m.state.StartTime) {
 352                 m.finished()
 353                 return nil
 354         }
 355         return m.processBlock(b)
 356 }
 357
 358 func (m *recoveryManager) finished() {
 359         m.db.Delete(recoveryKey)
 360         m.started = false
 361         m.addresses = make(map[bc.Hash]*account.CtrlProgram)
 362         m.state = newRecoveryState()
 363         m.stopXPubsRec()
 364 }
 365
 366 func (m *recoveryManager) LoadStatusInfo() error {
 367         m.mu.Lock()
 368         defer m.mu.Unlock()
 369
 370         rawStatus := m.db.Get(recoveryKey)
 371         if rawStatus == nil {
 372                 return nil
 373         }
 374
 375         if err := json.Unmarshal(rawStatus, m.state); err != nil {
 376                 return err
 377         }
 378
 379         if m.state.XPubs != nil && !m.tryStartXPubsRec() {
 380                 return ErrRecoveryBusy
 381         }
 382
 383         if err := m.restoreAddresses(); err != nil {
 384                 m.stopXPubsRec()
 385                 return err
 386         }
 387
 388         m.started = true
 389         return nil
 390 }
 391
 392 // restoreAddresses resume addresses for unfinished tasks
 393 func (m *recoveryManager) restoreAddresses() error {
 394         for _, state := range m.state.AccountsStatus {
 395                 for index := uint64(1); index <= state.InternalBranch.Horizon; index++ {
 396                         if err := m.extendAddress(state.Account, index, true); err != nil {
 397                                 return err
 398                         }
 399                 }
 400
 401                 for index := uint64(1); index <= state.ExternalBranch.Horizon; index++ {
 402                         if err := m.extendAddress(state.Account, index, false); err != nil {
 403                                 return err
 404                         }
 405                 }
 406         }
 407         return nil
 408 }
 409
 410 // reportFound found your own address operation.
 411 func (m *recoveryManager) reportFound(account *account.Account, cp *account.CtrlProgram) error {
 412         if m.state.XPubsStatus != nil && reflect.DeepEqual(m.state.XPubs, account.XPubs) {
 413                 //recovery from XPubs need save account to db.
 414                 if err := m.saveAccount(account); err != nil {
 415                         return err
 416                 }
 417
 418                 m.state.XPubsStatus.reportFound(account.KeyIndex)
 419                 if err := m.extendScanAccounts(); err != nil {
 420                         return err
 421                 }
 422         }
 423
 424         if cp.Change {
 425                 m.state.AccountsStatus[account.ID].InternalBranch.reportFound(cp.KeyIndex)
 426         } else {
 427                 m.state.AccountsStatus[account.ID].ExternalBranch.reportFound(cp.KeyIndex)
 428         }
 429
 430         if err := m.extendScanAddresses(account.ID, cp.Change); err != nil {
 431                 return err
 432         }
 433
 434         if err := m.accountMgr.CreateBatchAddresses(account.ID, cp.Change, cp.KeyIndex); err != nil {
 435                 return err
 436         }
 437
 438         return m.commitStatusInfo()
 439 }
 440
 441 func (m *recoveryManager) saveAccount(acct *account.Account) error {
 442         tmp, err := m.accountMgr.FindByID(acct.ID)
 443         if err != nil && errors.Root(err) != account.ErrFindAccount {
 444                 return err
 445         }
 446
 447         if tmp != nil {
 448                 return nil
 449         }
 450         return m.accountMgr.SaveAccount(acct)
 451 }
 452
 453 //tryStartXPubsRec guarantee that only one xPubs recovery is in progress.
 454 func (m *recoveryManager) tryStartXPubsRec() bool {
 455         return atomic.CompareAndSwapInt32(&m.locked, 0, 1)
 456 }
 457
 458 //stopXPubsRec release xPubs recovery lock.
 459 func (m *recoveryManager) stopXPubsRec() {
 460         m.state.XPubs = nil
 461         atomic.StoreInt32(&m.locked, 0)
 462 }