From aa061cf30e9cd7c35666abab764b026b4082e9ff 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 51f1deef..ba21cd54 100644
--- a/spectrum-network/src/protocol/upgrade.rs
+++ b/spectrum-network/src/protocol/upgrade.rs
@@ -57,6 +57,8 @@ pub enum ProtocolHandshakeErr {
 pub enum ProtocolUpgradeErr {
     #[error(transparent)]
     HandshakeErr(#[from] ProtocolHandshakeErr),
+    #[error("Unsupported {0:?}")]
+    UnsupportedProtocolVer(ProtocolVer),
 }
 
 #[derive(Debug, Clone)]
@@ -129,10 +131,11 @@ where
 
     fn upgrade_inbound(self, mut socket: Substream, negotiated_tag: Self::Info) -> Self::Future {
         Box::pin(async move {
+            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 {
@@ -228,10 +231,11 @@ where
 
     fn upgrade_outbound(self, mut socket: Substream, negotiated_tag: Self::Info) -> Self::Future {
         Box::pin(async move {
+            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 {