feat(code/network): Only enable the network behaviors which are actually required (#928)

* Only enable the network behaviours which are actually required

* Enable the sync channel in the Starknet app

* Enable the sync channel in the test app

Author: romac

code/crates/app/src/spawn.rs

@@ -200,5 +200,6 @@ fn make_gossip_config(cfg: &ConsensusConfig) -> NetworkConfig {
         },
         rpc_max_size: cfg.p2p.rpc_max_size.as_u64() as usize,
         pubsub_max_size: cfg.p2p.pubsub_max_size.as_u64() as usize,
+        enable_sync: true,
     }
 }

code/crates/config/src/lib.rs

@@ -379,10 +379,21 @@ pub enum VoteSyncMode {
     /// The lagging node sends a request to a peer for the missing votes
     #[default]
     RequestResponse,
+
     /// Nodes rebroadcast their last vote to all peers
     Rebroadcast,
 }
 
+impl VoteSyncMode {
+    pub fn is_request_response(&self) -> bool {
+        matches!(self, Self::RequestResponse)
+    }
+
+    pub fn is_rebroadcast(&self) -> bool {
+        matches!(self, Self::Rebroadcast)
+    }
+}
+
 /// Message types required by consensus to deliver the value being proposed
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "kebab-case")]

code/crates/network/src/behaviour.rs

@@ -2,6 +2,7 @@ use std::time::Duration;
 
 use libp2p::kad::{Addresses, KBucketKey, KBucketRef};
 use libp2p::request_response::{OutboundRequestId, ResponseChannel};
+use libp2p::swarm::behaviour::toggle::Toggle;
 use libp2p::swarm::NetworkBehaviour;
 use libp2p::{gossipsub, identify, ping};
 use libp2p_broadcast as broadcast;
