src/protocol.rs

   1 use std::error::Error;
   2 use std::io::Read;
   3
   4 const VERSION: u8 = 1;
   5
   6 #[derive(Debug, Clone)]
   7 pub struct DatagramEmptyError;
   8
   9 impl Error for DatagramEmptyError {}
  10
  11 impl std::fmt::Display for DatagramEmptyError {
  12     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
  13         write!(
  14             f,
  15             "datagram cannot be constructed when provided payload is empty"
  16         )
  17     }
  18 }
  19
  20 pub struct ConnectDatagram {
  21     version: u8,
  22     recipient: u16,
  23     data: Option<Vec<u8>>,
  24 }
  25
  26 impl ConnectDatagram {
  27     pub fn new(recipient: u16, data: Vec<u8>) -> Result<Self, DatagramEmptyError> {
  28         if data.len() > 0 {
  29             Ok(Self {
  30                 version: VERSION,
  31                 recipient,
  32                 data: Some(data),
  33             })
  34         } else {
  35             Err(DatagramEmptyError)
  36         }
  37     }
  38
  39     pub fn version(&self) -> u8 {
  40         self.version
  41     }
  42
  43     pub fn recipient(&self) -> u16 {
  44         self.recipient
  45     }
  46
  47     pub fn data(&self) -> Option<&Vec<u8>> {
  48         self.data.as_ref()
  49     }
  50
  51     pub fn take_data(&mut self) -> Option<Vec<u8>> {
  52         self.data.take()
  53     }
  54
  55     pub fn size(&self) -> usize {
  56         let data_len = if let Some(data) = self.data() {
  57             data.len()
  58         } else {
  59             0
  60         };
  61
  62         3 + data_len
  63     }
  64
  65     pub fn bytes(&self) -> Vec<u8> {
  66         let mut bytes = Vec::with_capacity(self.size());
  67
  68         bytes.extend(&self.version.to_be_bytes());
  69         bytes.extend(&self.recipient.to_be_bytes());
  70
  71         if let Some(data) = self.data() {
  72             bytes.extend(data.as_slice());
  73         }
  74
  75         return bytes;
  76     }
  77
  78     pub fn encode(&self) -> Vec<u8> {
  79         let size: u32 = (self.size()) as u32;
  80
  81         let mut bytes = Vec::from(size.to_be_bytes());
  82         bytes.extend(self.bytes());
  83
  84         return bytes;
  85     }
  86
  87     pub fn decode(source: &mut (dyn Read + Send + Sync)) -> anyhow::Result<Self> {
  88         // payload size
  89         let mut payload_size_bytes: [u8; 4] = [0; 4];
  90         source.read_exact(&mut payload_size_bytes)?;
  91         let payload_size = u32::from_be_bytes(payload_size_bytes);
  92
  93         // read whole payload
  94         let mut payload_bytes = vec![0; payload_size as usize];
  95         source.read_exact(payload_bytes.as_mut_slice())?;
  96
  97         // version
  98         let version_bytes = payload_bytes.remove(0);
  99         let version = u8::from_be(version_bytes);
 100
 101         // recipient
 102         let mut recipient_bytes: [u8; 2] = [0; 2];
 103         for i in 0..recipient_bytes.len() {
 104             recipient_bytes[i] = payload_bytes.remove(0);
 105         }
 106         let recipient = u16::from_be_bytes(recipient_bytes);
 107
 108         // data
 109         let data = payload_bytes;
 110
 111         if data.len() > 0 {
 112             Ok(Self {
 113                 version,
 114                 recipient,
 115                 data: Some(data),
 116             })
 117         } else {
 118             Err(anyhow::Error::from(DatagramEmptyError))
 119         }
 120     }
 121 }
 122
 123 #[cfg(test)]
 124 mod tests {
 125     use crate::protocol::ConnectDatagram;
 126     use std::io::Cursor;
 127
 128     #[test]
 129     fn encoded_size() -> anyhow::Result<()> {
 130         let mut data = Vec::new();
 131         for _ in 0..5 {
 132             data.push(1);
 133         }
 134         assert_eq!(5, data.len());
 135
 136         let sample = ConnectDatagram::new(1, data)?;
 137         assert_eq!(7 + 5, sample.encode().len());
 138
 139         Ok(())
 140     }
 141
 142     #[test]
 143     fn take_data() -> anyhow::Result<()> {
 144         let mut data = Vec::new();
 145         for _ in 0..5 {
 146             data.push(1);
 147         }
 148
 149         let mut sample = ConnectDatagram::new(1, data)?;
 150
 151         let taken_data = sample.take_data().unwrap();
 152         assert!(sample.data().is_none());
 153         assert_eq!(5, taken_data.len());
 154
 155         Ok(())
 156     }
 157
 158     #[async_std::test]
 159     async fn encode_and_decode() -> anyhow::Result<()> {
 160         let mut data = Vec::new();
 161         for _ in 0..5 {
 162             data.push(1);
 163         }
 164         assert_eq!(5, data.len());
 165
 166         let sample = ConnectDatagram::new(1, data)?;
 167
 168         let mut payload = sample.encode();
 169         assert_eq!(7 + 5, payload.len());
 170
 171         let mut cursor: Cursor<&mut [u8]> = Cursor::new(payload.as_mut());
 172         let sample_back_res = ConnectDatagram::decode(&mut cursor);
 173         assert!(sample_back_res.is_ok());
 174
 175         let sample_back = sample_back_res.unwrap();
 176         assert_eq!(sample_back.version(), 1);
 177         assert_eq!(sample_back.recipient(), 1);
 178         assert_eq!(sample_back.data().unwrap().len(), 5);
 179
 180         Ok(())
 181     }
 182 }