From fb9c4fabf49ec7dd56034a89d847a31d9fa61fce Mon Sep 17 00:00:00 2001
From: Jacek Sieka <jacek@status.im>
Date: Wed, 5 Feb 2020 12:41:46 +0100
Subject: [PATCH] fix state rewind

* rewind fast path comparison was not taking skipped slots into account
properly
* less messy blockref creation
---
 beacon_chain/block_pool.nim | 25 +++++++++----------
 tests/test_block_pool.nim   | 49 +++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 13 deletions(-)

diff --git a/beacon_chain/block_pool.nim b/beacon_chain/block_pool.nim
index 49137d7eb..f4c47987e 100644
--- a/beacon_chain/block_pool.nim
+++ b/beacon_chain/block_pool.nim
@@ -249,9 +249,10 @@ proc init*(T: type BlockPool, db: BeaconChainDB): BlockPool =
   res
 
 proc addResolvedBlock(
-    pool: var BlockPool, state: var StateData, blockRoot: Eth2Digest,
+    pool: var BlockPool, state: BeaconState, blockRoot: Eth2Digest,
     signedBlock: SignedBeaconBlock, parent: BlockRef): BlockRef =
   logScope: pcs = "block_resolution"
+  doAssert state.slot == signedBlock.message.slot, "state must match block"
 
   let blockRef = BlockRef.init(blockRoot, signedBlock.message)
   link(parent, blockRef)
@@ -262,17 +263,10 @@ proc addResolvedBlock(
   # Resolved blocks should be stored in database
   pool.db.putBlock(blockRoot, signedBlock)
 
-  # TODO this is a bit ugly - we update state.data outside of this function then
-  #      set the rest here - need a blockRef to update it. Clean this up -
-  #      hopefully it won't be necessary by the time hash caching and the rest
-  #      is done..
-  doAssert state.data.data.slot == blockRef.slot
-  state.blck = blockRef
-
   # This block *might* have caused a justification - make sure we stow away
   # that information:
   let justifiedSlot =
-    state.data.data.current_justified_checkpoint.epoch.compute_start_slot_at_epoch()
+    state.current_justified_checkpoint.epoch.compute_start_slot_at_epoch()
 
   var foundHead: Option[Head]
   for head in pool.heads.mitems():
@@ -389,9 +383,12 @@ proc add*(
 
       return
 
-    # Careful, pool.tmpState is now partially inconsistent and will be updated
-    # inside addResolvedBlock
-    return pool.addResolvedBlock(pool.tmpState, blockRoot, signedBlock, parent)
+    # Careful, tmpState.data has been updated but not blck - we need to create
+    # the BlockRef first!
+    pool.tmpState.blck = pool.addResolvedBlock(
+      pool.tmpState.data.data, blockRoot, signedBlock, parent)
+
+    return pool.tmpState.blck
 
   # TODO already checked hash though? main reason to keep this is because
   # the pending pool calls this function back later in a loop, so as long
@@ -603,7 +600,7 @@ proc rewindState(pool: BlockPool, state: var StateData, bs: BlockSlot):
   var ancestors = @[pool.get(bs.blck)]
   # Common case: the last block applied is the parent of the block to apply:
   if not bs.blck.parent.isNil and state.blck.root == bs.blck.parent.root and
-      state.data.data.slot < bs.slot:
+      state.data.data.slot < bs.blck.slot:
     return ancestors
 
   # It appears that the parent root of the proposed new block is different from
@@ -693,6 +690,7 @@ proc updateStateData*(pool: BlockPool, state: var StateData, bs: BlockSlot) =
     if state.data.data.slot != bs.slot:
       # Might be that we're moving to the same block but later slot
       process_slots(state.data, bs.slot)
+      # TODO we will not save if multiple slots are skipped here
       pool.maybePutState(state.data, bs.blck)
 
     return # State already at the right spot
@@ -715,6 +713,7 @@ proc updateStateData*(pool: BlockPool, state: var StateData, bs: BlockSlot) =
     doAssert ok, "Blocks in database should never fail to apply.."
 
   # TODO check if this triggers rest of state transition, or should
+  # TODO we will not save if multiple slots are skipped here
   process_slots(state.data, bs.slot)
   pool.maybePutState(state.data, bs.blck)
 
diff --git a/tests/test_block_pool.nim b/tests/test_block_pool.nim
index 4578e1f9b..c28de7280 100644
--- a/tests/test_block_pool.nim
+++ b/tests/test_block_pool.nim
@@ -187,6 +187,55 @@ when const_preset == "minimal": # Too much stack space used on mainnet
         pool.head.blck == b1Add
         pool.headState.data.data.slot == b1Add.slot
 
+    timedTest "updateStateData sanity" & preset():
+      let
+        b1Add = pool.add(b1Root, b1)
+        b2Add = pool.add(b2Root, b2)
+        bs1 = BlockSlot(blck: b1Add, slot: b1.message.slot)
+        bs1_3 = b1Add.atSlot(3.Slot)
+        bs2 = BlockSlot(blck: b2Add, slot: b2.message.slot)
+        bs2_3 = b2Add.atSlot(3.Slot)
+
+      var tmpState = pool.headState
+
+      # move to specific block
+      pool.updateStateData(tmpState, bs1)
+
+      check:
+        tmpState.blck == b1Add
+        tmpState.data.data.slot == bs1.slot
+
+      # Skip slots
+      pool.updateStateData(tmpState, bs1_3) # skip slots
+
+      check:
+        tmpState.blck == b1Add
+        tmpState.data.data.slot == bs1_3.slot
+
+      # Move back slots, but not blocks
+      pool.updateStateData(tmpState, bs1_3.parent())
+      check:
+        tmpState.blck == b1Add
+        tmpState.data.data.slot == bs1_3.parent().slot
+
+      # Move to different block and slot
+      pool.updateStateData(tmpState, bs2_3)
+      check:
+        tmpState.blck == b2Add
+        tmpState.data.data.slot == bs2_3.slot
+
+      # Move back slot and block
+      pool.updateStateData(tmpState, bs1)
+      check:
+        tmpState.blck == b1Add
+        tmpState.data.data.slot == bs1.slot
+
+      # Move back to genesis
+      pool.updateStateData(tmpState, bs1.parent())
+      check:
+        tmpState.blck == b1Add.parent
+        tmpState.data.data.slot == bs1.parent.slot
+
   suite "BlockPool finalization tests" & preset():
     setup:
       var