feat: avoid processing state witness multiple times

chain/chain-primitives/src/error.rs

@@ -133,6 +133,8 @@ pub enum Error {
     InvalidChunkState(Box<ChunkState>),
     #[error("Invalid Chunk State Witness: {0}")]
     InvalidChunkStateWitness(String),
+    #[error("Duplicate Chunk State Witness: {0}")]
+    DuplicateChunkStateWitness(String),
     #[error("Invalid Chunk Endorsement")]
     InvalidChunkEndorsement,
     /// Invalid chunk mask
@@ -270,6 +272,7 @@ impl Error {
             | Error::InvalidChunkProofs(_)
             | Error::InvalidChunkState(_)
             | Error::InvalidChunkStateWitness(_)
+            | Error::DuplicateChunkStateWitness(_)
             | Error::InvalidChunkEndorsement
             | Error::InvalidChunkMask
             | Error::InvalidStateRoot
@@ -344,6 +347,7 @@ impl Error {
             Error::InvalidChunkProofs(_) => "invalid_chunk_proofs",
             Error::InvalidChunkState(_) => "invalid_chunk_state",
             Error::InvalidChunkStateWitness(_) => "invalid_chunk_state_witness",
+            Error::DuplicateChunkStateWitness(_) => "duplicate_chunk_state_witness",
             Error::InvalidChunkEndorsement => "invalid_chunk_endorsement",
             Error::InvalidChunkMask => "invalid_chunk_mask",
             Error::InvalidStateRoot => "invalid_state_root",

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

@@ -27,8 +27,8 @@ use near_primitives::merkle::merklize;
 use near_primitives::receipt::Receipt;
 use near_primitives::sharding::{ChunkHash, ReceiptProof, ShardChunkHeader};
 use near_primitives::stateless_validation::{
-    ChunkEndorsement, ChunkStateWitness, ChunkStateWitnessAck, ChunkStateWitnessSize,
-    SignedEncodedChunkStateWitness,
+    ChunkEndorsement, ChunkProductionKey, ChunkStateWitness, ChunkStateWitnessAck,
+    ChunkStateWitnessSize, SignedEncodedChunkStateWitness,
 };
 use near_primitives::transaction::SignedTransaction;
 use near_primitives::types::chunk_extra::ChunkExtra;
@@ -46,6 +46,8 @@ use std::sync::Arc;
 // Keeping a threshold of 5 block producers should be sufficient for most scenarios.
 const NUM_NEXT_BLOCK_PRODUCERS_TO_SEND_CHUNK_ENDORSEMENT: u64 = 5;
 
+const RECEIVED_STATE_WITNESSES_CACHE_SIZE: usize = 100;
+
 /// A module that handles chunk validation logic. Chunk validation refers to a
 /// critical process of stateless validation, where chunk validators (certain
 /// validators selected to validate the chunk) verify that the chunk's state
@@ -60,6 +62,7 @@ pub struct ChunkValidator {
     chunk_endorsement_tracker: Arc<ChunkEndorsementTracker>,
     orphan_witness_pool: OrphanStateWitnessPool,
     validation_spawner: Arc<dyn AsyncComputationSpawner>,
+    received_witnesses: lru::LruCache<ChunkProductionKey, ()>,
 }
 
 impl ChunkValidator {
@@ -80,6 +83,7 @@ impl ChunkValidator {
             chunk_endorsement_tracker,
             orphan_witness_pool: OrphanStateWitnessPool::new(orphan_witness_pool_size),
             validation_spawner,
+            received_witnesses: lru::LruCache::new(RECEIVED_STATE_WITNESSES_CACHE_SIZE),
         }
     }
 
@@ -89,7 +93,7 @@ impl ChunkValidator {
     /// The chunk is validated asynchronously, if you want to wait for the processing to finish
     /// you can use the `processing_done_tracker` argument (but it's optional, it's safe to pass None there).
     pub fn start_validating_chunk(
-        &self,
+        &mut self,
         state_witness: ChunkStateWitness,
         chain: &Chain,
         processing_done_tracker: Option<ProcessingDoneTracker>,
@@ -103,6 +107,8 @@ impl ChunkValidator {
             )));
         }
 
+        self.check_duplicate_witness(&state_witness)?;
+
         let pre_validation_result = pre_validate_chunk_state_witness(
             &state_witness,
             chain,
@@ -142,6 +148,24 @@ impl ChunkValidator {
         });
         Ok(())
     }
+
+    /// Verifies that we haven't already processed state witness for the corresponding
+    /// (shard_id, epoch_id, height_created). This protects against malicious chunk
+    /// producers wasting stateless validator resources by making it apply chunk multiple
+    /// times.
+    fn check_duplicate_witness(&mut self, state_witness: &ChunkStateWitness) -> Result<(), Error> {
+        let chunk_production_key = ChunkProductionKey::from_witness(&state_witness);
+        if self.received_witnesses.contains(&chunk_production_key) {
+            return Err(Error::DuplicateChunkStateWitness(format!(
+                "Already received witness for height {} shard {} epoch {:?}",
+                chunk_production_key.height_created,
+                chunk_production_key.shard_id,
+                chunk_production_key.epoch_id,
+            )));
+        }
+        self.received_witnesses.push(chunk_production_key, ());
+        Ok(())
+    }
 }
 
 /// Checks that proposed `transactions` are valid for a chunk with `chunk_header`.

core/primitives/src/stateless_validation.rs

@@ -376,6 +376,27 @@ impl ChunkValidatorAssignments {
     }
 }
 
+/// This struct contains combination of fields that uniquely identify chunk production.
+/// It means that for a given instance only one chunk could be produced.
+/// The main use case it to track processed state witness in order to protect stateless
+/// validator against malicious chunk producers.
+#[derive(Hash, PartialEq, Eq)]
+pub struct ChunkProductionKey {
+    pub shard_id: ShardId,
+    pub epoch_id: EpochId,
+    pub height_created: BlockHeight,
+}
+
+impl ChunkProductionKey {
+    pub fn from_witness(witness: &ChunkStateWitness) -> Self {
+        Self {
+            shard_id: witness.chunk_header.shard_id(),
+            epoch_id: witness.epoch_id.clone(),
+            height_created: witness.chunk_header.height_created(),
+        }
+    }
+}
+
 fn decompress_with_limit(data: &[u8], limit: usize) -> std::io::Result<Vec<u8>> {
     let mut buf = Vec::new().limit(limit).writer();
     match zstd::stream::copy_decode(data, &mut buf) {