From 7836378e64136a7aa827ff34f6f4cadbe16c57e5 Mon Sep 17 00:00:00 2001 From: jbesraa Date: Wed, 3 Jul 2024 18:14:28 +0300 Subject: [PATCH] Change `Sv2Frame` from struct to enum As `payload` and `serialised` can not be defined in the same time, `enum` is more appropriate here. --- protocols/v2/framing-sv2/src/framing.rs | 95 +++++++++++-------------- 1 file changed, 43 insertions(+), 52 deletions(-) diff --git a/protocols/v2/framing-sv2/src/framing.rs b/protocols/v2/framing-sv2/src/framing.rs index 616d533542..555b3004cd 100644 --- a/protocols/v2/framing-sv2/src/framing.rs +++ b/protocols/v2/framing-sv2/src/framing.rs @@ -40,11 +40,9 @@ impl From> for Frame { /// Abstraction for a SV2 Frame. #[derive(Debug, Clone)] -pub struct Sv2Frame { - header: Header, - payload: Option, - /// Serialized header + payload - serialized: Option, +pub enum Sv2Frame { + Raw(Header, B), + Payload(Header, T), } impl + AsRef<[u8]>> Sv2Frame { @@ -53,23 +51,23 @@ impl + AsRef<[u8]>> Sv2Frame { /// When called on a 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 { - #[cfg(not(feature = "with_serde"))] - to_writer(self.header, dst).map_err(Error::BinarySv2Error)?; - #[cfg(not(feature = "with_serde"))] - to_writer(payload, &mut dst[Header::SIZE..]).map_err(Error::BinarySv2Error)?; - #[cfg(feature = "with_serde")] - to_writer(&self.header, dst.as_mut()).map_err(Error::BinarySv2Error)?; - #[cfg(feature = "with_serde")] - to_writer(&payload, &mut dst.as_mut()[Header::SIZE..]) - .map_err(Error::BinarySv2Error)?; - Ok(()) - } else { - // Sv2Frame always has a payload or a serialized payload - panic!("Impossible state") + match self { + Sv2Frame::Raw(header, mut serialized) => { + dst.swap_with_slice(serialized.as_mut()); + Ok(()) + } + Sv2Frame::Payload(header, payload) => { + #[cfg(not(feature = "with_serde"))] + to_writer(header, dst).map_err(Error::BinarySv2Error)?; + #[cfg(not(feature = "with_serde"))] + to_writer(payload, &mut dst[Header::SIZE..]).map_err(Error::BinarySv2Error)?; + #[cfg(feature = "with_serde")] + to_writer(&header, dst.as_mut()).map_err(Error::BinarySv2Error)?; + #[cfg(feature = "with_serde")] + to_writer(&payload, &mut dst.as_mut()[Header::SIZE..]) + .map_err(Error::BinarySv2Error)?; + Ok(()) + } } } @@ -79,17 +77,18 @@ impl + AsRef<[u8]>> Sv2Frame { /// already serialized 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 self { + Sv2Frame::Raw(_, serialized) => &mut serialized.as_mut()[Header::SIZE..], + Sv2Frame::Payload(_, _payload) => panic!("fixme"), } } /// `Sv2Frame` always returns `Some(self.header)`. pub fn get_header(&self) -> Option { - Some(self.header) + match self { + Sv2Frame::Raw(header, _) => Some(*header), + Sv2Frame::Payload(header, _) => Some(*header), + } } /// Tries to build a `Sv2Frame` from raw bytes, assuming they represent a serialized `Sv2Frame` frame (`Self.serialized`). @@ -110,11 +109,7 @@ impl + AsRef<[u8]>> Sv2Frame { pub fn from_bytes_unchecked(mut bytes: B) -> Self { // Unchecked function caller is supposed to already know that the passed bytes are valid let header = Header::from_bytes(bytes.as_mut()).expect("Invalid header"); - Self { - header, - payload: None, - serialized: Some(bytes), - } + Sv2Frame::Raw(header, bytes) } /// After parsing `bytes` into a `Header`, this function helps to determine if the `msg_length` @@ -147,13 +142,9 @@ impl + AsRef<[u8]>> Sv2Frame { /// otherwise, returns 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 { + Sv2Frame::Raw(_, serialized) => serialized.as_ref().len(), + Sv2Frame::Payload(_, payload) => payload.get_size() + Header::SIZE, } } @@ -167,11 +158,15 @@ impl + AsRef<[u8]>> Sv2Frame { ) -> Option { let extension_type = update_extension_type(extension_type, channel_msg); let len = message.get_size() as u32; - Header::from_len(len, message_type, extension_type).map(|header| Self { - header, - payload: Some(message), - serialized: None, - }) + match len.checked_add(Header::SIZE as u32) { + Some(len) => { + if let Some(header) = Header::from_len(len, message_type, extension_type) { + return Some(Sv2Frame::Payload(header, message)); + } + return None; + } + None => return None, + } } } @@ -179,13 +174,9 @@ 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, + match self { + Sv2Frame::Raw(header, serialized) => Sv2Frame::Raw(header, serialized), + Sv2Frame::Payload(header, payload) => Sv2Frame::Payload(header, fun(payload)), } } }