Prevent sync lookups from reverting to awaiting block (#6443)

* Prevent sync lookups from reverting to awaiting block

* Remove stale comment

beacon_node/network/src/sync/block_lookups/common.rs

@@ -13,7 +13,7 @@ use std::sync::Arc;
 use types::blob_sidecar::FixedBlobSidecarList;
 use types::{DataColumnSidecarList, SignedBeaconBlock};
 
-use super::single_block_lookup::DownloadResult;
+use super::single_block_lookup::{ComponentRequests, DownloadResult};
 use super::SingleLookupId;
 
 #[derive(Debug, Copy, Clone)]
@@ -42,7 +42,7 @@ pub trait RequestState<T: BeaconChainTypes> {
         &self,
         id: Id,
         peer_id: PeerId,
-        downloaded_block: Option<Arc<SignedBeaconBlock<T::EthSpec>>>,
+        expected_blobs: usize,
         cx: &mut SyncNetworkContext<T>,
     ) -> Result<LookupRequestResult, LookupRequestError>;
 
@@ -61,7 +61,7 @@ pub trait RequestState<T: BeaconChainTypes> {
     fn response_type() -> ResponseType;
 
     /// A getter for the `BlockRequestState` or `BlobRequestState` associated with this trait.
-    fn request_state_mut(request: &mut SingleBlockLookup<T>) -> &mut Self;
+    fn request_state_mut(request: &mut SingleBlockLookup<T>) -> Result<&mut Self, &'static str>;
 
     /// A getter for a reference to the `SingleLookupRequestState` associated with this trait.
     fn get_state(&self) -> &SingleLookupRequestState<Self::VerifiedResponseType>;
@@ -77,7 +77,7 @@ impl<T: BeaconChainTypes> RequestState<T> for BlockRequestState<T::EthSpec> {
         &self,
         id: SingleLookupId,
         peer_id: PeerId,
-        _: Option<Arc<SignedBeaconBlock<T::EthSpec>>>,
+        _: usize,
         cx: &mut SyncNetworkContext<T>,
     ) -> Result<LookupRequestResult, LookupRequestError> {
         cx.block_lookup_request(id, peer_id, self.requested_block_root)
@@ -107,8 +107,8 @@ impl<T: BeaconChainTypes> RequestState<T> for BlockRequestState<T::EthSpec> {
     fn response_type() -> ResponseType {
         ResponseType::Block
     }
-    fn request_state_mut(request: &mut SingleBlockLookup<T>) -> &mut Self {
-        &mut request.block_request_state
+    fn request_state_mut(request: &mut SingleBlockLookup<T>) -> Result<&mut Self, &'static str> {
+        Ok(&mut request.block_request_state)
     }
     fn get_state(&self) -> &SingleLookupRequestState<Self::VerifiedResponseType> {
         &self.state
@@ -125,10 +125,10 @@ impl<T: BeaconChainTypes> RequestState<T> for BlobRequestState<T::EthSpec> {
         &self,
         id: Id,
         peer_id: PeerId,
-        downloaded_block: Option<Arc<SignedBeaconBlock<T::EthSpec>>>,
+        expected_blobs: usize,
         cx: &mut SyncNetworkContext<T>,
     ) -> Result<LookupRequestResult, LookupRequestError> {
-        cx.blob_lookup_request(id, peer_id, self.block_root, downloaded_block)
+        cx.blob_lookup_request(id, peer_id, self.block_root, expected_blobs)
             .map_err(LookupRequestError::SendFailedNetwork)
     }
 
@@ -150,8 +150,13 @@ impl<T: BeaconChainTypes> RequestState<T> for BlobRequestState<T::EthSpec> {
     fn response_type() -> ResponseType {
         ResponseType::Blob
     }
-    fn request_state_mut(request: &mut SingleBlockLookup<T>) -> &mut Self {
-        &mut request.blob_request_state
+    fn request_state_mut(request: &mut SingleBlockLookup<T>) -> Result<&mut Self, &'static str> {
+        match &mut request.component_requests {
+            ComponentRequests::WaitingForBlock => Err("waiting for block"),
+            ComponentRequests::ActiveBlobRequest(request, _) => Ok(request),
+            ComponentRequests::ActiveCustodyRequest { .. } => Err("expecting custody request"),
+            ComponentRequests::NotNeeded { .. } => Err("not needed"),
+        }
     }
     fn get_state(&self) -> &SingleLookupRequestState<Self::VerifiedResponseType> {
         &self.state
@@ -169,10 +174,10 @@ impl<T: BeaconChainTypes> RequestState<T> for CustodyRequestState<T::EthSpec> {
         id: Id,
         // TODO(das): consider selecting peers that have custody but are in this set
         _peer_id: PeerId,
-        downloaded_block: Option<Arc<SignedBeaconBlock<T::EthSpec>>>,
+        _: usize,
         cx: &mut SyncNetworkContext<T>,
     ) -> Result<LookupRequestResult, LookupRequestError> {
-        cx.custody_lookup_request(id, self.block_root, downloaded_block)
+        cx.custody_lookup_request(id, self.block_root)
             .map_err(LookupRequestError::SendFailedNetwork)
     }
 
@@ -200,8 +205,13 @@ impl<T: BeaconChainTypes> RequestState<T> for CustodyRequestState<T::EthSpec> {
     fn response_type() -> ResponseType {
         ResponseType::CustodyColumn
     }
-    fn request_state_mut(request: &mut SingleBlockLookup<T>) -> &mut Self {
-        &mut request.custody_request_state
+    fn request_state_mut(request: &mut SingleBlockLookup<T>) -> Result<&mut Self, &'static str> {
+        match &mut request.component_requests {
+            ComponentRequests::WaitingForBlock => Err("waiting for block"),
+            ComponentRequests::ActiveBlobRequest { .. } => Err("expecting blob request"),
+            ComponentRequests::ActiveCustodyRequest(request) => Ok(request),
+            ComponentRequests::NotNeeded { .. } => Err("not needed"),
+        }
     }
     fn get_state(&self) -> &SingleLookupRequestState<Self::VerifiedResponseType> {
         &self.state