Add sync lookup custody request state (#6257)

* Add sync lookup custody request state

* Review PR

* clippy

* Merge branch 'unstable' of https://github.com/sigp/lighthouse into peerdas-network-lookup

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

@@ -1,14 +1,16 @@
 use crate::sync::block_lookups::single_block_lookup::{
     LookupRequestError, SingleBlockLookup, SingleLookupRequestState,
 };
-use crate::sync::block_lookups::{BlobRequestState, BlockRequestState, PeerId};
+use crate::sync::block_lookups::{
+    BlobRequestState, BlockRequestState, CustodyRequestState, PeerId,
+};
 use crate::sync::network_context::{LookupRequestResult, SyncNetworkContext};
 use beacon_chain::block_verification_types::RpcBlock;
 use beacon_chain::BeaconChainTypes;
 use lighthouse_network::service::api_types::Id;
 use std::sync::Arc;
 use types::blob_sidecar::FixedBlobSidecarList;
-use types::SignedBeaconBlock;
+use types::{DataColumnSidecarList, SignedBeaconBlock};
 
 use super::single_block_lookup::DownloadResult;
 use super::SingleLookupId;
@@ -17,6 +19,7 @@ use super::SingleLookupId;
 pub enum ResponseType {
     Block,
     Blob,
+    CustodyColumn,
 }
 
 /// This trait unifies common single block lookup functionality across blocks and blobs. This
@@ -38,7 +41,7 @@ pub trait RequestState<T: BeaconChainTypes> {
         &self,
         id: Id,
         peer_id: PeerId,
-        downloaded_block_expected_blobs: Option<usize>,
+        downloaded_block: Option<Arc<SignedBeaconBlock<T::EthSpec>>>,
         cx: &mut SyncNetworkContext<T>,
     ) -> Result<LookupRequestResult, LookupRequestError>;
 
@@ -73,7 +76,7 @@ impl<T: BeaconChainTypes> RequestState<T> for BlockRequestState<T::EthSpec> {
         &self,
         id: SingleLookupId,
         peer_id: PeerId,
-        _: Option<usize>,
+        _: Option<Arc<SignedBeaconBlock<T::EthSpec>>>,
         cx: &mut SyncNetworkContext<T>,
     ) -> Result<LookupRequestResult, LookupRequestError> {
         cx.block_lookup_request(id, peer_id, self.requested_block_root)
@@ -121,16 +124,11 @@ impl<T: BeaconChainTypes> RequestState<T> for BlobRequestState<T::EthSpec> {
         &self,
         id: Id,
         peer_id: PeerId,
-        downloaded_block_expected_blobs: Option<usize>,
+        downloaded_block: Option<Arc<SignedBeaconBlock<T::EthSpec>>>,
         cx: &mut SyncNetworkContext<T>,
     ) -> Result<LookupRequestResult, LookupRequestError> {
-        cx.blob_lookup_request(
-            id,
-            peer_id,
-            self.block_root,
-            downloaded_block_expected_blobs,
-        )
-        .map_err(LookupRequestError::SendFailedNetwork)
+        cx.blob_lookup_request(id, peer_id, self.block_root, downloaded_block)
+            .map_err(LookupRequestError::SendFailedNetwork)
     }
 
     fn send_for_processing(
@@ -161,3 +159,47 @@ impl<T: BeaconChainTypes> RequestState<T> for BlobRequestState<T::EthSpec> {
         &mut self.state
     }
 }
+
+impl<T: BeaconChainTypes> RequestState<T> for CustodyRequestState<T::EthSpec> {
+    type VerifiedResponseType = DataColumnSidecarList<T::EthSpec>;
+
+    fn make_request(
+        &self,
+        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>>>,
+        cx: &mut SyncNetworkContext<T>,
+    ) -> Result<LookupRequestResult, LookupRequestError> {
+        cx.custody_lookup_request(id, self.block_root, downloaded_block)
+            .map_err(LookupRequestError::SendFailedNetwork)
+    }
+
+    fn send_for_processing(
+        id: Id,
+        download_result: DownloadResult<Self::VerifiedResponseType>,
+        cx: &SyncNetworkContext<T>,
+    ) -> Result<(), LookupRequestError> {
+        let DownloadResult {
+            value,
+            block_root,
+            seen_timestamp,
+            ..
+        } = download_result;
+        cx.send_custody_columns_for_processing(id, block_root, value, seen_timestamp)
+            .map_err(LookupRequestError::SendFailedProcessor)
+    }
+
+    fn response_type() -> ResponseType {
+        ResponseType::CustodyColumn
+    }
+    fn request_state_mut(request: &mut SingleBlockLookup<T>) -> &mut Self {
+        &mut request.custody_request_state
+    }
+    fn get_state(&self) -> &SingleLookupRequestState<Self::VerifiedResponseType> {
+        &self.state
+    }
+    fn get_state_mut(&mut self) -> &mut SingleLookupRequestState<Self::VerifiedResponseType> {
+        &mut self.state
+    }
+}