Remove unwrap in NetworkController

spectrum-finance · Nov 21, 2022 · 30721df · 30721df
1 parent 081fd1c
commit 30721df
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
@@ -281,24 +281,36 @@ where
                 }) = self.enabled_peers.get_mut(&peer_id)
                 {
                     let protocol_id = protocol_tag.protocol_id();
-                    let (_, prot_handler) = self.supported_protocols.get(&protocol_id).unwrap();
-                    match enabled_protocols.entry(protocol_id) {
-                        Entry::Vacant(entry) => {
-                            entry.insert((EnabledProtocol::PendingApprove, prot_handler.clone()));
-                            prot_handler.protocol_requested(peer_id, protocol_tag.protocol_ver(), handshake);
-                        }
-                        Entry::Occupied(_) => {
-                            warn!(
-                                "Peer {:?} opened already enabled protocol {:?}",
-                                peer_id, protocol_id
-                            );
-                            self.pending_actions
-                                .push_back(NetworkBehaviourAction::NotifyHandler {
+                    if let Some((_, prot_handler)) = self.supported_protocols.get(&protocol_id) {
+                        match enabled_protocols.entry(protocol_id) {
+                            Entry::Vacant(entry) => {
+                                entry.insert((EnabledProtocol::PendingApprove, prot_handler.clone()));
+                                prot_handler.protocol_requested(
                                     peer_id,
-                                    handler: NotifyHandler::One(connection),
-                                    event: ConnHandlerIn::Close(protocol_id),
-                                })
+                                    protocol_tag.protocol_ver(),
+                                    handshake,
+                                );
+                            }
+                            Entry::Occupied(_) => {
+                                warn!(
+                                    "Peer {:?} opened already enabled protocol {:?}",
+                                    peer_id, protocol_id
+                                );
+                                self.pending_actions
+                                    .push_back(NetworkBehaviourAction::NotifyHandler {
+                                        peer_id,
+                                        handler: NotifyHandler::One(connection),
+                                        event: ConnHandlerIn::Close(protocol_id),
+                                    })
+                            }
                         }
+                    } else {
+                        self.pending_actions
+                            .push_back(NetworkBehaviourAction::NotifyHandler {
+                                peer_id,
+                                handler: NotifyHandler::One(connection),
+                                event: ConnHandlerIn::Close(protocol_id),
+                            })
                     }
                 }
             }