[stateless_validation] Versioned chunk endorsement (#11856)

This PR is a pre-step for the work around merging validation of chunk
endorsement and partial witness.

This PR has the following changes
- New protocol feature ChunkEndorsementV2 which would introduce a new
network message and versioning of ChunkEndorsement.
- ChunkEndorsement is changed from struct to a versioned enum. The
original/old chunk endorsement message would work with
ChunkEndorsementV1 till deprecated.
- The outgoing network stack would use the original
RoutedMessageBody::ChunkEndorsement for V1 while we would use
RoutedMessageBody::VersionedChunkEndorsement for futue versions.

No functional change in this PR.

chain/chain/src/test_utils/kv_runtime.rs

@@ -32,7 +32,7 @@ use near_primitives::shard_layout;
 use near_primitives::shard_layout::{ShardLayout, ShardUId};
 use near_primitives::sharding::{ChunkHash, ShardChunkHeader};
 use near_primitives::state_part::PartId;
-use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsement;
+use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsementV1;
 use near_primitives::stateless_validation::partial_witness::PartialEncodedStateWitness;
 use near_primitives::stateless_validation::validator_assignment::ChunkValidatorAssignments;
 use near_primitives::transaction::{
@@ -945,7 +945,7 @@ impl EpochManagerAdapter for MockEpochManager {
     fn verify_chunk_endorsement(
         &self,
         _chunk_header: &ShardChunkHeader,
-        _endorsement: &ChunkEndorsement,
+        _endorsement: &ChunkEndorsementV1,
     ) -> Result<bool, Error> {
         Ok(true)
     }

chain/chain/src/tests/simple_chain.rs

@@ -32,7 +32,7 @@ fn build_chain() {
     //     cargo insta test --accept -p near-chain --features nightly -- tests::simple_chain::build_chain
     let hash = chain.head().unwrap().last_block_hash;
     if cfg!(feature = "nightly") {
-        insta::assert_snapshot!(hash, @"C3zeKRZubVungxfrSdq379TSCYnuz2YzjEkcJTdm3pU4");
+        insta::assert_snapshot!(hash, @"H9qfaTFJ7jFMEmJX9eYQ3kJQ6vwATAhidfbjwa9PvCPK");
     } else {
         insta::assert_snapshot!(hash, @"EeGa9BFTyrPoM56iZBteXZfhq86g4k1tjMCPprqYyfKF");
     }
@@ -50,7 +50,7 @@ fn build_chain() {
 
     let hash = chain.head().unwrap().last_block_hash;
     if cfg!(feature = "nightly") {
-        insta::assert_snapshot!(hash, @"EjLaoHRiAdRp2NcDqwbMcAYYxGfcv5R7GuYUNfRpaJvB");
+        insta::assert_snapshot!(hash, @"HPRdNbp8dJL2JA5bRSmLg6GJpGXx6xmw472vVAWLpXhi");
     } else {
         insta::assert_snapshot!(hash, @"Em7E1W5xPMTToyiGbRAES2iAKJQDsZ5m8ymvy4S7adxu");
     }

chain/client/src/stateless_validation/chunk_endorsement_tracker.rs

@@ -1,6 +1,8 @@
 use lru::LruCache;
 use near_chain::ChainStoreAccess;
-use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsement;
+use near_primitives::stateless_validation::chunk_endorsement::{
+    ChunkEndorsement, ChunkEndorsementV1,
+};
 use near_primitives::stateless_validation::validator_assignment::EndorsementStats;
 use std::collections::HashMap;
 use std::num::NonZeroUsize;
@@ -45,36 +47,51 @@ struct ChunkEndorsementTrackerInner {
     /// This is keyed on chunk_hash and account_id of validator to avoid duplicates.
     /// Chunk endorsements would later be used as a part of block production.
     chunk_endorsements:
-        LruCache<ChunkHash, (ShardChunkHeader, HashMap<AccountId, ChunkEndorsement>)>,
+        LruCache<ChunkHash, (ShardChunkHeader, HashMap<AccountId, ChunkEndorsementV1>)>,
     /// We store chunk endorsements to be processed later because we did not have
     /// chunks ready at the time we received that endorsements from validators.
     /// This is keyed on chunk_hash and account_id of validator to avoid duplicates.
-    pending_chunk_endorsements: LruCache<ChunkHash, HashMap<AccountId, ChunkEndorsement>>,
+    pending_chunk_endorsements: LruCache<ChunkHash, HashMap<AccountId, ChunkEndorsementV1>>,
 }
 
 impl Client {
     pub fn process_chunk_endorsement(
         &mut self,
         endorsement: ChunkEndorsement,
+    ) -> Result<(), Error> {
+        // TODO(ChunkEndorsementV2): Remove chunk_header once tracker_v1 is deprecated
+        let chunk_header =
+            match self.chain.chain_store().get_partial_chunk(endorsement.chunk_hash()) {
+                Ok(chunk) => Some(chunk.cloned_header()),
+                Err(Error::ChunkMissing(_)) => None,
+                Err(error) => return Err(error),
+            };
+        match endorsement {
+            ChunkEndorsement::V1(endorsement) => {
+                self.chunk_endorsement_tracker.process_chunk_endorsement(endorsement, chunk_header)
+            }
+        }
+    }
+}
+
+impl ChunkEndorsementTracker {
+    pub fn process_chunk_endorsement(
+        &self,
+        endorsement: ChunkEndorsementV1,
+        chunk_header: Option<ShardChunkHeader>,
     ) -> Result<(), Error> {
         // We need the chunk header in order to process the chunk endorsement.
         // If we don't have the header, then queue it up for when we do have the header.
         // We must use the partial chunk (as opposed to the full chunk) in order to get
         // the chunk header, because we may not be tracking that shard.
-        match self.chain.chain_store().get_partial_chunk(endorsement.chunk_hash()) {
-            Ok(chunk) => self
-                .chunk_endorsement_tracker
-                .process_chunk_endorsement(&chunk.cloned_header(), endorsement),
-            Err(Error::ChunkMissing(_)) => {
-                tracing::debug!(target: "client", ?endorsement, "Endorsement arrived before chunk.");
-                self.chunk_endorsement_tracker.add_chunk_endorsement_to_pending_cache(endorsement)
+        match chunk_header {
+            Some(chunk_header) => {
+                self.process_chunk_endorsement_with_chunk_header(&chunk_header, endorsement)
             }
-            Err(error) => return Err(error),
+            None => self.add_chunk_endorsement_to_pending_cache(endorsement),
         }
     }
-}
 
-impl ChunkEndorsementTracker {
     pub fn new(epoch_manager: Arc<dyn EpochManagerAdapter>) -> Self {
         Self {
             epoch_manager: epoch_manager.clone(),
@@ -100,18 +117,18 @@ impl ChunkEndorsementTracker {
     /// Add the chunk endorsement to a cache of pending chunk endorsements (if not yet there).
     pub(crate) fn add_chunk_endorsement_to_pending_cache(
         &self,
-        endorsement: ChunkEndorsement,
+        endorsement: ChunkEndorsementV1,
     ) -> Result<(), Error> {
         self.inner.lock().unwrap().process_chunk_endorsement_impl(endorsement, None, false)
     }
 
     /// Function to process an incoming chunk endorsement from chunk validators.
     /// We first verify the chunk endorsement and then store it in a cache.
     /// We would later include the endorsements in the block production.
-    pub(crate) fn process_chunk_endorsement(
+    fn process_chunk_endorsement_with_chunk_header(
         &self,
         chunk_header: &ShardChunkHeader,
-        endorsement: ChunkEndorsement,
+        endorsement: ChunkEndorsementV1,
     ) -> Result<(), Error> {
         let _span = tracing::debug_span!(target: "client", "process_chunk_endorsement", chunk_hash=?chunk_header.chunk_hash(), shard_id=chunk_header.shard_id()).entered();
         // Validate the endorsement before locking the mutex to improve performance.
@@ -164,7 +181,7 @@ impl ChunkEndorsementTrackerInner {
     /// Otherwise, we store the endorsement in a separate cache of endorsements to be processed when the chunk is ready.
     fn process_chunk_endorsement_impl(
         &mut self,
-        endorsement: ChunkEndorsement,
+        endorsement: ChunkEndorsementV1,
         chunk_header: Option<&ShardChunkHeader>,
         already_validated: bool,
     ) -> Result<(), Error> {

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

@@ -16,7 +16,9 @@ use near_epoch_manager::EpochManagerAdapter;
 use near_network::types::{NetworkRequests, PeerManagerMessageRequest};
 use near_o11y::log_assert;
 use near_primitives::sharding::ShardChunkHeader;
-use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsement;
+use near_primitives::stateless_validation::chunk_endorsement::{
+    ChunkEndorsement, ChunkEndorsementV1,
+};
 use near_primitives::stateless_validation::state_witness::{
     ChunkStateWitness, ChunkStateWitnessAck, ChunkStateWitnessSize,
 };
@@ -227,12 +229,12 @@ pub(crate) fn send_chunk_endorsement_to_block_producers(
         "send_chunk_endorsement",
     );
 
-    let endorsement = ChunkEndorsement::new(chunk_header.chunk_hash(), signer);
+    let endorsement = ChunkEndorsementV1::new(chunk_header.chunk_hash(), signer);
     for block_producer in block_producers {
         if signer.validator_id() == &block_producer {
             // Our own endorsements are not always valid (see issue #11750).
             if let Err(err) = chunk_endorsement_tracker
-                .process_chunk_endorsement(chunk_header, endorsement.clone())
+                .process_chunk_endorsement(endorsement.clone(), Some(chunk_header.clone()))
             {
                 tracing::warn!(
                     target: "client",
@@ -242,7 +244,10 @@ pub(crate) fn send_chunk_endorsement_to_block_producers(
             }
         } else {
             network_sender.send(PeerManagerMessageRequest::NetworkRequests(
-                NetworkRequests::ChunkEndorsement(block_producer, endorsement.clone()),
+                NetworkRequests::ChunkEndorsement(
+                    block_producer,
+                    ChunkEndorsement::V1(endorsement.clone()),
+                ),
             ));
         }
     }

chain/client/src/test_utils/client.rs

@@ -21,7 +21,7 @@ use near_primitives::block::Block;
 use near_primitives::hash::CryptoHash;
 use near_primitives::merkle::{merklize, PartialMerkleTree};
 use near_primitives::sharding::{EncodedShardChunk, ShardChunk};
-use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsement;
+use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsementV1;
 use near_primitives::transaction::SignedTransaction;
 use near_primitives::types::{BlockHeight, ShardId};
 use near_primitives::utils::MaybeValidated;
@@ -249,7 +249,7 @@ pub fn create_chunk(
     let mut block_merkle_tree = PartialMerkleTree::clone(&block_merkle_tree);
 
     let signer = client.validator_signer.get().unwrap();
-    let endorsement = ChunkEndorsement::new(chunk.cloned_header().chunk_hash(), signer.as_ref());
+    let endorsement = ChunkEndorsementV1::new(chunk.cloned_header().chunk_hash(), signer.as_ref());
     block_merkle_tree.insert(*last_block.hash());
     let block = Block::produce(
         PROTOCOL_VERSION,

chain/client/src/test_utils/test_env.rs

@@ -27,7 +27,9 @@ use near_primitives::epoch_manager::RngSeed;
 use near_primitives::errors::InvalidTxError;
 use near_primitives::hash::CryptoHash;
 use near_primitives::sharding::{ChunkHash, PartialEncodedChunk};
-use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsement;
+use near_primitives::stateless_validation::chunk_endorsement::{
+    ChunkEndorsement, ChunkEndorsementV1,
+};
 use near_primitives::stateless_validation::state_witness::ChunkStateWitness;
 use near_primitives::test_utils::create_test_signer;
 use near_primitives::transaction::{Action, FunctionCallAction, SignedTransaction};
@@ -432,11 +434,13 @@ impl TestEnv {
     /// Wait until an endorsement for `chunk_hash` appears in the network messages send by
     /// the Client with index `client_idx`. Times out after CHUNK_ENDORSEMENTS_TIMEOUT.
     /// Doesn't process or consume the message, it just waits until the message appears on the network_adapter.
+    /// TODO(ChunkEndorsementV2): This function is only used by orphan_chunk_state_witnesses test.
+    /// Can remove once we shift to ChunkEndorsementV2.
     pub fn wait_for_chunk_endorsement(
         &mut self,
         client_idx: usize,
         chunk_hash: &ChunkHash,
-    ) -> Result<ChunkEndorsement, TimeoutError> {
+    ) -> Result<ChunkEndorsementV1, TimeoutError> {
         let start_time = Instant::now();
         let network_adapter = self.network_adapters[client_idx].clone();
         loop {
@@ -445,8 +449,13 @@ impl TestEnv {
                 match &request {
                     PeerManagerMessageRequest::NetworkRequests(
                         NetworkRequests::ChunkEndorsement(_receiver_account_id, endorsement),
-                    ) if endorsement.chunk_hash() == chunk_hash => {
-                        endorsement_opt = Some(endorsement.clone());
+                    ) => {
+                        let endorsement = match endorsement {
+                            ChunkEndorsement::V1(endorsement) => endorsement,
+                        };
+                        if endorsement.chunk_hash() == chunk_hash {
+                            endorsement_opt = Some(endorsement.clone());
+                        }
                     }
                     _ => {}
                 };

chain/epoch-manager/src/adapter.rs

@@ -13,7 +13,7 @@ use near_primitives::errors::EpochError;
 use near_primitives::hash::CryptoHash;
 use near_primitives::shard_layout::{account_id_to_shard_id, ShardLayout, ShardLayoutError};
 use near_primitives::sharding::{ChunkHash, ShardChunkHeader};
-use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsement;
+use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsementV1;
 use near_primitives::stateless_validation::partial_witness::PartialEncodedStateWitness;
 use near_primitives::stateless_validation::validator_assignment::ChunkValidatorAssignments;
 use near_primitives::types::validator_stake::ValidatorStake;
@@ -418,7 +418,7 @@ pub trait EpochManagerAdapter: Send + Sync {
     fn verify_chunk_endorsement(
         &self,
         chunk_header: &ShardChunkHeader,
-        endorsement: &ChunkEndorsement,
+        endorsement: &ChunkEndorsementV1,
     ) -> Result<bool, Error>;
 
     fn verify_partial_witness_signature(
@@ -1049,7 +1049,7 @@ impl EpochManagerAdapter for EpochManagerHandle {
     fn verify_chunk_endorsement(
         &self,
         chunk_header: &ShardChunkHeader,
-        endorsement: &ChunkEndorsement,
+        endorsement: &ChunkEndorsementV1,
     ) -> Result<bool, Error> {
         if &chunk_header.chunk_hash() != endorsement.chunk_hash() {
             return Err(Error::InvalidChunkEndorsement);

chain/epoch-manager/src/tests/mod.rs

@@ -19,7 +19,7 @@ use near_primitives::epoch_manager::EpochConfig;
 use near_primitives::hash::hash;
 use near_primitives::shard_layout::ShardLayout;
 use near_primitives::sharding::{ShardChunkHeader, ShardChunkHeaderV3};
-use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsement;
+use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsementV1;
 use near_primitives::stateless_validation::partial_witness::PartialEncodedStateWitness;
 use near_primitives::types::ValidatorKickoutReason::{
     NotEnoughBlocks, NotEnoughChunkEndorsements, NotEnoughChunks,
@@ -2818,15 +2818,15 @@ fn test_verify_chunk_endorsements() {
     let chunk_header = test_chunk_header(&h, signer.as_ref());
 
     // check chunk endorsement validity
-    let mut chunk_endorsement = ChunkEndorsement::new(chunk_header.chunk_hash(), signer.as_ref());
+    let mut chunk_endorsement = ChunkEndorsementV1::new(chunk_header.chunk_hash(), signer.as_ref());
     assert!(epoch_manager.verify_chunk_endorsement(&chunk_header, &chunk_endorsement).unwrap());
 
     // check invalid chunk endorsement signature
     chunk_endorsement.signature = Signature::default();
     assert!(!epoch_manager.verify_chunk_endorsement(&chunk_header, &chunk_endorsement).unwrap());
 
     // check chunk endorsement invalidity when chunk header and chunk endorsement don't match
-    let chunk_endorsement = ChunkEndorsement::new(h[3].into(), signer.as_ref());
+    let chunk_endorsement = ChunkEndorsementV1::new(h[3].into(), signer.as_ref());
     let err =
         epoch_manager.verify_chunk_endorsement(&chunk_header, &chunk_endorsement).unwrap_err();
     match err {
@@ -2836,7 +2836,7 @@ fn test_verify_chunk_endorsements() {
 
     // check chunk endorsement invalidity when signer is not chunk validator
     let bad_signer = Arc::new(create_test_signer("test2"));
-    let chunk_endorsement = ChunkEndorsement::new(chunk_header.chunk_hash(), bad_signer.as_ref());
+    let chunk_endorsement = ChunkEndorsementV1::new(chunk_header.chunk_hash(), bad_signer.as_ref());
     let err =
         epoch_manager.verify_chunk_endorsement(&chunk_header, &chunk_endorsement).unwrap_err();
     match err {

chain/network/src/network_protocol/mod.rs

@@ -8,6 +8,7 @@ mod proto_conv;
 mod state_sync;
 pub use edge::*;
 use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsement;
+use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsementV1;
 use near_primitives::stateless_validation::partial_witness::PartialEncodedStateWitness;
 use near_primitives::stateless_validation::state_witness::ChunkStateWitnessAck;
 pub use peer::*;
@@ -531,10 +532,12 @@ pub enum RoutedMessageBody {
     _UnusedVersionedStateResponse,
     PartialEncodedChunkForward(PartialEncodedChunkForwardMsg),
     _UnusedChunkStateWitness,
-    ChunkEndorsement(ChunkEndorsement),
+    /// TODO(ChunkEndorsementV2): Deprecate once we move to VersionedChunkEndorsement
+    ChunkEndorsement(ChunkEndorsementV1),
     ChunkStateWitnessAck(ChunkStateWitnessAck),
     PartialEncodedStateWitness(PartialEncodedStateWitness),
     PartialEncodedStateWitnessForward(PartialEncodedStateWitness),
+    VersionedChunkEndorsement(ChunkEndorsement),
 }
 
 impl RoutedMessageBody {
@@ -607,6 +610,9 @@ impl fmt::Debug for RoutedMessageBody {
             RoutedMessageBody::PartialEncodedStateWitnessForward(_) => {
                 write!(f, "PartialEncodedStateWitnessForward")
             }
+            RoutedMessageBody::VersionedChunkEndorsement(_) => {
+                write!(f, "VersionedChunkEndorsement")
+            }
         }
     }
 }

chain/network/src/peer/peer_actor.rs

@@ -45,6 +45,7 @@ use near_o11y::{handler_debug_span, log_assert, WithSpanContext};
 use near_performance_metrics_macros::perf;
 use near_primitives::hash::CryptoHash;
 use near_primitives::network::{AnnounceAccount, PeerId};
+use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsement;
 use near_primitives::types::EpochId;
 use near_primitives::utils::DisplayOption;
 use near_primitives::version::{
@@ -1028,6 +1029,7 @@ impl PeerActor {
                 None
             }
             RoutedMessageBody::ChunkEndorsement(endorsement) => {
+                let endorsement = ChunkEndorsement::V1(endorsement);
                 network_state.client.send_async(ChunkEndorsementMessage(endorsement)).await.ok();
                 None
             }
@@ -1043,6 +1045,10 @@ impl PeerActor {
                     .send(PartialEncodedStateWitnessForwardMessage(witness));
                 None
             }
+            RoutedMessageBody::VersionedChunkEndorsement(endorsement) => {
+                network_state.client.send_async(ChunkEndorsementMessage(endorsement)).await.ok();
+                None
+            }
             body => {
                 tracing::error!(target: "network", "Peer received unexpected message type: {:?}", body);
                 None

chain/network/src/peer_manager/peer_manager_actor.rs

@@ -30,6 +30,7 @@ use near_o11y::{handler_debug_span, handler_trace_span, WithSpanContext};
 use near_performance_metrics_macros::perf;
 use near_primitives::block::GenesisId;
 use near_primitives::network::{AnnounceAccount, PeerId};
+use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsement;
 use near_primitives::views::{
     ConnectionInfoView, EdgeView, KnownPeerStateView, NetworkGraphView, PeerStoreView,
     RecentOutboundConnectionsView, SnapshotHostInfoView, SnapshotHostsView,
@@ -980,11 +981,12 @@ impl PeerManagerActor {
                 NetworkResponses::NoResponse
             }
             NetworkRequests::ChunkEndorsement(target, endorsement) => {
-                self.state.send_message_to_account(
-                    &self.clock,
-                    &target,
-                    RoutedMessageBody::ChunkEndorsement(endorsement),
-                );
+                let msg = match endorsement {
+                    ChunkEndorsement::V1(endorsement) => {
+                        RoutedMessageBody::ChunkEndorsement(endorsement)
+                    }
+                };
+                self.state.send_message_to_account(&self.clock, &target, msg);
                 NetworkResponses::NoResponse
             }
             NetworkRequests::PartialEncodedStateWitness(validator_witness_tuple) => {

chain/network/src/rate_limits/messages_limits.rs

@@ -217,6 +217,7 @@ fn get_key_and_token_cost(message: &PeerMessage) -> Option<(RateLimitedPeerMessa
             RoutedMessageBody::PartialEncodedStateWitnessForward(_) => {
                 Some((PartialEncodedStateWitnessForward, 1))
             }
+            RoutedMessageBody::VersionedChunkEndorsement(_) => Some((ChunkEndorsement, 1)),
             RoutedMessageBody::Ping(_)
             | RoutedMessageBody::Pong(_)
             | RoutedMessageBody::_UnusedChunkStateWitness