From c6b8a2498f3013fe36409bee7c9f3463459f3136 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Moraczy=C5=84ski?= <marekm2504@gmail.com>
Date: Wed, 27 Jul 2022 09:07:19 +0200
Subject: [PATCH] fix terminal block checking (#4318)

* fix terminal block checking

* add test

* fix + tests
---
 .../Nethermind.Merge.Plugin/PoSSwitcher.cs    | 31 +++++++------------
 .../SyncServerTests.cs                        |  9 +++---
 2 files changed, 17 insertions(+), 23 deletions(-)

diff --git a/src/Nethermind/Nethermind.Merge.Plugin/PoSSwitcher.cs b/src/Nethermind/Nethermind.Merge.Plugin/PoSSwitcher.cs
index f2a5204541a..a7916dbfe30 100644
--- a/src/Nethermind/Nethermind.Merge.Plugin/PoSSwitcher.cs
+++ b/src/Nethermind/Nethermind.Merge.Plugin/PoSSwitcher.cs
@@ -176,6 +176,15 @@ public void ForkchoiceUpdated(BlockHeader newHeadHash, Keccak finalizedHash)
                 _logger.Trace(
                     $"GetBlockConsensusInfo {header.ToString(BlockHeader.Format.FullHashAndNumber)} header.IsPostMerge: {header.IsPostMerge} header.TotalDifficulty {header.TotalDifficulty} header.Difficulty {header.Difficulty} TTD: {_specProvider.TerminalTotalDifficulty} MergeBlockNumber {_specProvider.MergeBlockNumber}, TransitionFinished: {TransitionFinished}");
 
+            if ((header.TotalDifficulty ?? 0) != 0 && dontTrustTotalDifficulty && header.IsGenesis == false)
+            {
+                BlockHeader? parentHeader = _blockTree.FindParentHeader(header, BlockTreeLookupOptions.None);
+                if (parentHeader != null && parentHeader.TotalDifficulty != 0)
+                    header.TotalDifficulty = parentHeader.TotalDifficulty + header.Difficulty;
+                else
+                    header.TotalDifficulty = null;
+            }
+
             bool isTerminal = false, isPostMerge;
             if (header.IsPostMerge) // block from Engine API, there is no need to check more cases
             {
@@ -206,25 +215,9 @@ public void ForkchoiceUpdated(BlockHeader newHeadHash, Keccak finalizedHash)
                 }
                 else
                 {
-                    if (header.TotalDifficulty >= _specProvider.TerminalTotalDifficulty && dontTrustTotalDifficulty && header.IsGenesis == false)
-                    {
-                        BlockHeader? parentHeader = _blockTree.FindParentHeader(header, BlockTreeLookupOptions.None);
-                        if (parentHeader != null && parentHeader.TotalDifficulty != 0)
-                            header.TotalDifficulty = parentHeader.TotalDifficulty + header.Difficulty;
-                        else
-                            header.TotalDifficulty = null;
-                    }
-
-                    if (header.TotalDifficulty == null)
-                    {
-                        isPostMerge = header.Difficulty == 0;
-                        isTerminal = false; // we can't say if block isTerminal if we don't have TD
-                    }
-                    else
-                    {
-                        isTerminal = header.IsTerminalBlock(_specProvider); // we're checking if block is terminal if not it should be PostMerge block
-                        isPostMerge = !isTerminal;
-                    }
+
+                    isTerminal = header.IsTerminalBlock(_specProvider); // we're checking if block is terminal if not it should be PostMerge block
+                    isPostMerge = !isTerminal;
                 }
             }
 
diff --git a/src/Nethermind/Nethermind.Synchronization.Test/SyncServerTests.cs b/src/Nethermind/Nethermind.Synchronization.Test/SyncServerTests.cs
index 9c44eb79689..4dea43e8606 100644
--- a/src/Nethermind/Nethermind.Synchronization.Test/SyncServerTests.cs
+++ b/src/Nethermind/Nethermind.Synchronization.Test/SyncServerTests.cs
@@ -23,6 +23,7 @@
 using Nethermind.Core.Crypto;
 using Nethermind.Core.Test.Builders;
 using Nethermind.Db;
+using Nethermind.Int256;
 using Nethermind.Logging;
 using Nethermind.Merge.Plugin;
 using Nethermind.Merge.Plugin.Handlers;
@@ -206,15 +207,15 @@ public void Terminal_block_with_lower_td_should_not_change_best_suggested_but_sh
             Assert.AreEqual(remoteBestBlock.Hash, localBlockTree.FindBlock(remoteBestBlock.Hash, BlockTreeLookupOptions.None)!.Hash);
         }
 
-        [Test]
-        public void Terminal_blocks_with_incorrect_td_from_peer()
+        [TestCase(10000000)]
+        [TestCase(20000000)]
+        public void Fake_total_difficulty_from_peer_does_not_trick_the_node(long ttd)
         {
             Context ctx = new();
             BlockTree remoteBlockTree = Build.A.BlockTree().OfChainLength(10).TestObject;
             BlockTree localBlockTree = Build.A.BlockTree().OfChainLength(9).TestObject;
             TestSpecProvider testSpecProvider = new(London.Instance);
-            testSpecProvider.TerminalTotalDifficulty = 10000000;
-
+            testSpecProvider.TerminalTotalDifficulty = (UInt256)ttd;
 
             PoSSwitcher poSSwitcher = new(new MergeConfig() { Enabled = true }, new SyncConfig(), new MemDb(), localBlockTree, testSpecProvider, LimboLogs.Instance);
             MergeSealEngine sealEngine = new MergeSealEngine(new SealEngine(new NethDevSealEngine(), Always.Valid), poSSwitcher, new MergeSealValidator(poSSwitcher, Always.Valid), LimboLogs.Instance);