Mallory - Single commit

beacon_node/lighthouse_network/src/rpc/mod.rs

@@ -29,12 +29,14 @@ use self::protocol::RPCProtocol;
 use self::self_limiter::SelfRateLimiter;
 use crate::rpc::rate_limiter::RateLimiterItem;
 use crate::rpc::response_limiter::ResponseLimiter;
+use crate::MalloryConfig;
 pub use handler::SubstreamId;
 pub use methods::{
     BlocksByRangeRequest, BlocksByRootRequest, GoodbyeReason, LightClientBootstrapRequest,
     ResponseTermination, RpcErrorResponse, StatusMessage,
 };
-pub use protocol::{Protocol, RPCError};
+pub use methods::{RawMode, RawRequest};
+pub use protocol::{Protocol, RPCError, SupportedProtocol};
 
 pub(crate) mod codec;
 pub mod config;
@@ -149,6 +151,16 @@ pub struct NetworkParams {
     pub resp_timeout: Duration,
 }
 
+/// Additional configurations for the RPC Behaviour.
+#[derive(Clone, Copy)]
+pub struct MalloryLocalConfig {
+    /// Timeout in seconds for inbound connections.
+    pub inbound_timeout: u64,
+    /// Timeout for outbound connections.
+    pub outbound_timeout: u64,
+    pub self_handle_ping: bool,
+}
+
 /// Implements the libp2p `NetworkBehaviour` trait and therefore manages network-level
 /// logic.
 pub struct RPC<Id: ReqId, E: EthSpec> {
@@ -166,6 +178,9 @@ pub struct RPC<Id: ReqId, E: EthSpec> {
     network_params: NetworkParams,
     /// A sequential counter indicating when data gets modified.
     seq_number: u64,
+
+    /// Mallory Config
+    config: MalloryLocalConfig,
 }
 
 impl<Id: ReqId, E: EthSpec> RPC<Id, E> {
@@ -182,6 +197,7 @@ impl<Id: ReqId, E: EthSpec> RPC<Id, E> {
         outbound_rate_limiter_config: Option<OutboundRateLimiterConfig>,
         network_params: NetworkParams,
         seq_number: u64,
+        mallory_config: &MalloryConfig,
     ) -> Self {
         let response_limiter = inbound_rate_limiter_config.map(|config| {
             debug!(?config, "Using response rate limiting params");
@@ -193,6 +209,16 @@ impl<Id: ReqId, E: EthSpec> RPC<Id, E> {
             SelfRateLimiter::new(outbound_rate_limiter_config, fork_context.clone())
                 .expect("Outbound limiter configuration parameters are valid");
 
+        let mallory_config = MalloryLocalConfig {
+            inbound_timeout: mallory_config
+                .inbound_rpc_timeout
+                .unwrap_or(network_params.resp_timeout.as_secs()),
+            outbound_timeout: mallory_config
+                .outbound_rpc_timeout
+                .unwrap_or(network_params.resp_timeout.as_secs()),
+            self_handle_ping: mallory_config.user_handle_ping,
+        };
+
         RPC {
             response_limiter,
             outbound_request_limiter,
@@ -202,6 +228,7 @@ impl<Id: ReqId, E: EthSpec> RPC<Id, E> {
             enable_light_client_server,
             network_params,
             seq_number,
+            config: mallory_config,
         }
     }
 
@@ -344,6 +371,15 @@ impl<Id: ReqId, E: EthSpec> RPC<Id, E> {
         trace!(%peer_id, "Sending Ping");
         self.send_request(peer_id, id, RequestType::Ping(ping));
     }
+
+    /// Sends a pong response
+    pub fn pong(&mut self, peer_id: PeerId, inbound_request_id: InboundRequestId, data: u64) {
+        self.send_response(
+            inbound_request_id,
+            RpcResponse::Success(RpcSuccessResponse::Pong(Ping { data })),
+        )
+        .expect("request should exist");
+    }
 }
 
 impl<Id, E> NetworkBehaviour for RPC<Id, E>
@@ -378,6 +414,7 @@ where
             self.network_params.resp_timeout,
             peer_id,
             connection_id,
+            self.config.clone(),
         );
 
         Ok(handler)
@@ -408,6 +445,7 @@ where
             self.network_params.resp_timeout,
             peer_id,
             connection_id,
+            self.config,
         );
 
         Ok(handler)