Remove unwrap calls in [In|Out]boundUpgrade

spectrum-finance · Nov 9, 2022 · aa061cf · aa061cf
1 parent 3270fab
commit aa061cf
Showing 1 changed file with 8 additions and 4 deletions.
diff --git 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 {