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 {