Remove unwrap in NetworkController::poll

spectrum-finance · Nov 21, 2022 · 93bb2ca · 93bb2ca
1 parent d98e1c2
commit 93bb2ca
Showing 1 changed file with 28 additions and 15 deletions.
diff --git a/spectrum-network/src/network_controller.rs b/spectrum-network/src/network_controller.rs
@@ -114,7 +114,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>
@@ -158,7 +160,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");
@@ -443,20 +445,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);
+                                    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(_)
@@ -548,3 +555,9 @@ where
         }
     }
 }
+
+#[derive(Debug, thiserror::Error)]
+pub enum NetworkControllerError {
+    #[error("Unsupported protocol: {0:?}")]
+    UnsupportedProtocol(ProtocolId),
+}