Add extra justification change check

beacon_node/beacon_chain/src/builder.rs

@@ -391,6 +391,7 @@ where
                 .ok_or_else(|| "fork_choice_backend requires a genesis_block_root")?;
 
             let backend = ProtoArrayForkChoice::new(
+                finalized_checkpoint.beacon_block.slot,
                 // Note: here we set the `justified_epoch` to be the same as the epoch of the
                 // finalized checkpoint. Whilst this finalized checkpoint may actually point to
                 // a _later_ justified checkpoint, that checkpoint won't yet exist in the fork

beacon_node/beacon_chain/src/fork_choice.rs

@@ -20,6 +20,7 @@ pub enum Error {
     BeaconStateError(BeaconStateError),
     StoreError(StoreError),
     BeaconChainError(Box<BeaconChainError>),
+    UnknownBlockSlot(Hash256),
 }
 
 #[derive(PartialEq, Clone, Encode, Decode)]
@@ -91,6 +92,7 @@ impl JustificationManager {
         &mut self,
         state: &BeaconState<T::EthSpec>,
         chain: &BeaconChain<T>,
+        proto_array: &ProtoArrayForkChoice,
     ) -> Result<()> {
         let new_checkpoint = &state.current_justified_checkpoint;
 
@@ -115,12 +117,18 @@ impl JustificationManager {
                     .start_slot(T::EthSpec::slots_per_epoch()),
             )?;
 
+            let new_justified_block_slot = proto_array
+                .block_slot(&new_checkpoint.root)
+                .ok_or_else(|| Error::UnknownBlockSlot(new_checkpoint.root))?;
+
             // If the new justified checkpoint is an ancestor of the current justified checkpoint,
             // it is always safe to change it.
-            //
-            // TODO: check the slot of the block to see if we can update it. Might involve DB read
-            // or maybe we can add it to proto-array.
-            if new_checkpoint_ancestor == Some(self.justified_checkpoint.root) {
+            if new_checkpoint_ancestor == Some(self.justified_checkpoint.root)
+                && new_justified_block_slot
+                    >= new_checkpoint
+                        .epoch
+                        .start_slot(T::EthSpec::slots_per_epoch())
+            {
                 self.justified_checkpoint = new_checkpoint_balances.clone()
             }
 
@@ -236,7 +244,6 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
                 justified_checkpoint.epoch,
                 remove_alias(justified_checkpoint.root),
                 finalized_checkpoint.epoch,
-                remove_alias(finalized_checkpoint.root),
                 &justified_checkpoint.balances,
             )
             .map_err(Into::into);
@@ -261,7 +268,7 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
 
         self.justification_manager
             .write()
-            .process_state(state, chain)?;
+            .process_state(state, chain, &self.backend)?;
         self.justification_manager.write().update(chain)?;
 
         // TODO: be more stringent about changing the finalized checkpoint (i.e., check for
@@ -285,6 +292,7 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
         // This does not apply a vote to the block, it just makes fork choice aware of the block so
         // it can still be identified as the head even if it doesn't have any votes.
         self.backend.process_block(
+            block.slot,
             block_root,
             block.parent_root,
             state.current_justified_checkpoint.epoch,

eth2/proto_array_fork_choice/src/lib.rs

@@ -7,7 +7,7 @@ use ssz::{Decode, Encode};
 use ssz_container::SszContainer;
 use ssz_derive::{Decode, Encode};
 use std::collections::HashMap;
-use types::{Epoch, Hash256};
+use types::{Epoch, Hash256, Slot};
 
 pub const DEFAULT_PRUNE_THRESHOLD: usize = 256;
 
@@ -96,6 +96,7 @@ impl PartialEq for ProtoArrayForkChoice {
 
 impl ProtoArrayForkChoice {
     pub fn new(
+        finalized_block_slot: Slot,
         justified_epoch: Epoch,
         finalized_epoch: Epoch,
         finalized_root: Hash256,
@@ -110,7 +111,13 @@ impl ProtoArrayForkChoice {
         };
 
         proto_array
-            .on_new_block(finalized_root, None, justified_epoch, finalized_epoch)
+            .on_new_block(
+                finalized_block_slot,
+                finalized_root,
+                None,
+                justified_epoch,
+                finalized_epoch,
+            )
             .map_err(|e| format!("Failed to add finalized block to proto_array: {:?}", e))?;
 
         Ok(Self {
@@ -139,6 +146,7 @@ impl ProtoArrayForkChoice {
 
     pub fn process_block(
         &self,
+        slot: Slot,
         block_root: Hash256,
         parent_root: Hash256,
         justified_epoch: Epoch,
@@ -147,6 +155,7 @@ impl ProtoArrayForkChoice {
         self.proto_array
             .write()
             .on_new_block(
+                slot,
                 block_root,
                 Some(parent_root),
                 justified_epoch,
@@ -160,7 +169,6 @@ impl ProtoArrayForkChoice {
         justified_epoch: Epoch,
         justified_root: Hash256,
         finalized_epoch: Epoch,
-        finalized_root: Hash256,
         justified_state_balances: &[u64],
     ) -> Result<Hash256, String> {
         let mut proto_array = self.proto_array.write();
@@ -211,6 +219,15 @@ impl ProtoArrayForkChoice {
         self.proto_array.read().indices.contains_key(block_root)
     }
 
+    pub fn block_slot(&self, block_root: &Hash256) -> Option<Slot> {
+        let proto_array = self.proto_array.read();
+
+        let i = proto_array.indices.get(block_root)?;
+        let block = proto_array.nodes.get(*i)?;
+
+        Some(block.slot)
+    }
+
     pub fn latest_message(&self, validator_index: usize) -> Option<(Hash256, Epoch)> {
         let votes = self.votes.read();
 

eth2/proto_array_fork_choice/src/proto_array.rs

@@ -1,10 +1,13 @@
 use crate::Error;
 use ssz_derive::{Decode, Encode};
 use std::collections::HashMap;
-use types::{Epoch, Hash256};
+use types::{Epoch, Hash256, Slot};
 
 #[derive(Clone, PartialEq, Debug, Encode, Decode)]
 pub struct ProtoNode {
+    /// The `slot` is not necessary for `ProtoArray`, it just exists so external components can
+    /// easily query the block slot. This is useful for upstream fork choice logic.
+    pub slot: Slot,
     root: Hash256,
     parent: Option<usize>,
     justified_epoch: Epoch,
@@ -131,6 +134,7 @@ impl ProtoArray {
     /// It is only sane to supply a `None` parent for the genesis block.
     pub fn on_new_block(
         &mut self,
+        slot: Slot,
         root: Hash256,
         parent_opt: Option<Hash256>,
         justified_epoch: Epoch,
@@ -139,6 +143,7 @@ impl ProtoArray {
         let node_index = self.nodes.len();
 
         let node = ProtoNode {
+            slot,
             root,
             parent: parent_opt.and_then(|parent| self.indices.get(&parent).copied()),
             justified_epoch,
@@ -376,7 +381,7 @@ impl ProtoArray {
                     // There is no current best-child and the child is viable.
                     change_to_child
                 } else {
-                    // There is no current-best child but the child is not viable.
+                    // There is no current best-child but the child is not viable.
                     no_change
                 }
             };
@@ -388,6 +393,7 @@ impl ProtoArray {
 
         parent.best_child = new_best_child;
         parent.best_descendant = new_best_descendant;
+
         Ok(())
     }