Reorg on the light wallet

Breeze/src/Breeze.Wallet/LightWalletSyncManager.cs

@@ -9,6 +9,7 @@ using Stratis.Bitcoin.Notifications;
 using Stratis.Bitcoin.Utilities;
 using Stratis.Bitcoin.Wallet;
 using Stratis.Bitcoin.Wallet.Notifications;
+using System.Collections.Generic;
 
 namespace Breeze.Wallet
 {
@@ -21,6 +22,8 @@ namespace Breeze.Wallet
         private readonly ILogger logger;
         private readonly Signals signals;
 
+        private ChainedBlock lastReceivedBlock;
+
         public LightWalletSyncManager(ILoggerFactory loggerFactory, IWalletManager walletManager, ConcurrentChain chain, Network network,
             BlockNotification blockNotification, Signals signals)// :base(loggerFactory, walletManager, chain, network)
         {
@@ -35,6 +38,10 @@ namespace Breeze.Wallet
         /// <inheritdoc />
         public async Task Initialize()
         {
+            this.lastReceivedBlock = this.chain.GetBlock(this.walletManager.LastReceivedBlock);
+            if (this.lastReceivedBlock == null)
+                throw new WalletException("the wallet tip was not found in the main chain");
+
             // get the chain headers. This needs to be up-to-date before we really do anything
             await this.WaitForChainDownloadAsync();
 
@@ -106,16 +113,41 @@ namespace Breeze.Wallet
         /// <inheritdoc />
         public void ProcessBlock(Block block)
         {
-            var chainedBlock = this.chain.GetBlock(block.GetHash());
-
-            // if the newly received block is too far forward, ignore it (for example when we start receiving blocks before the wallets start their syncing)
-            if (chainedBlock.Height > this.walletManager.LastBlockHeight() + 1)
+            // if the new block previous hash is the same as the 
+            // wallet hash then just pass the block to the manager 
+            if (block.Header.HashPrevBlock != this.lastReceivedBlock.HashBlock)
             {
-                this.logger.LogDebug($"block received with height: {chainedBlock.Height} and hash: {block.Header.GetHash()} is too far in advance. Ignoring.");
-                return;
+                // if previous block does not match there might have 
+                // been a reorg, check if the wallet is still on the main chain
+                var current = this.chain.GetBlock(this.lastReceivedBlock.HashBlock);
+                if (current == null)
+                {
+                    // the current wallet hash was not found on the main chain
+                    // a reorg happenend so bring the wallet back top the last known fork
+
+                    var blockstoremove = new List<uint256>();
+                    var fork = this.lastReceivedBlock;
+
+                    // we walk back the chained block object to find the fork
+                    while (this.chain.GetBlock(fork.HashBlock) == null)
+                    {
+                        blockstoremove.Add(fork.HashBlock);
+                        fork = fork.Previous;
+                    }
+
+                    this.walletManager.RemoveBlocks(fork);
+                }
+                else if (current.Height > this.lastReceivedBlock.Height)
+                {
+                    // the wallet is falling behind we need to catch up
+                    this.logger.LogDebug($"block received with height: {current.Height} and hash: {block.Header.GetHash()} is too far in advance. Ignoring.");
+                    this.SyncFrom(this.lastReceivedBlock.Height);
+                    return;
+                }
             }
 
-            this.walletManager.ProcessBlock(block);
+            var chainedBlock = this.chain.GetBlock(block.GetHash());
+            this.walletManager.ProcessBlock(block, chainedBlock);
         }
 
         /// <inheritdoc />