From 09494e07aaf151435eb2b2818b88ca408f9ed66c Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Wed, 16 Sep 2020 16:41:18 +1000
Subject: [PATCH] Add comments, safety to block shuffle id

---
 beacon_node/beacon_chain/src/beacon_chain.rs    |  3 ++-
 beacon_node/beacon_chain/src/errors.rs          |  1 +
 beacon_node/beacon_chain/src/shuffling_cache.rs | 16 +++++++++++-----
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
index 9d699cbf3b0..9c94da5e76a 100644
--- a/beacon_node/beacon_chain/src/beacon_chain.rs
+++ b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -2069,7 +2069,8 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             current: head_block.current_epoch_shuffling_id.clone(),
             next: head_block.next_epoch_shuffling_id.clone(),
         }
-        .id_for_epoch(shuffling_epoch);
+        .id_for_epoch(shuffling_epoch)
+        .ok_or_else(|| Error::InvalidShufflingId)?;
 
         // Obtain the shuffling cache, timing how long we wait.
         let cache_wait_timer =
diff --git a/beacon_node/beacon_chain/src/errors.rs b/beacon_node/beacon_chain/src/errors.rs
index 96f1c9a8411..37e1b3c8044 100644
--- a/beacon_node/beacon_chain/src/errors.rs
+++ b/beacon_node/beacon_chain/src/errors.rs
@@ -83,6 +83,7 @@ pub enum BeaconChainError {
     ObservedBlockProducersError(ObservedBlockProducersError),
     PruningError(PruningError),
     ArithError(ArithError),
+    InvalidShufflingId,
 }
 
 easy_from_to!(SlotProcessingError, BeaconChainError);
diff --git a/beacon_node/beacon_chain/src/shuffling_cache.rs b/beacon_node/beacon_chain/src/shuffling_cache.rs
index 28130a3c311..b9b41830ca4 100644
--- a/beacon_node/beacon_chain/src/shuffling_cache.rs
+++ b/beacon_node/beacon_chain/src/shuffling_cache.rs
@@ -47,21 +47,27 @@ impl ShufflingCache {
     }
 }
 
+/// Contains the shuffling IDs for a beacon block.
 pub struct BlockShufflingIds {
     pub current: ShufflingId,
     pub next: ShufflingId,
 }
 
 impl BlockShufflingIds {
-    pub fn id_for_epoch(&self, epoch: Epoch) -> ShufflingId {
+    /// Returns the shuffling ID for the given epoch.
+    ///
+    /// Returns `None` if `epoch` is prior to `self.current.shuffling_epoch`.
+    pub fn id_for_epoch(&self, epoch: Epoch) -> Option<ShufflingId> {
         if epoch == self.current.shuffling_epoch {
-            self.current.clone()
+            Some(self.current.clone())
         } else if epoch == self.next.shuffling_epoch {
-            self.next.clone()
-        } else {
+            Some(self.next.clone())
+        } else if epoch > self.next.shuffling_epoch {
             let mut shuffling_id = self.next.clone();
             shuffling_id.shuffling_epoch = epoch;
-            shuffling_id
+            Some(shuffling_id)
+        } else {
+            None
         }
     }
 }