Merge pull request #63 from stroomnetwork/feat/add-manual-rollback

feat: Add manual rollback

Author: diminio1

run/btcwallet.go

@@ -39,7 +39,8 @@ type BtcwalletConfig struct {
 	InitTimeout    time.Duration
 	FeeCoefficient float64
 	// RescanStartBlock is the block height to start rescan from; 0 means from the wallet sync block
-	RescanStartBlock uint64
+	RescanStartBlock       uint64
+	RollbackStateFromBlock uint64
 }
 
 func SafeInitWallet(config *BtcwalletConfig) (*wallet.Wallet, error) {
@@ -184,11 +185,18 @@ func doInit(config *BtcwalletConfig) (*wallet.Wallet, error) {
 		w.FeeCoefficient = config.FeeCoefficient
 
 		if config.RescanStartBlock != 0 {
-			rescanStartStamp, err := w.GetBlockStamp(config.RescanStartBlock)
+			stamp, err := w.GetBlockStamp(config.RescanStartBlock)
 			if err != nil {
 				return nil, fmt.Errorf("cannot get rescan block stamp: %w", err)
 			}
-			w.RescanStartStamp = rescanStartStamp
+			w.RescanStartStamp = stamp
+		}
+		if config.RollbackStateFromBlock != 0 {
+			stamp, err := w.GetBlockStamp(config.RollbackStateFromBlock)
+			if err != nil {
+				return nil, fmt.Errorf("cannot get rescan block stamp: %w", err)
+			}
+			w.RollbackFromStamp = stamp
 		}
 	}
 

wallet/wallet.go

@@ -162,6 +162,7 @@ type Wallet struct {
 	rescanFinished      chan *RescanFinishedMsg
 	rescanLock          sync.Mutex
 	RescanStartStamp    *waddrmgr.BlockStamp
+	RollbackFromStamp   *waddrmgr.BlockStamp
 
 	// Channel for transaction creation requests.
 	createTxRequests chan createTxRequest
@@ -507,28 +508,34 @@ func (w *Wallet) syncWithChain(birthdayStamp *waddrmgr.BlockStamp) error {
 		addrmgrNs := tx.ReadWriteBucket(waddrmgrNamespaceKey)
 		txmgrNs := tx.ReadWriteBucket(wtxmgrNamespaceKey)
 
-		for height := rollbackStamp.Height; true; height-- {
-			hash, err := w.Manager.BlockHash(addrmgrNs, height)
-			if err != nil {
-				return err
-			}
-			chainHash, err := chainClient.GetBlockHash(int64(height))
-			if err != nil {
-				return err
-			}
-			header, err := chainClient.GetBlockHeader(chainHash)
-			if err != nil {
-				return err
-			}
+		// NB: if the rollback specified in config we will do rollback from here
+		if w.RollbackFromStamp != nil {
+			rollbackStamp = *w.RollbackFromStamp
+			rollback = true
+		} else {
+			for height := rollbackStamp.Height; true; height-- {
+				hash, err := w.Manager.BlockHash(addrmgrNs, height)
+				if err != nil {
+					return err
+				}
+				chainHash, err := chainClient.GetBlockHash(int64(height))
+				if err != nil {
+					return err
+				}
+				header, err := chainClient.GetBlockHeader(chainHash)
+				if err != nil {
+					return err
+				}
 
-			rollbackStamp.Hash = *chainHash
-			rollbackStamp.Height = height
-			rollbackStamp.Timestamp = header.Timestamp
+				rollbackStamp.Hash = *chainHash
+				rollbackStamp.Height = height
+				rollbackStamp.Timestamp = header.Timestamp
 
-			if bytes.Equal(hash[:], chainHash[:]) {
-				break
+				if bytes.Equal(hash[:], chainHash[:]) {
+					break
+				}
+				rollback = true
 			}
-			rollback = true
 		}
 
 		// If a rollback did not happen, we can proceed safely.