Resolve merge conflicts

beacon_node/lighthouse_network/tests/rpc_tests.rs

@@ -5,8 +5,12 @@ use crate::common::spec_with_all_forks_enabled;
 use crate::common::{Protocol, build_tracing_subscriber};
 use bls::Signature;
 use fixed_bytes::FixedBytesExtended;
+use libp2p::PeerId;
 use lighthouse_network::rpc::{RequestType, methods::*};
-use lighthouse_network::service::api_types::AppRequestId;
+use lighthouse_network::service::api_types::{
+    AppRequestId, BlobsByRangeRequestId, BlocksByRangeRequestId, ComponentsByRangeRequestId,
+    DataColumnsByRangeRequestId, DataColumnsByRangeRequester, RangeRequestId, SyncRequestId,
+};
 use lighthouse_network::{NetworkEvent, ReportSource, Response};
 use ssz::Encode;
 use ssz_types::{RuntimeVariableList, VariableList};
@@ -1785,3 +1789,157 @@ fn test_active_requests() {
         }
     })
 }
+
+// Test that when a node receives an invalid BlocksByRange request exceeding the maximum count,
+// it bans the sender.
+#[test]
+fn test_request_too_large_blocks_by_range() {
+    let spec = Arc::new(spec_with_all_forks_enabled());
+
+    test_request_too_large(
+        AppRequestId::Sync(SyncRequestId::BlocksByRange(BlocksByRangeRequestId {
+            id: 1,
+            parent_request_id: ComponentsByRangeRequestId {
+                id: 1,
+                requester: RangeRequestId::RangeSync {
+                    chain_id: 1,
+                    batch_id: Epoch::new(1),
+                },
+            },
+        })),
+        RequestType::BlocksByRange(OldBlocksByRangeRequest::new(
+            0,
+            spec.max_request_blocks(ForkName::Base) as u64 + 1, // exceeds the max request defined in the spec.
+            1,
+        )),
+    );
+}
+
+// Test that when a node receives an invalid BlobsByRange request exceeding the maximum count,
+// it bans the sender.
+#[test]
+fn test_request_too_large_blobs_by_range() {
+    let spec = Arc::new(spec_with_all_forks_enabled());
+
+    let max_request_blobs_count = spec.max_request_blob_sidecars(ForkName::Base) as u64
+        / spec.max_blobs_per_block_within_fork(ForkName::Base);
+    test_request_too_large(
+        AppRequestId::Sync(SyncRequestId::BlobsByRange(BlobsByRangeRequestId {
+            id: 1,
+            parent_request_id: ComponentsByRangeRequestId {
+                id: 1,
+                requester: RangeRequestId::RangeSync {
+                    chain_id: 1,
+                    batch_id: Epoch::new(1),
+                },
+            },
+        })),
+        RequestType::BlobsByRange(BlobsByRangeRequest {
+            start_slot: 0,
+            count: max_request_blobs_count + 1, // exceeds the max request defined in the spec.
+        }),
+    );
+}
+
+// Test that when a node receives an invalid DataColumnsByRange request exceeding the columns count,
+// it bans the sender.
+#[test]
+fn test_request_too_large_data_columns_by_range() {
+    test_request_too_large(
+        AppRequestId::Sync(SyncRequestId::DataColumnsByRange(
+            DataColumnsByRangeRequestId {
+                id: 1,
+                parent_request_id: DataColumnsByRangeRequester::ComponentsByRange(
+                    ComponentsByRangeRequestId {
+                        id: 1,
+                        requester: RangeRequestId::RangeSync {
+                            chain_id: 1,
+                            batch_id: Epoch::new(1),
+                        },
+                    },
+                ),
+                peer: PeerId::random(),
+            },
+        )),
+        RequestType::DataColumnsByRange(DataColumnsByRangeRequest {
+            start_slot: 0,
+            count: 0,
+            // exceeds the max request defined in the spec.
+            columns: vec![0; E::number_of_columns() + 1],
+        }),
+    );
+}
+
+fn test_request_too_large(app_request_id: AppRequestId, request: RequestType<E>) {
+    // Set up the logging.
+    let log_level = "debug";
+    let enable_logging = true;
+    let _subscriber = build_tracing_subscriber(log_level, enable_logging);
+    let rt = Arc::new(Runtime::new().unwrap());
+    let spec = Arc::new(spec_with_all_forks_enabled());
+
+    rt.block_on(async {
+        let (mut sender, mut receiver) = common::build_node_pair(
+            Arc::downgrade(&rt),
+            ForkName::Base,
+            spec,
+            Protocol::Tcp,
+            false,
+            None,
+        )
+        .await;
+
+        // Build the sender future
+        let sender_future = async {
+            loop {
+                match sender.next_event().await {
+                    NetworkEvent::PeerConnectedOutgoing(peer_id) => {
+                        debug!(?request, %peer_id, "Sending RPC request");
+                        sender
+                            .send_request(peer_id, app_request_id, request.clone())
+                            .unwrap();
+                    }
+                    NetworkEvent::ResponseReceived {
+                        app_request_id,
+                        response,
+                        ..
+                    } => {
+                        debug!(?app_request_id, ?response, "Received response");
+                    }
+                    NetworkEvent::RPCFailed { error, .. } => {
+                        // This variant should be unreachable, as the receiver doesn't respond with an error when a request exceeds the limit.
+                        debug!(?error, "RPC failed");
+                        unreachable!();
+                    }
+                    NetworkEvent::PeerDisconnected(peer_id) => {
+                        // The receiver should disconnect as a result of the invalid request.
+                        debug!(%peer_id, "Peer disconnected");
+                        // End the test.
+                        return;
+                    }
+                    _ => {}
+                }
+            }
+        }
+        .instrument(info_span!("Sender"));
+
+        // Build the receiver future
+        let receiver_future = async {
+            loop {
+                if let NetworkEvent::RequestReceived { .. } = receiver.next_event().await {
+                    // This event should be unreachable, as the handler drops the invalid request.
+                    unreachable!();
+                }
+            }
+        }
+        .instrument(info_span!("Receiver"));
+
+        tokio::select! {
+            _ = sender_future => {}
+            _ = receiver_future => {}
+            _ = sleep(Duration::from_secs(30)) => {
+                panic!("Future timed out");
+            }
+        }
+    });
+}