feat: add partial witness epoch check (#11430)

This is mostly to prevent malicious chunk producer from sending a
witness with a valid height and old epoch id.

Part of #11301.

chain/client/src/stateless_validation/chunk_validator/orphan_witness_handling.rs

@@ -9,7 +9,7 @@ use crate::Client;
 use near_chain::Block;
 use near_chain_primitives::Error;
 use near_primitives::stateless_validation::ChunkStateWitness;
-use near_primitives::types::{BlockHeight, EpochId};
+use near_primitives::types::BlockHeight;
 use std::ops::Range;
 
 /// We keep only orphan witnesses that are within this distance of
@@ -68,18 +68,6 @@ impl Client {
             return Ok(HandleOrphanWitnessOutcome::TooBig(witness_size));
         }
 
-        // Try to find the EpochId to which this witness will belong based on its height.
-        // It's not always possible to determine the exact epoch_id because the exact
-        // starting height of the next epoch isn't known until it actually starts,
-        // so things can get unclear around epoch boundaries.
-        // Let's collect the epoch_ids in which the witness might possibly be.
-        let possible_epochs =
-            self.epoch_manager.possible_epochs_of_height_around_tip(&chain_head, witness_height)?;
-
-        if !possible_epochs.contains(&witness.epoch_id) {
-            return Ok(HandleOrphanWitnessOutcome::UnsupportedEpochId(witness.epoch_id));
-        }
-
         // Orphan witness is OK, save it to the pool
         tracing::debug!(target: "client", "Saving an orphaned ChunkStateWitness to orphan pool");
         self.chunk_validator.orphan_witness_pool.add_orphan_state_witness(witness, witness_size);
@@ -145,5 +133,4 @@ pub enum HandleOrphanWitnessOutcome {
     SavedToPool,
     TooBig(usize),
     TooFarFromHead { head_height: BlockHeight, witness_height: BlockHeight },
-    UnsupportedEpochId(EpochId),
 }

chain/client/src/stateless_validation/partial_witness/partial_witness_actor.rs

@@ -298,6 +298,8 @@ impl PartialWitnessActor {
     /// Function to validate the partial encoded state witness. We check the following
     /// - shard_id is valid
     /// - we are one of the validators for the chunk
+    /// - height_created is in (last_final_height..chain_head_height + MAX_HEIGHTS_AHEAD] range
+    /// - epoch_id is within epoch_manager's possible_epochs_of_height_around_tip
     /// - part_ord is valid and within range of the number of expected parts for this chunk
     /// - partial_witness signature is valid and from the expected chunk_producer
     /// TODO(stateless_validation): Include checks from handle_orphan_state_witness in orphan_witness_handling.rs
@@ -366,7 +368,25 @@ impl PartialWitnessActor {
                     head.height,
                 )));
             }
+
+            // Try to find the EpochId to which this witness will belong based on its height.
+            // It's not always possible to determine the exact epoch_id because the exact
+            // starting height of the next epoch isn't known until it actually starts,
+            // so things can get unclear around epoch boundaries.
+            // Let's collect the epoch_ids in which the witness might possibly be.
+            let possible_epochs = self
+                .epoch_manager
+                .possible_epochs_of_height_around_tip(&head, partial_witness.height_created())?;
+            if !possible_epochs.contains(&partial_witness.epoch_id()) {
+                return Err(Error::InvalidPartialChunkStateWitness(format!(
+                    "EpochId {:?} in PartialEncodedStateWitness at height {} is not in the possible list of epochs {:?}",
+                    partial_witness.epoch_id(),
+                    partial_witness.height_created(),
+                    possible_epochs
+                )));
+            }
         }
+
         if !self.epoch_manager.verify_partial_witness_signature(&partial_witness)? {
             return Err(Error::InvalidPartialChunkStateWitness("Invalid signature".to_string()));
         }