diff --git a/sv2/codec-sv2/examples/encrypted.rs b/sv2/codec-sv2/examples/encrypted.rs index 58780d8527..37800f260e 100644 --- a/sv2/codec-sv2/examples/encrypted.rs +++ b/sv2/codec-sv2/examples/encrypted.rs @@ -228,9 +228,7 @@ fn main() { } // Parse the decoded frame header and payload - let decoded_frame_header = decoded_frame - .get_header() - .expect("Failed to get the frame header"); + let decoded_frame_header = decoded_frame.get_header(); let decoded_msg: CustomMessage = binary_sv2::from_bytes(decoded_frame.payload()) .expect("Failed to extract the message from the payload"); diff --git a/sv2/codec-sv2/examples/unencrypted.rs b/sv2/codec-sv2/examples/unencrypted.rs index f38033891f..d4aa1e40d1 100644 --- a/sv2/codec-sv2/examples/unencrypted.rs +++ b/sv2/codec-sv2/examples/unencrypted.rs @@ -80,9 +80,7 @@ fn main() { let mut decoded_frame = receiver_side(stream_receiver); // Parse the decoded frame header and payload - let decoded_frame_header = decoded_frame - .get_header() - .expect("Failed to get the frame header"); + let decoded_frame_header = decoded_frame.get_header(); let decoded_msg: CustomMessage = binary_sv2::from_bytes(decoded_frame.payload()) .expect("Failed to extract the message from the payload"); diff --git a/sv2/codec-sv2/src/decoder.rs b/sv2/codec-sv2/src/decoder.rs index bbc8656081..33fcd66853 100644 --- a/sv2/codec-sv2/src/decoder.rs +++ b/sv2/codec-sv2/src/decoder.rs @@ -515,7 +515,7 @@ mod prop_tests { None => return TestResult::discard(), }; - let expected_ext_type = frame.get_header().unwrap().ext_type(); + let expected_ext_type = frame.get_header().ext_type(); let mut encoder = Encoder::::new(); let encoded = match encoder.encode(frame) { @@ -526,10 +526,7 @@ mod prop_tests { let mut decoder = StandardDecoder::::new(); match decode_frame(&mut decoder, encoded.as_ref(), None) { Some(mut decoded_frame) => { - let header = match decoded_frame.get_header() { - Some(h) => h, - None => return TestResult::failed(), - }; + let header = decoded_frame.get_header(); let actual_msg_type = header.msg_type(); let actual_ext_type = header.ext_type(); let decoded_msg: TestMessage = match binary_sv2::from_bytes(decoded_frame.payload()) @@ -725,7 +722,7 @@ mod prop_tests { Some(f) => f, None => return TestResult::discard(), }; - let expected_ext = sv2_frame.get_header().unwrap().ext_type(); + let expected_ext = sv2_frame.get_header().ext_type(); let frame = Frame::Sv2(sv2_frame); let mut encoder = NoiseEncoder::::new(); @@ -738,10 +735,7 @@ mod prop_tests { let encrypted_bytes: &[u8] = encrypted.as_ref(); match decode_noise_frame(&mut decoder, &mut receiver_state, encrypted_bytes) { Some(mut decoded) => { - let header = match decoded.get_header() { - Some(h) => h, - None => return TestResult::failed(), - }; + let header = decoded.get_header(); let decoded_msg: TestMessage = match binary_sv2::from_bytes(decoded.payload()) { Ok(m) => m, Err(_) => return TestResult::failed(), diff --git a/sv2/framing-sv2/examples/sv2_frame.rs b/sv2/framing-sv2/examples/sv2_frame.rs index 333913e930..9d8ade08b7 100644 --- a/sv2/framing-sv2/examples/sv2_frame.rs +++ b/sv2/framing-sv2/examples/sv2_frame.rs @@ -51,9 +51,7 @@ fn main() { .expect("Failed to deserialize frame"); // Assert that deserialized header has the original content - let deserialized_header = deserialized_frame - .get_header() - .expect("Frame has no header"); + let deserialized_header = deserialized_frame.get_header(); assert_eq!(deserialized_header.msg_type(), MSG_TYPE); assert_eq!(deserialized_header.ext_type(), EXT_TYPE); diff --git a/sv2/framing-sv2/src/framing.rs b/sv2/framing-sv2/src/framing.rs index 5d7758bd6e..627500ba81 100644 --- a/sv2/framing-sv2/src/framing.rs +++ b/sv2/framing-sv2/src/framing.rs @@ -58,6 +58,18 @@ impl From> for Frame { } } +/// The body of a [`Sv2Frame`], which is either a deserialized payload or +/// serialized bytes. Exactly one variant is always present, eliminating the +/// need for dual-`Option` fields and impossible-state panic branches. +#[derive(Debug, Clone)] +pub enum FrameBody { + /// A deserialized message payload, not yet serialized for transmission. + Payload(T), + /// Serialized bytes (header + payload), received from the wire or + /// previously encoded. + Serialized(B), +} + /// Abstraction for a Sv2 frame. /// /// Represents a regular Sv2 frame, used for all communication outside of the Noise protocol @@ -66,9 +78,7 @@ impl From> for Frame { #[derive(Debug, Clone)] pub struct Sv2Frame { header: Header, - payload: Option, - // Serialized header + payload - serialized: Option, + body: FrameBody, } impl + AsRef<[u8]>> Sv2Frame { @@ -78,16 +88,16 @@ impl + AsRef<[u8]>> Sv2Frame { /// non serialized frame, it is not so cheap (because it serializes it). #[inline] pub fn serialize(self, dst: &mut [u8]) -> Result<(), Error> { - if let Some(mut serialized) = self.serialized { - dst.swap_with_slice(serialized.as_mut()); - Ok(()) - } else if let Some(payload) = self.payload { - to_writer(self.header, dst).map_err(Error::BinarySv2Error)?; - to_writer(payload, &mut dst[Header::SIZE..]).map_err(Error::BinarySv2Error)?; - Ok(()) - } else { - // Sv2Frame always has a payload or a serialized payload - panic!("Impossible state") + match self.body { + FrameBody::Serialized(mut serialized) => { + dst.swap_with_slice(serialized.as_mut()); + Ok(()) + } + FrameBody::Payload(payload) => { + to_writer(self.header, dst).map_err(Error::BinarySv2Error)?; + to_writer(payload, &mut dst[Header::SIZE..]).map_err(Error::BinarySv2Error)?; + Ok(()) + } } } @@ -100,17 +110,18 @@ impl + AsRef<[u8]>> Sv2Frame { /// payload. If the frame has not yet been serialized, this function should never be used (it /// will panic). pub fn payload(&mut self) -> &mut [u8] { - if let Some(serialized) = self.serialized.as_mut() { - &mut serialized.as_mut()[Header::SIZE..] - } else { - // panic here is the expected behaviour - panic!("Sv2Frame is not yet serialized.") + match &mut self.body { + FrameBody::Serialized(serialized) => &mut serialized.as_mut()[Header::SIZE..], + FrameBody::Payload(_) => { + // panic here is the expected behaviour + panic!("Sv2Frame is not yet serialized.") + } } } - /// [`Sv2Frame`] always returns `Some(self.header)`. - pub fn get_header(&self) -> Option { - Some(self.header) + /// [`Sv2Frame`] always has a header. + pub fn get_header(&self) -> crate::header::Header { + self.header } /// Tries to build a [`Sv2Frame`] from raw bytes. @@ -137,8 +148,7 @@ impl + AsRef<[u8]>> Sv2Frame { let header = Header::from_bytes(bytes.as_mut()).expect("Invalid header"); Self { header, - payload: None, - serialized: Some(bytes), + body: FrameBody::Serialized(bytes), } } @@ -172,13 +182,9 @@ impl + AsRef<[u8]>> Sv2Frame { /// the length of `self.payload`. #[inline] pub fn encoded_length(&self) -> usize { - if let Some(serialized) = self.serialized.as_ref() { - serialized.as_ref().len() - } else if let Some(payload) = self.payload.as_ref() { - payload.get_size() + Header::SIZE - } else { - // Sv2Frame always has a payload or a serialized payload - panic!("Impossible state") + match &self.body { + FrameBody::Serialized(serialized) => serialized.as_ref().len(), + FrameBody::Payload(payload) => payload.get_size() + Header::SIZE, } } @@ -195,8 +201,7 @@ impl + AsRef<[u8]>> Sv2Frame { let len = message.get_size() as u32; Header::from_len(len, message_type, extension_type).map(|header| Self { header, - payload: Some(message), - serialized: None, + body: FrameBody::Payload(message), }) } } @@ -205,14 +210,12 @@ impl Sv2Frame { /// Maps a `Sv2Frame` to `Sv2Frame` by applying `fun`, which is assumed to be a /// closure that converts `A` to `C` pub fn map(self, fun: fn(A) -> C) -> Sv2Frame { - let serialized = self.serialized; let header = self.header; - let payload = self.payload.map(fun); - Sv2Frame { - header, - payload, - serialized, - } + let body = match self.body { + FrameBody::Payload(p) => FrameBody::Payload(fun(p)), + FrameBody::Serialized(s) => FrameBody::Serialized(s), + }; + Sv2Frame { header, body } } } @@ -245,8 +248,8 @@ impl HandShakeFrame { /// Builds a [`HandShakeFrame`] from raw bytes. Nothing is assumed or checked about the /// correctness of the payload. - pub fn from_bytes(bytes: Slice) -> Result { - Ok(Self::from_bytes_unchecked(bytes)) + pub fn from_bytes(bytes: Slice) -> Self { + Self::from_bytes_unchecked(bytes) } #[inline]