diff --git a/mixnet/client/src/error.rs b/mixnet/client/src/error.rs index f1613e5f..caaae8a8 100644 --- a/mixnet/client/src/error.rs +++ b/mixnet/client/src/error.rs @@ -11,6 +11,8 @@ pub enum MixnetClientError { UnexpectedStreamBody, #[error("invalid payload")] InvalidPayload, + #[error("invalid fragment")] + InvalidFragment, #[error("invalid routing address: {0}")] InvalidRoutingAddress(#[from] NymNodeRoutingAddressError), #[error("{0}")] diff --git a/mixnet/client/src/receiver.rs b/mixnet/client/src/receiver.rs index 65ebd716..885d64ea 100644 --- a/mixnet/client/src/receiver.rs +++ b/mixnet/client/src/receiver.rs @@ -90,7 +90,7 @@ impl Receiver { match next { Ok(fragment) => { if let Some(message) = - Self::try_reconstruct_message(fragment, message_reconstructor) + Self::try_reconstruct_message(fragment, message_reconstructor)? { return Ok(message); } @@ -108,14 +108,14 @@ impl Receiver { fn try_reconstruct_message( fragment: Fragment, message_reconstructor: &mut MessageReconstructor, - ) -> Option> { + ) -> Result>> { let reconstruction_result = message_reconstructor.insert_new_fragment(fragment); match reconstruction_result { Some((padded_message, _)) => { - let message = Self::remove_padding(padded_message).unwrap(); - Some(message) + let message = Self::remove_padding(padded_message)?; + Ok(Some(message)) } - None => None, + None => Ok(None), } } @@ -126,7 +126,7 @@ impl Receiver { match padded_message.remove_padding(dummy_num_mix_hops)? { NymMessage::Plain(msg) => Ok(msg), - _ => todo!("return error"), + _ => Err(MixnetClientError::InvalidFragment), } } }