From 82a73748b5302c1b72434758e4366410a1e3ccf7 Mon Sep 17 00:00:00 2001
From: Tim Ling <791016+kettlebell@users.noreply.github.com>
Date: Thu, 10 Nov 2022 00:29:36 +1100
Subject: [PATCH] Remove `unwrap` calls in `[In|Out]boundUpgrade`

---
 spectrum-network/src/protocol_upgrade.rs | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/spectrum-network/src/protocol_upgrade.rs b/spectrum-network/src/protocol_upgrade.rs
index df1757f0..807e81c2 100644
--- a/spectrum-network/src/protocol_upgrade.rs
+++ b/spectrum-network/src/protocol_upgrade.rs
@@ -33,6 +33,8 @@ pub enum ProtocolHandshakeErr {
 pub enum ProtocolUpgradeErr {
     #[error(transparent)]
     HandshakeErr(#[from] ProtocolHandshakeErr),
+    #[error("Unsupported {0:?}")]
+    UnsupportedProtocolVer(ProtocolVer),
 }
 
 #[derive(Debug, Clone)]
@@ -104,10 +106,11 @@ where
         Box::pin(async move {
             let target = format!("Inbound({})", negotiated_tag);
             trace!(target: &target, "upgrade_inbound()");
+            let protocol_ver = negotiated_tag.protocol_ver();
             let pspec = self
                 .supported_versions
-                .get(&negotiated_tag.protocol_ver())
-                .unwrap();
+                .get(&protocol_ver)
+                .ok_or(ProtocolUpgradeErr::UnsupportedProtocolVer(protocol_ver))?;
             let mut codec = UviBytes::default();
             codec.set_max_len(pspec.max_message_size);
             let handshake = if pspec.handshake_required {
@@ -211,10 +214,11 @@ where
         Box::pin(async move {
             let target = format!("Outbound({})", negotiated_tag);
             trace!(target: &target, "upgrade_outbound()");
+            let protocol_ver = negotiated_tag.protocol_ver();
             let pspec = self
                 .supported_versions
-                .get(&negotiated_tag.protocol_ver())
-                .unwrap();
+                .get(&protocol_ver)
+                .ok_or(ProtocolUpgradeErr::UnsupportedProtocolVer(protocol_ver))?;
             let mut codec = UviBytes::default();
             codec.set_max_len(pspec.max_message_size);
             if let Some(handshake) = &pspec.handshake {