Remove unwrap in NetworkController::poll

spectrum-finance · Nov 25, 2022 · 0e8dc27 · 0e8dc27
1 parent 212d180
commit 0e8dc27
Showing 1 changed file with 28 additions and 16 deletions.
diff --git a/spectrum-network/src/network_controller.rs b/spectrum-network/src/network_controller.rs
@@ -222,7 +222,9 @@ pub struct NetworkController<TPeers, TPeerManager, THandler> {
     peer_manager: TPeerManager,
     enabled_peers: HashMap<PeerId, ConnectedPeer<THandler>>,
     requests_recv: UnboundedReceiver<NetworkControllerIn>,
-    pending_actions: VecDeque<NetworkBehaviourAction<NetworkControllerOut, PartialPeerConnHandler>>,
+    pending_actions: VecDeque<
+        NetworkBehaviourAction<Result<NetworkControllerOut, NetworkControllerError>, PartialPeerConnHandler>,
+    >,
 }
 
 impl<TPeers, TPeerManager, THandler> NetworkController<TPeers, TPeerManager, THandler>
@@ -266,7 +268,7 @@ where
     THandler: ProtocolEvents + Clone + 'static,
 {
     type ConnectionHandler = PartialPeerConnHandler;
-    type OutEvent = NetworkControllerOut;
+    type OutEvent = Result<NetworkControllerOut, NetworkControllerError>;
 
     fn new_handler(&mut self) -> Self::ConnectionHandler {
         trace!("New handler is created");
@@ -562,21 +564,25 @@ where
                                 ConnectedPeer::Connected {
                                     enabled_protocols, ..
                                 } => {
-                                    let (_, prot_handler) = self.supported_protocols.get(&protocol).unwrap();
-                                    match enabled_protocols.entry(protocol) {
-                                        Entry::Occupied(_) => warn!(
-                                            "PM requested already enabled protocol {:?} with peer {:?}",
-                                            protocol, pid
-                                        ),
-                                        Entry::Vacant(protocol_entry) => {
-                                            protocol_entry.insert((
-                                                EnabledProtocol::PendingEnable,
-                                                prot_handler.clone(),
-                                            ));
-                                            prot_handler.protocol_requested_local(pid);
-                                            self.protocol_pending_enable(pid, protocol);
+                                    if let Some((_, prot_handler)) = self.supported_protocols.get(&protocol) {
+                                        match enabled_protocols.entry(protocol) {
+                                            Entry::Occupied(_) => warn!(
+                                                "PM requested already enabled protocol {:?} with peer {:?}",
+                                                protocol, pid
+                                            ),
+                                            Entry::Vacant(protocol_entry) => {
+                                                protocol_entry.insert((
+                                                    EnabledProtocol::PendingEnable,
+                                                    prot_handler.clone(),
+                                                ));
+                                                prot_handler.protocol_requested_local(pid);
+                                            }
                                         }
-                                    };
+                                    } else {
+                                        return Poll::Ready(NetworkBehaviourAction::GenerateEvent(Err(
+                                            NetworkControllerError::UnsupportedProtocol(protocol),
+                                        )));
+                                    }
                                 }
                                 ConnectedPeer::PendingConnect
                                 | ConnectedPeer::PendingApprove(_)
@@ -673,3 +679,9 @@ where
         }
     }
 }
+
+#[derive(Debug, thiserror::Error)]
+pub enum NetworkControllerError {
+    #[error("Unsupported protocol: {0:?}")]
+    UnsupportedProtocol(ProtocolId),
+}