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 {