Add block receive height limit (#1870)

* Add block cache limit * Remove known hashes * Optimize * Update Blockchain.cs * Clean code * Rename method * Prevent process far blocks * Block ACL * Remove size control #1870 (comment) * Move knownHashes * Remove no sender * Rename * Rename * Update RemoteNode.ProtocolHandler.cs Co-authored-by: erikzhang <[email protected]>
neo-project · Sep 6, 2020 · 5bea103 · 5bea103
1 parent c071898
commit 5bea103
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 14 deletions.
diff --git a/src/neo/Ledger/Blockchain.cs b/src/neo/Ledger/Blockchain.cs
@@ -1,5 +1,6 @@
 using Akka.Actor;
 using Akka.Configuration;
+using Akka.IO;
 using Neo.Cryptography.ECC;
 using Neo.IO;
 using Neo.IO.Actors;
@@ -28,6 +29,7 @@ public class FillMemoryPool { public IEnumerable<Transaction> Transactions; }
         public class FillCompleted { }
         internal class PreverifyCompleted { public Transaction Transaction; public VerifyResult Result; public bool Relay; }
         public class RelayResult { public IInventory Inventory; public VerifyResult Result; }
+        private class UnverifiedBlocksList { public LinkedList<Block> Blocks = new LinkedList<Block>(); public HashSet<IActorRef> Nodes = new HashSet<IActorRef>(); }
 
         public static readonly uint MillisecondsPerBlock = ProtocolSettings.Default.MillisecondsPerBlock;
         public static readonly TimeSpan TimePerBlock = TimeSpan.FromMilliseconds(MillisecondsPerBlock);
@@ -61,7 +63,7 @@ public class RelayResult { public IInventory Inventory; public VerifyResult Resu
         private readonly List<UInt256> header_index = new List<UInt256>();
         private uint stored_header_count = 0;
         private readonly Dictionary<UInt256, Block> block_cache = new Dictionary<UInt256, Block>();
-        private readonly Dictionary<uint, LinkedList<Block>> block_cache_unverified = new Dictionary<uint, LinkedList<Block>>();
+        private readonly Dictionary<uint, UnverifiedBlocksList> block_cache_unverified = new Dictionary<uint, UnverifiedBlocksList>();
         internal readonly RelayCache RelayCache = new RelayCache(100);
         private SnapshotView currentSnapshot;
 
@@ -273,19 +275,30 @@ private void OnImport(IEnumerable<Block> blocks, bool verify)
         private void AddUnverifiedBlockToCache(Block block)
         {
             // Check if any block proposal for height `block.Index` exists
-            if (!block_cache_unverified.TryGetValue(block.Index, out LinkedList<Block> blocks))
+            if (!block_cache_unverified.TryGetValue(block.Index, out var list))
             {
-                // There are no blocks, a new LinkedList is created and, consequently, the current block is added to the list
-                blocks = new LinkedList<Block>();
-                block_cache_unverified.Add(block.Index, blocks);
+                // There are no blocks, a new UnverifiedBlocksList is created and, consequently, the current block is added to the list
+                list = new UnverifiedBlocksList();
+                block_cache_unverified.Add(block.Index, list);
             }
-            // Check if any block with the hash being added already exists on possible candidates to be processed
-            foreach (var unverifiedBlock in blocks)
+            else
             {
-                if (block.Hash == unverifiedBlock.Hash)
+                // Check if any block with the hash being added already exists on possible candidates to be processed
+                foreach (var unverifiedBlock in list.Blocks)
+                {
+                    if (block.Hash == unverifiedBlock.Hash)
+                        return;
+                }
+
+                if (!list.Nodes.Add(Sender))
+                {
+                    // Same index with different hash
+                    Sender.Tell(Tcp.Abort.Instance);
                     return;
+                }
             }
-            blocks.AddLast(block);
+
+            list.Blocks.AddLast(block);
         }
 
         private void OnFillMemoryPool(IEnumerable<Transaction> transactions)
@@ -340,9 +353,9 @@ private VerifyResult OnNewBlock(Block block)
                 system.LocalNode.Tell(Message.Create(MessageCommand.Ping, PingPayload.Create(Singleton.Height + 1)));
                 Persist(block);
                 SaveHeaderHashList();
-                if (block_cache_unverified.TryGetValue(Height + 1, out LinkedList<Block> unverifiedBlocks))
+                if (block_cache_unverified.TryGetValue(Height + 1, out var unverifiedBlocks))
                 {
-                    foreach (var unverifiedBlock in unverifiedBlocks)
+                    foreach (var unverifiedBlock in unverifiedBlocks.Blocks)
                         Self.Tell(unverifiedBlock, ActorRefs.NoSender);
                     block_cache_unverified.Remove(Height + 1);
                 }

diff --git a/src/neo/Network/P2P/RemoteNode.ProtocolHandler.cs b/src/neo/Network/P2P/RemoteNode.ProtocolHandler.cs
@@ -287,18 +287,19 @@ private void OnGetHeadersMessageReceived(GetBlockByIndexPayload payload)
         private void OnInventoryReceived(IInventory inventory)
         {
             pendingKnownHashes.Remove(inventory.Hash);
-            knownHashes.Add(inventory.Hash);
-            system.TaskManager.Tell(inventory);
-            system.Blockchain.Tell(inventory, ActorRefs.NoSender);
             switch (inventory)
             {
                 case Transaction transaction:
                     system.Consensus?.Tell(transaction);
                     break;
                 case Block block:
+                    if (block.Index > Blockchain.Singleton.Height + InvPayload.MaxHashesCount) return;
                     UpdateLastBlockIndex(block.Index, false);
                     break;
             }
+            knownHashes.Add(inventory.Hash);
+            system.TaskManager.Tell(inventory);
+            system.Blockchain.Tell(inventory);
         }
 
         private void OnInvMessageReceived(InvPayload payload)