From 365be93bd4b6650a39d8198f622d283fffd03de8 Mon Sep 17 00:00:00 2001
From: Tim Ling <791016+kettlebell@users.noreply.github.com>
Date: Mon, 21 Nov 2022 01:22:36 +1100
Subject: [PATCH] Remove unwrap in `NetworkController`

---
 spectrum-network/src/network_controller.rs | 44 ++++++++++++++--------
 1 file changed, 28 insertions(+), 16 deletions(-)

diff --git a/spectrum-network/src/network_controller.rs b/spectrum-network/src/network_controller.rs
index 8b52d6bf..c66c5175 100644
--- 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),
+                            })
                     }
                 }
             }