@@ -66,10 +67,10 @@ impl From<discovery::NetworkEvent> for NetworkEvent {
 pub struct Behaviour {
     pub identify: identify::Behaviour,
     pub ping: ping::Behaviour,
-    pub gossipsub: gossipsub::Behaviour,
-    pub broadcast: broadcast::Behaviour,
-    pub sync: sync::Behaviour,
-    pub discovery: discovery::Behaviour,
+    pub gossipsub: Toggle<gossipsub::Behaviour>,
+    pub broadcast: Toggle<broadcast::Behaviour>,
+    pub sync: Toggle<sync::Behaviour>,
+    pub discovery: Toggle<discovery::Behaviour>,
 }
 
 /// Dummy implementation of Debug for Behaviour.
@@ -82,6 +83,8 @@ impl std::fmt::Debug for Behaviour {
 impl discovery::DiscoveryClient for Behaviour {
     fn add_address(&mut self, peer: &PeerId, address: Multiaddr) -> libp2p::kad::RoutingUpdate {
         self.discovery
+            .as_mut()
+            .expect("Discovery behaviour should be available")
             .kademlia
             .as_mut()
             .expect("Kademlia behaviour should be available")
@@ -90,22 +93,32 @@ impl discovery::DiscoveryClient for Behaviour {
 
     fn kbuckets(&mut self) -> impl Iterator<Item = KBucketRef<'_, KBucketKey<PeerId>, Addresses>> {
         self.discovery
+            .as_mut()
+            .expect("Discovery behaviour should be available")
             .kademlia
             .as_mut()
             .expect("Kademlia behaviour should be available")
             .kbuckets()
     }
 
     fn send_request(&mut self, peer_id: &PeerId, req: discovery::Request) -> OutboundRequestId {
-        self.discovery.request_response.send_request(peer_id, req)
+        self.discovery
+            .as_mut()
+            .expect("Discovery behaviour should be available")
+            .request_response
+            .send_request(peer_id, req)
     }
 
     fn send_response(
         &mut self,
         ch: ResponseChannel<discovery::Response>,
         rs: discovery::Response,
     ) -> Result<(), discovery::Response> {
-        self.discovery.request_response.send_response(ch, rs)
+        self.discovery
+            .as_mut()
+            .expect("Discovery behaviour should be available")
+            .request_response
+            .send_response(ch, rs)
     }
 }
 
@@ -144,35 +157,45 @@ impl Behaviour {
 
         let ping = ping::Behaviour::new(ping::Config::new().with_interval(Duration::from_secs(5)));
 
-        let gossipsub = gossipsub::Behaviour::new_with_metrics(
-            gossipsub::MessageAuthenticity::Signed(keypair.clone()),
-            gossipsub_config(config.gossipsub, config.pubsub_max_size),
-            registry.sub_registry_with_prefix("gossipsub"),
-            Default::default(),
-        )
-        .unwrap();
-
-        let broadcast = broadcast::Behaviour::new_with_metrics(
-            broadcast::Config {
-                max_buf_size: config.pubsub_max_size,
-            },
-            registry.sub_registry_with_prefix("broadcast"),
-        );
-
-        let sync = sync::Behaviour::new_with_metrics(
-            sync::Config::default().with_max_response_size(config.rpc_max_size),
-            registry.sub_registry_with_prefix("sync"),
-        );
-
-        let discovery = discovery::Behaviour::new(keypair, config.discovery);
+        let gossipsub = config.pubsub_protocol.is_gossipsub().then(|| {
+            gossipsub::Behaviour::new_with_metrics(
+                gossipsub::MessageAuthenticity::Signed(keypair.clone()),
+                gossipsub_config(config.gossipsub, config.pubsub_max_size),
+                registry.sub_registry_with_prefix("gossipsub"),
+                Default::default(),
+            )
+            .unwrap()
+        });
+
+        let enable_broadcast = config.pubsub_protocol.is_broadcast() || config.enable_sync;
+        let broadcast = enable_broadcast.then(|| {
+            broadcast::Behaviour::new_with_metrics(
+                broadcast::Config {
+                    max_buf_size: config.pubsub_max_size,
+                },
+                registry.sub_registry_with_prefix("broadcast"),
+            )
+        });
+
+        let sync = config.enable_sync.then(|| {
+            sync::Behaviour::new_with_metrics(
+                sync::Config::default().with_max_response_size(config.rpc_max_size),
+                registry.sub_registry_with_prefix("sync"),
+            )
+        });
+
+        let discovery = config
+            .discovery
+            .enabled
+            .then(|| discovery::Behaviour::new(keypair, config.discovery));
 
         Self {
             identify,
             ping,
-            gossipsub,
-            broadcast,
-            sync,
-            discovery,
+            sync: Toggle::from(sync),
+            gossipsub: Toggle::from(gossipsub),
+            broadcast: Toggle::from(broadcast),
+            discovery: Toggle::from(discovery),
         }
     }
 }

code/crates/network/src/lib.rs

@@ -94,6 +94,7 @@ pub struct Config {
     pub pubsub_protocol: PubSubProtocol,
     pub rpc_max_size: usize,
     pub pubsub_max_size: usize,
+    pub enable_sync: bool,
 }
 
 impl Config {
@@ -232,10 +233,12 @@ async fn run(
         return;
     };
 
-    if let Err(e) = pubsub::subscribe(&mut swarm, PubSubProtocol::Broadcast, &[Channel::Sync]) {
-        error!("Error subscribing to Sync channel: {e}");
-        return;
-    };
+    if config.enable_sync {
+        if let Err(e) = pubsub::subscribe(&mut swarm, PubSubProtocol::Broadcast, &[Channel::Sync]) {
+            error!("Error subscribing to Sync channel: {e}");
+            return;
+        };
+    }
 
     loop {
         let result = tokio::select! {
@@ -295,6 +298,11 @@ async fn handle_ctrl_msg(
         }
 
         CtrlMsg::Broadcast(channel, data) => {
+            if channel == Channel::Sync && !config.enable_sync {
+                trace!("Ignoring broadcast message to Sync channel: Sync not enabled");
+                return ControlFlow::Continue(());
+            }
+
             let msg_size = data.len();
             let result = pubsub::publish(swarm, PubSubProtocol::Broadcast, channel, data);
 
@@ -307,10 +315,12 @@ async fn handle_ctrl_msg(
         }
 
         CtrlMsg::SyncRequest(peer_id, request, reply_to) => {
-            let request_id = swarm
-                .behaviour_mut()
-                .sync
-                .send_request(peer_id.to_libp2p(), request);
+            let Some(sync) = swarm.behaviour_mut().sync.as_mut() else {
+                error!("Cannot request Sync from peer: Sync not enabled");
+                return ControlFlow::Continue(());
+            };
+
+            let request_id = sync.send_request(peer_id.to_libp2p(), request);
 
             if let Err(e) = reply_to.send(request_id) {
                 error!(%peer_id, "Error sending Sync request: {e}");
@@ -320,12 +330,17 @@ async fn handle_ctrl_msg(
         }
 
         CtrlMsg::SyncReply(request_id, data) => {
+            let Some(sync) = swarm.behaviour_mut().sync.as_mut() else {
+                error!("Cannot send Sync response to peer: Sync not enabled");
+                return ControlFlow::Continue(());
+            };
+
             let Some(channel) = state.sync_channels.remove(&request_id) else {
                 error!(%request_id, "Received Sync reply for unknown request ID");
                 return ControlFlow::Continue(());
             };
 
-            let result = swarm.behaviour_mut().sync.send_response(channel, data);
+            let result = sync.send_response(channel, data);
 
             match result {
                 Ok(()) => debug!(%request_id, "Replied to Sync request"),

code/crates/network/src/pubsub.rs

@@ -11,19 +11,21 @@ pub fn subscribe(
 ) -> Result<(), eyre::Report> {
     match protocol {
         PubSubProtocol::GossipSub => {
-            for channel in channels {
-                swarm
-                    .behaviour_mut()
-                    .gossipsub
-                    .subscribe(&channel.to_gossipsub_topic())?;
+            if let Some(gossipsub) = swarm.behaviour_mut().gossipsub.as_mut() {
+                for channel in channels {
+                    gossipsub.subscribe(&channel.to_gossipsub_topic())?;
+                }
+            } else {
+                return Err(eyre::eyre!("GossipSub not enabled"));
             }
         }
         PubSubProtocol::Broadcast => {
-            for channel in channels {
-                swarm
-                    .behaviour_mut()
-                    .broadcast
-                    .subscribe(channel.to_broadcast_topic());
+            if let Some(broadcast) = swarm.behaviour_mut().broadcast.as_mut() {
+                for channel in channels {
+                    broadcast.subscribe(channel.to_broadcast_topic());
+                }
+            } else {
+                return Err(eyre::eyre!("Broadcast not enabled"));
             }
         }
     }
@@ -39,16 +41,18 @@ pub fn publish(
 ) -> Result<(), eyre::Report> {
     match protocol {
         PubSubProtocol::GossipSub => {
-            swarm
-                .behaviour_mut()
-                .gossipsub
-                .publish(channel.to_gossipsub_topic(), data)?;
+            if let Some(gossipsub) = swarm.behaviour_mut().gossipsub.as_mut() {
+                gossipsub.publish(channel.to_gossipsub_topic(), data)?;
+            } else {
+                return Err(eyre::eyre!("GossipSub not enabled"));
+            }
         }
         PubSubProtocol::Broadcast => {
-            swarm
-                .behaviour_mut()
-                .broadcast
-                .broadcast(&channel.to_broadcast_topic(), data);
+            if let Some(broadcast) = swarm.behaviour_mut().broadcast.as_mut() {
+                broadcast.broadcast(&channel.to_broadcast_topic(), data);
+            } else {
+                return Err(eyre::eyre!("Broadcast not enabled"));
+            }
         }
     }
 

code/crates/network/test/src/lib.rs

@@ -162,6 +162,7 @@ impl<const N: usize> Test<N> {
             pubsub_protocol: malachitebft_network::PubSubProtocol::default(),
             rpc_max_size: 10 * 1024 * 1024,   // 10 MiB
             pubsub_max_size: 4 * 1024 * 1024, // 4 MiB
+            enable_sync: false,
         })
     }
 

code/crates/starknet/host/src/spawn.rs

@@ -247,6 +247,7 @@ async fn spawn_network_actor(
         },
         rpc_max_size: cfg.consensus.p2p.rpc_max_size.as_u64() as usize,
         pubsub_max_size: cfg.consensus.p2p.pubsub_max_size.as_u64() as usize,
+        enable_sync: true,
     };
 
     let keypair = make_keypair(private_key);