fix two bugs that regarding verifying shard_id (#8977)

Co-authored-by: Bowen Wang <bowenwang1996@users.noreply.github.com>
near · Apr 26, 2023 · 8c1b0df · 8c1b0df
1 parent 055a9e2
commit 8c1b0df
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 16 deletions.
diff --git a/chain/chain/src/chain.rs b/chain/chain/src/chain.rs
@@ -1107,6 +1107,9 @@ impl Chain {
                     byzantine_assert!(false);
                     return Err(Error::InvalidChunk);
                 }
+                if chunk_header.shard_id() != shard_id as ShardId {
+                    return Err(Error::InvalidShardId(chunk_header.shard_id()));
+                }
             }
         }
         block.check_validity().map_err(|e| e.into())

diff --git a/chain/chunks/src/lib.rs b/chain/chunks/src/lib.rs
@@ -1248,22 +1248,28 @@ impl ShardsManager {
             }
         };
 
-        match self.epoch_manager.verify_chunk_header_signature(header, &epoch_id, &ancestor_hash) {
-            Ok(false) => {
-                return if epoch_id_confirmed {
-                    byzantine_assert!(false);
-                    Err(Error::InvalidChunkSignature)
-                } else {
-                    // we are not sure if we are using the correct epoch id for validation, so
-                    // we can't be sure if the chunk header is actually invalid. Let's return
-                    // DbNotFoundError for now, which means we don't have all needed information yet
-                    Err(DBNotFoundErr(format!("block {:?}", header.prev_block_hash())).into())
-                };
-            }
-            Ok(true) => (),
-            Err(chain_error) => {
-                return Err(chain_error.into());
-            }
+        if !self.epoch_manager.verify_chunk_header_signature(header, &epoch_id, &ancestor_hash)? {
+            return if epoch_id_confirmed {
+                byzantine_assert!(false);
+                Err(Error::InvalidChunkSignature)
+            } else {
+                // we are not sure if we are using the correct epoch id for validation, so
+                // we can't be sure if the chunk header is actually invalid. Let's return
+                // DbNotFoundError for now, which means we don't have all needed information yet
+                Err(DBNotFoundErr(format!("block {:?}", header.prev_block_hash())).into())
+            };
+        }
+
+        if header.shard_id() >= self.epoch_manager.num_shards(&epoch_id)? {
+            return if epoch_id_confirmed {
+                byzantine_assert!(false);
+                Err(Error::InvalidChunkShardId)
+            } else {
+                // we are not sure if we are using the correct epoch id for validation, so
+                // we can't be sure if the chunk header is actually invalid. Let's return
+                // DbNotFoundError for now, which means we don't have all needed information yet
+                Err(DBNotFoundErr(format!("block {:?}", header.prev_block_hash())).into())
+            };
         }
 
         // 2. check protocol version