src/protocol.rs

   1 use std::array::TryFromSliceError;
   2 use std::convert::TryInto;
   3 use std::error::Error;
   4
   5 const VERSION: u16 = 1;
   6
   7 /// Encountered when there is an issue constructing, serializing, or deserializing a [`ConnectDatagram`].
   8 ///
   9 #[derive(Debug, Clone)]
  10 pub enum DatagramError {
  11     /// Tried to construct a [`ConnectDatagram`] with an empty message body.
  12     EmptyBody,
  13
  14     /// Did not provide the complete byte-string necessary to deserialize the [`ConnectDatagram`].
  15     IncompleteBytes,
  16
  17     BytesParseFail(TryFromSliceError),
  18 }
  19
  20 impl Error for DatagramError {}
  21
  22 impl std::fmt::Display for DatagramError {
  23     fn fmt(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
  24         match self {
  25             DatagramError::EmptyBody => formatter.write_str("tried to construct a `ConnectDatagram` with an empty message body"),
  26             DatagramError::IncompleteBytes => formatter.write_str("did not provide the complete byte-string necessary to deserialize the `ConnectDatagram`"),
  27             DatagramError::BytesParseFail(err) => std::fmt::Display::fmt(err, formatter),
  28         }
  29     }
  30 }
  31
  32 /// A simple size-prefixed packet format containing a version tag, recipient tag, and message body.
  33 ///
  34 #[derive(Clone)]
  35 pub struct ConnectDatagram {
  36     version: u16,
  37     recipient: u16,
  38     data: Option<Vec<u8>>,
  39 }
  40
  41 impl ConnectDatagram {
  42     /// Creates a new [`ConnectDatagram`] based on an intended recipient and message body.
  43     ///
  44     /// This will return a [EmptyBody](`DatagramError::EmptyBody`) error if the `data` parameter
  45     /// contains no bytes, or in other words, when there is no message body.
  46     ///
  47     /// The version field is decided by the library version and used to maintain backwards
  48     /// compatibility with previous datagram formats.
  49     ///
  50     pub fn new(recipient: u16, data: Vec<u8>) -> Result<Self, DatagramError> {
  51         if data.len() > 0 {
  52             Ok(Self {
  53                 version: VERSION,
  54                 recipient,
  55                 data: Some(data),
  56             })
  57         } else {
  58             Err(DatagramError::EmptyBody)
  59         }
  60     }
  61
  62     /// Gets the version number of the datagram.
  63     ///
  64     pub fn version(&self) -> u16 {
  65         self.version
  66     }
  67
  68     /// Gets the recipient of the datagram.
  69     ///
  70     pub fn recipient(&self) -> u16 {
  71         self.recipient
  72     }
  73
  74     /// Gets the message body of the datagram.
  75     ///
  76     pub fn data(&self) -> Option<&Vec<u8>> {
  77         self.data.as_ref()
  78     }
  79
  80     /// Takes ownership of the message body of the datagram.
  81     ///
  82     pub fn take_data(&mut self) -> Option<Vec<u8>> {
  83         self.data.take()
  84     }
  85
  86     /// Calculates the size-prefixed serialized byte-size of the datagram.
  87     ///
  88     /// This will include the byte-size of the size-prefix.
  89     ///
  90     pub fn size(&self) -> usize {
  91         let data_len = if let Some(data) = self.data() {
  92             data.len()
  93         } else {
  94             0
  95         };
  96
  97         8 + data_len
  98     }
  99
 100     /// Constructs a serialized representation of the datagram contents.
 101     ///
 102     pub(crate) fn bytes(&self) -> Vec<u8> {
 103         let mut bytes = Vec::with_capacity(self.size());
 104
 105         bytes.extend(&self.version.to_be_bytes());
 106         bytes.extend(&self.recipient.to_be_bytes());
 107
 108         if let Some(data) = self.data() {
 109             bytes.extend(data.as_slice());
 110         }
 111
 112         return bytes;
 113     }
 114
 115     /// Serializes the datagram.
 116     ///
 117     pub fn encode(self) -> Vec<u8> {
 118         let content_encoded = self.bytes();
 119         let size: u32 = (content_encoded.len()) as u32;
 120
 121         let mut bytes = Vec::from(size.to_be_bytes());
 122         bytes.extend(content_encoded);
 123
 124         return bytes;
 125     }
 126
 127     /// Deserializes the datagram from a buffer.
 128     ///
 129     /// The buffer **should not** contain the size-prefix, and only contain the byte contents of the
 130     /// struct (version, recipient, and message body).
 131     ///
 132     pub fn decode(mut buffer: Vec<u8>) -> Result<Self, DatagramError> {
 133         if buffer.len() > 4 {
 134             let mem_size = std::mem::size_of::<u16>();
 135             let data = buffer.split_off(mem_size * 2);
 136
 137             let (version_bytes, recipient_bytes) = buffer.split_at(mem_size);
 138
 139             match version_bytes.try_into() {
 140                 Ok(version_slice) => match recipient_bytes.try_into() {
 141                     Ok(recipient_slice) => {
 142                         let version = u16::from_be_bytes(version_slice);
 143                         let recipient = u16::from_be_bytes(recipient_slice);
 144
 145                         Ok(Self {
 146                             version,
 147                             recipient,
 148                             data: Some(data),
 149                         })
 150                     }
 151
 152                     Err(err) => Err(DatagramError::BytesParseFail(err)),
 153                 },
 154
 155                 Err(err) => Err(DatagramError::BytesParseFail(err)),
 156             }
 157         } else {
 158             Err(DatagramError::IncompleteBytes)
 159         }
 160     }
 161 }
 162
 163 #[cfg(test)]
 164 mod tests {
 165     use crate::protocol::ConnectDatagram;
 166
 167     #[test]
 168     fn serialized_size() -> anyhow::Result<()> {
 169         let mut data = Vec::new();
 170         for _ in 0..5 {
 171             data.push(1);
 172         }
 173         assert_eq!(5, data.len());
 174
 175         let sample = ConnectDatagram::new(1, data)?;
 176         assert_eq!(8 + 5, sample.encode().len());
 177
 178         Ok(())
 179     }
 180
 181     #[test]
 182     fn take_data() -> anyhow::Result<()> {
 183         let mut data = Vec::new();
 184         for _ in 0..5 {
 185             data.push(1);
 186         }
 187
 188         let mut sample = ConnectDatagram::new(1, data)?;
 189
 190         let taken_data = sample.take_data().unwrap();
 191         assert!(sample.data().is_none());
 192         assert_eq!(5, taken_data.len());
 193
 194         Ok(())
 195     }
 196
 197     #[async_std::test]
 198     async fn encode_and_decode() -> anyhow::Result<()> {
 199         let mut data = Vec::new();
 200         for _ in 0..5 {
 201             data.push(1);
 202         }
 203         assert_eq!(5, data.len());
 204
 205         let sample = ConnectDatagram::new(1, data)?;
 206         let serialized_size = sample.size();
 207         assert_eq!(8 + 5, serialized_size);
 208
 209         let mut payload = sample.encode();
 210         assert_eq!(serialized_size, payload.len());
 211
 212         let payload = payload.split_off(std::mem::size_of::<u32>());
 213         let sample_back_res = ConnectDatagram::decode(payload);
 214         assert!(sample_back_res.is_ok());
 215
 216         let sample_back = sample_back_res.unwrap();
 217         assert_eq!(sample_back.version(), 1);
 218         assert_eq!(sample_back.recipient(), 1);
 219         assert_eq!(sample_back.data().unwrap().len(), 5);
 220
 221         Ok(())
 222     }
 223 }