Merge pull request meshtastic#3 from lukipuki/simplify-magic-buffer

Simplify the encoding of size

src/connections/stream_api.rs

@@ -278,7 +278,7 @@ impl<State> ConnectedStreamApi<State> {
     ///
     pub async fn send_raw(&mut self, data: EncodedToRadioPacket) -> Result<(), Error> {
         let channel = self.write_input_tx.clone();
-        let data_with_header = utils::format_data_packet(data);
+        let data_with_header = utils::format_data_packet(data)?;
 
         channel
             .send(data_with_header)

src/connections/stream_buffer.rs

@@ -182,8 +182,7 @@ impl StreamBuffer {
 
         // Combine MSB and LSB of incoming packet size bytes
         // Recall that packet size doesn't include the first four magic bytes
-        let shifted_msb: u32 = (*msb as u32).checked_shl(8).unwrap_or(0);
-        let incoming_packet_size: usize = 4 + shifted_msb as usize + (*lsb as usize);
+        let incoming_packet_size: usize = usize::from(4 + u16::from_le_bytes([*lsb, *msb]));
 
         // Defer decoding until the correct number of bytes are received
         if self.buffer.len() < incoming_packet_size {
@@ -223,11 +222,8 @@ impl StreamBuffer {
             });
         }
 
-        // Get index of next packet after removing current packet from buffer
-        let start_of_next_packet_idx: usize = 3 + (shifted_msb as usize) + ((*lsb) as usize) + 1;
-
         // Remove current packet from buffer based on start location of next packet
-        self.buffer = self.buffer[start_of_next_packet_idx..].to_vec();
+        self.buffer = self.buffer[incoming_packet_size..].to_vec();
 
         // Attempt to decode the current packet
         let decoded_packet = protobufs::FromRadio::decode(packet.as_slice())?;

src/errors_internal.rs

@@ -32,6 +32,10 @@ pub enum Error {
         description: String,
     },
 
+    /// An error indicating that too much data is being sent.
+    #[error("Trying to send too much data")]
+    InvalidaDataSize { data_length: usize },
+
     /// An error indicating that the method failed to remove a packet header from a packet buffer
     /// due to the packet buffer being too small to contain a header.
     #[error("Failed to remove packet header from packet buffer due to insufficient data length: {packet}")]

src/utils_internal.rs

@@ -268,14 +268,19 @@ where
 ///
 /// None
 ///
-pub fn format_data_packet(packet: EncodedToRadioPacket) -> EncodedToRadioPacketWithHeader {
+pub fn format_data_packet(
+    packet: EncodedToRadioPacket,
+) -> Result<EncodedToRadioPacketWithHeader, Error> {
     let data = packet.data();
-    let (msb, _) = data.len().overflowing_shr(8);
-    let lsb = (data.len() & 0xff) as u8;
-
-    let magic_buffer = [0x94, 0xc3, msb as u8, lsb];
+    if data.len() >= 1 << 16 {
+        return Err(Error::InvalidaDataSize {
+            data_length: data.len(),
+        });
+    }
+    let [lsb, msb, ..] = data.len().to_le_bytes();
+    let magic_buffer = [0x94, 0xc3, msb, lsb];
 
-    [&magic_buffer, data].concat().into()
+    Ok([&magic_buffer, data].concat().into())
 }
 
 /// A helper function that takes a vector of bytes (u8) representing an encoded packet with a 4-byte header,
@@ -382,7 +387,7 @@ mod tests {
         let data = vec![];
         let serial_data = format_data_packet(data.into());
 
-        assert_eq!(serial_data.data(), vec![0x94, 0xc3, 0x00, 0x00]);
+        assert_eq!(serial_data.data(), Some(vec![0x94, 0xc3, 0x00, 0x00]));
     }
 
     #[test]
@@ -392,7 +397,7 @@ mod tests {
 
         assert_eq!(
             serial_data.data(),
-            vec![0x94, 0xc3, 0x00, 0x03, 0x00, 0xff, 0x88]
+            Some(vec![0x94, 0xc3, 0x00, 0x03, 0x00, 0xff, 0x88])
         );
     }
 
@@ -401,6 +406,19 @@ mod tests {
         let data = vec![0x00; 0x100];
         let serial_data = format_data_packet(data.into());
 
-        assert_eq!(serial_data.data()[..4], vec![0x94, 0xc3, 0x01, 0x00]);
+        assert_eq!(serial_data.data()[..4], Some(vec![0x94, 0xc3, 0x01, 0x00]));
+    }
+
+    #[test]
+    fn invalid_too_large_packet() {
+        let data = vec![0x00; 0x10000];
+        let serial_data = format_data_packet(data.into());
+
+        assert_eq!(
+            serial_data,
+            Err(Error::InvalidaDataSize {
+                data_length: 0x10000
+            })
+        );
     }
 }