fix(chunks): process partial encoded chunk properly handles missing blocks

chain/chunks/src/lib.rs

@@ -38,8 +38,7 @@ use crate::chunk_cache::{EncodedChunksCache, EncodedChunksCacheEntry};
 pub use crate::types::Error;
 
 mod chunk_cache;
-#[cfg(test)]
-mod test_utils;
+pub mod test_utils;
 mod types;
 
 const CHUNK_PRODUCER_BLACKLIST_SIZE: usize = 100;
@@ -949,10 +948,26 @@ impl ShardsManager {
         let chunk_hash = header.chunk_hash();
 
         // 1. Checking signature validity
-        if !self.runtime_adapter.verify_chunk_header_signature(&header)? {
-            byzantine_assert!(false);
-            return Err(Error::InvalidChunkSignature);
-        }
+        match self.runtime_adapter.verify_chunk_header_signature(&header) {
+            Ok(false) => {
+                byzantine_assert!(false);
+                return Err(Error::InvalidChunkSignature);
+            }
+            Ok(true) => (),
+            Err(chain_error) => {
+                return match chain_error.kind() {
+                    near_chain::ErrorKind::BlockMissing(_)
+                    | near_chain::ErrorKind::DBNotFoundErr(_) => {
+                        // We can't check if this chunk came from a valid chunk producer because
+                        // we don't know `prev_block`, so return that we need a block.
+                        Ok(ProcessPartialEncodedChunkResult::NeedBlock)
+                    }
+                    // Some other error kind happened during the signature check, we don't
+                    // know how to handle it.
+                    _ => Err(Error::ChainError(chain_error)),
+                };
+            }
+        };
 
         // 2. Leave if we received known chunk
         if let Some(entry) = self.encoded_chunks.get(&chunk_hash) {

chain/chunks/src/test_utils.rs

@@ -5,15 +5,23 @@ use chrono::Utc;
 use near_chain::test_utils::KeyValueRuntime;
 use near_chain::types::RuntimeAdapter;
 use near_chain::ChainStore;
+use near_crypto::KeyType;
+use near_network::test_utils::MockNetworkAdapter;
 use near_primitives::block::BlockHeader;
 use near_primitives::hash::{self, CryptoHash};
-use near_primitives::sharding::ChunkHash;
-use near_primitives::types::BlockHeight;
+use near_primitives::merkle;
+use near_primitives::sharding::{
+    ChunkHash, PartialEncodedChunk, PartialEncodedChunkPart, ReedSolomonWrapper, ShardChunkHeader,
+};
 use near_primitives::types::{AccountId, ShardId};
+use near_primitives::types::{BlockHeight, MerkleHash};
+use near_primitives::validator_signer::InMemoryValidatorSigner;
 use near_primitives::version::PROTOCOL_VERSION;
 use near_store::Store;
 
-use crate::{Seal, SealsManager, ACCEPTING_SEAL_PERIOD_MS, PAST_SEAL_HEIGHT_HORIZON};
+use crate::{
+    Seal, SealsManager, ShardsManager, ACCEPTING_SEAL_PERIOD_MS, PAST_SEAL_HEIGHT_HORIZON,
+};
 
 pub struct SealsManagerTestFixture {
     pub mock_chunk_producer: AccountId,
@@ -31,15 +39,7 @@ impl Default for SealsManagerTestFixture {
     fn default() -> Self {
         let store = near_store::test_utils::create_test_store();
         // 12 validators, 3 shards => 4 validators per shard
-        let letters = (b'A'..=b'Z')
-            .take(12)
-            .map(|c| {
-                let mut s = String::with_capacity(1);
-                s.push(c as char);
-                s
-            })
-            .collect();
-        let validators = vec![letters];
+        let validators = make_validators(12);
         let mock_runtime =
             KeyValueRuntime::new_with_validators(Arc::clone(&store), validators, 1, 3, 5);
 
@@ -130,3 +130,159 @@ impl SealsManagerTestFixture {
         demur.sent = demur.sent - d;
     }
 }
+
+pub struct ChunkForwardingTestFixture {
+    pub mock_runtime: Arc<KeyValueRuntime>,
+    pub mock_network: Arc<MockNetworkAdapter>,
+    pub chain_store: ChainStore,
+    pub mock_part_ords: Vec<u64>,
+    pub mock_chunk_part_owner: AccountId,
+    pub mock_shard_tracker: AccountId,
+    pub mock_chunk_header: ShardChunkHeader,
+    pub mock_chunk_parts: Vec<PartialEncodedChunkPart>,
+    pub rs: ReedSolomonWrapper,
+}
+
+impl Default for ChunkForwardingTestFixture {
+    fn default() -> Self {
+        let store = near_store::test_utils::create_test_store();
+        // 12 validators, 3 shards => 4 validators per shard
+        let validators = make_validators(12);
+        let mock_runtime = Arc::new(KeyValueRuntime::new_with_validators(
+            Arc::clone(&store),
+            validators.clone(),
+            1,
+            3,
+            5,
+        ));
+        let mock_network = Arc::new(MockNetworkAdapter::default());
+
+        let data_parts = mock_runtime.num_data_parts();
+        let parity_parts = mock_runtime.num_total_parts() - data_parts;
+        let mut rs = ReedSolomonWrapper::new(data_parts, parity_parts);
+        let mock_parent_hash = CryptoHash::default();
+        let mock_height: BlockHeight = 1;
+        let mock_shard_id: ShardId = 0;
+        let mock_epoch_id = mock_runtime.get_epoch_id_from_prev_block(&mock_parent_hash).unwrap();
+        let mock_chunk_producer =
+            mock_runtime.get_chunk_producer(&mock_epoch_id, mock_height, mock_shard_id).unwrap();
+        let signer = InMemoryValidatorSigner::from_seed(
+            &mock_chunk_producer,
+            KeyType::ED25519,
+            &mock_chunk_producer,
+        );
+        let mock_shard_tracker = validators
+            .iter()
+            .flatten()
+            .find(|v| {
+                if *v == &mock_chunk_producer {
+                    false
+                } else {
+                    let tracks_shard = mock_runtime.cares_about_shard(
+                        Some(*v),
+                        &mock_parent_hash,
+                        mock_shard_id,
+                        false,
+                    ) || mock_runtime.will_care_about_shard(
+                        Some(*v),
+                        &mock_parent_hash,
+                        mock_shard_id,
+                        false,
+                    );
+                    tracks_shard
+                }
+            })
+            .cloned()
+            .unwrap();
+        let mock_chunk_part_owner = validators
+            .into_iter()
+            .flatten()
+            .find(|v| v != &mock_chunk_producer && v != &mock_shard_tracker)
+            .unwrap();
+
+        let mut producer_shard_manager = ShardsManager::new(
+            Some(mock_chunk_producer.clone()),
+            mock_runtime.clone(),
+            mock_network.clone(),
+        );
+        let receipts = Vec::new();
+        let receipts_hashes = mock_runtime.build_receipts_hashes(&receipts);
+        let (receipts_root, _) = merkle::merklize(&receipts_hashes);
+        let (mock_chunk, mock_merkles) = producer_shard_manager
+            .create_encoded_shard_chunk(
+                mock_parent_hash,
+                Default::default(),
+                Default::default(),
+                mock_height,
+                mock_shard_id,
+                0,
+                1000,
+                0,
+                Vec::new(),
+                Vec::new(),
+                &receipts,
+                receipts_root,
+                MerkleHash::default(),
+                &signer,
+                &mut rs,
+            )
+            .unwrap();
+
+        let all_part_ords: Vec<u64> =
+            (0..mock_chunk.content.parts.len()).map(|p| p as u64).collect();
+        let mock_part_ords = all_part_ords
+            .iter()
+            .copied()
+            .filter(|p| {
+                mock_runtime.get_part_owner(&mock_parent_hash, *p).unwrap() == mock_chunk_part_owner
+            })
+            .collect();
+        let encoded_chunk = mock_chunk.create_partial_encoded_chunk(
+            all_part_ords.clone(),
+            Vec::new(),
+            &mock_merkles,
+        );
+        let chain_store = ChainStore::new(store, 0);
+
+        ChunkForwardingTestFixture {
+            mock_runtime,
+            mock_network,
+            chain_store,
+            mock_part_ords,
+            mock_chunk_part_owner,
+            mock_shard_tracker,
+            mock_chunk_header: encoded_chunk.header,
+            mock_chunk_parts: encoded_chunk.parts,
+            rs,
+        }
+    }
+}
+
+impl ChunkForwardingTestFixture {
+    pub fn make_partial_encoded_chunk(&self, part_ords: &[u64]) -> PartialEncodedChunk {
+        let parts = part_ords
+            .iter()
+            .copied()
+            .flat_map(|ord| self.mock_chunk_parts.iter().find(|part| part.part_ord == ord))
+            .cloned()
+            .collect();
+        PartialEncodedChunk { header: self.mock_chunk_header.clone(), parts, receipts: Vec::new() }
+    }
+}
+
+fn make_validators(n: usize) -> Vec<Vec<AccountId>> {
+    if n > 26 {
+        panic!("I can't make that many validators!");
+    }
+
+    let letters = (b'a'..=b'z')
+        .take(n)
+        .map(|c| {
+            let mut s = String::with_capacity(1);
+            s.push(c as char);
+            s
+        })
+        .collect();
+
+    vec![letters]
+}

chain/client/src/client.rs

@@ -1398,6 +1398,8 @@ mod test {
 
     use near_chain::{ChainGenesis, RuntimeAdapter};
     use near_chain_configs::Genesis;
+    use near_chunks::test_utils::ChunkForwardingTestFixture;
+    use near_chunks::ProcessPartialEncodedChunkResult;
     use near_crypto::KeyType;
     use near_primitives::block::{Approval, ApprovalInner};
     use near_primitives::hash::hash;
@@ -1407,17 +1409,21 @@ mod test {
 
     use crate::test_utils::TestEnv;
 
-    #[test]
-    fn test_pending_approvals() {
+    fn create_runtimes() -> Vec<Arc<dyn RuntimeAdapter>> {
         let store = create_test_store();
         let genesis = Genesis::test(vec!["test0", "test1"], 1);
-        let runtimes: Vec<Arc<dyn RuntimeAdapter>> = vec![Arc::new(neard::NightshadeRuntime::new(
+        vec![Arc::new(neard::NightshadeRuntime::new(
             Path::new("."),
             store,
             Arc::new(genesis),
             vec![],
             vec![],
-        ))];
+        ))]
+    }
+
+    #[test]
+    fn test_pending_approvals() {
+        let runtimes = create_runtimes();
         let mut env = TestEnv::new_with_runtime(ChainGenesis::test(), 1, 1, runtimes);
         let signer = InMemoryValidatorSigner::from_seed("test1", KeyType::ED25519, "test1");
         let parent_hash = hash(&[1]);
@@ -1428,4 +1434,25 @@ mod test {
         let expected = vec![("test1".to_string(), approval)].into_iter().collect::<HashMap<_, _>>();
         assert_eq!(approvals, Some(expected));
     }
+
+    #[test]
+    fn test_process_partial_encoded_chunk_with_missing_block() {
+        let runtimes = create_runtimes();
+        let mut env = TestEnv::new_with_runtime(ChainGenesis::test(), 1, 1, runtimes);
+        let client = &mut env.clients[0];
+        let chunk_producer = ChunkForwardingTestFixture::default();
+        let mut mock_chunk = chunk_producer.make_partial_encoded_chunk(&[0]);
+        // change the prev_block to some unknown block
+        mock_chunk.header.inner.prev_block_hash = hash(b"some_prev_block");
+        mock_chunk.header.init();
+
+        // process_partial_encoded_chunk should return Ok(NeedBlock) if the chunk is
+        // based on a missing block.
+        let result = client.shards_mgr.process_partial_encoded_chunk(
+            mock_chunk,
+            client.chain.mut_store(),
+            &mut client.rs,
+        );
+        assert!(matches!(result, Ok(ProcessPartialEncodedChunkResult::NeedBlock)));
+    }
 }