Flesh out individual message format

Signed-off-by: Lee Smet <lee.smet@hotmail.com>

src/message.rs

@@ -46,6 +46,12 @@ const FLAG_MESSAGE_READ: u16 = 0b0000_1000_0000_0000;
 /// transmitted again by the sender.
 const FLAG_MESSAGE_ACK: u16 = 0b0000_0001_0000_0000;
 
+/// Length of a message checksum in bytes.
+const MESSAGE_CHECKSUM_LENGTH: usize = 32;
+
+/// A message checksum. In practice this is a 32 byte blake3 digest of the entire message.
+pub type Checksum = [u8; MESSAGE_CHECKSUM_LENGTH];
+
 pub struct MessageStack {
     data_plane: Arc<DataPlane>,
     inbox: Arc<Mutex<MessageInbox>>,
@@ -189,7 +195,7 @@ impl MessageId {
 }
 
 /// An owned [`PacketBuffer`] for working with messages.
-struct MessagePacket {
+pub struct MessagePacket {
     packet: PacketBuffer,
 }
 
@@ -221,9 +227,6 @@ impl MessagePacket {
     /// Get a mutable reference to the header of the `MessagePacket`.
     pub fn header_mut(&mut self) -> MessagePacketHeaderMut {
         MessagePacketHeaderMut {
-            // header: &mut self.packet.buffer_mut()[..MESSAGE_HEADER_SIZE]
-            //     .try_into()
-            //     .expect("Packet contains enough size for a header; qed"),
             header: <&mut [u8] as TryInto<&mut [u8; MESSAGE_HEADER_SIZE]>>::try_into(
                 &mut self.packet.buffer_mut()[..MESSAGE_HEADER_SIZE],
             )

src/message/chunk.rs

@@ -1 +1,224 @@
-pub struct MessageChunk {}
+use std::fmt;
+
+use super::MessagePacket;
+
+/// A message representing a "chunk" message.
+///
+/// The body of a chunk message has the following structure:
+///   - 8 bytes: chunk index
+///   - 8 bytes: chunk offset
+///   - 8 bytes: chunk size
+///   - remainder: chunk data of length based on field 3
+pub struct MessageChunk {
+    buffer: MessagePacket,
+}
+
+impl MessageChunk {
+    /// Create a new `MessageChunk` in the provided [`MessagePacket`].
+    pub fn new(buffer: MessagePacket) -> Self {
+        Self { buffer }
+    }
+
+    /// Return the index of the chunk in the message, as written in the body.
+    pub fn chunk_idx(&self) -> u64 {
+        u64::from_be_bytes(
+            self.buffer.buffer()[..8]
+                .try_into()
+                .expect("Buffer contains a size field of valid length; qed"),
+        )
+    }
+
+    /// Set the index of the chunk in the message body.
+    pub fn set_chunk_idx(&mut self, chunk_idx: u64) {
+        self.buffer.buffer_mut()[..8].copy_from_slice(&chunk_idx.to_be_bytes())
+    }
+
+    /// Return the chunk offset in the message, as written in the body.
+    pub fn chunk_offset(&self) -> u64 {
+        u64::from_be_bytes(
+            self.buffer.buffer()[8..16]
+                .try_into()
+                .expect("Buffer contains a size field of valid length; qed"),
+        )
+    }
+
+    /// Set the offset of the chunk in the message body.
+    pub fn set_chunk_offset(&mut self, chunk_offset: u64) {
+        self.buffer.buffer_mut()[8..16].copy_from_slice(&chunk_offset.to_be_bytes())
+    }
+
+    /// Return the size of the chunk in the message, as written in the body.
+    pub fn chunk_size(&self) -> u64 {
+        u64::from_be_bytes(
+            self.buffer.buffer()[16..24]
+                .try_into()
+                .expect("Buffer contains a size field of valid length; qed"),
+        )
+    }
+
+    /// Set the size of the chunk in the message body.
+    pub fn set_chunk_size(&mut self, chunk_size: u64) {
+        self.buffer.buffer_mut()[16..24].copy_from_slice(&chunk_size.to_be_bytes())
+    }
+
+    /// Return a reference to the chunk data in the message.
+    pub fn data(&self) -> &[u8] {
+        &self.buffer.buffer()[24..24 + self.chunk_size() as usize]
+    }
+
+    /// Set the chunk data in this message. This will also set the size field to the proper value.
+    pub fn set_chunk_data(&mut self, data: &[u8]) -> Result<(), InsufficientChunkSpace> {
+        let buf = self.buffer.buffer_mut();
+        let available_space = buf.len() - 24;
+
+        if data.len() > available_space {
+            return Err(InsufficientChunkSpace {
+                available: available_space,
+                needed: data.len(),
+            });
+        }
+
+        // Slicing based on data.len() is fine here as we just checked to make sure we can handle
+        // this capacity.
+        buf[24..24 + data.len()].copy_from_slice(data);
+
+        self.set_chunk_size(data.len() as u64);
+
+        Ok(())
+    }
+}
+
+/// An error indicating not enough space is availbe in a message to set the chunk data.
+#[derive(Debug)]
+pub struct InsufficientChunkSpace {
+    /// Amount of space available in the chunk.
+    pub available: usize,
+    /// Amount of space needed to set the chunk data
+    pub needed: usize,
+}
+
+impl fmt::Display for InsufficientChunkSpace {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "Insufficient capacity available, needed {} bytes, have {} bytes",
+            self.needed, self.available
+        )
+    }
+}
+
+impl std::error::Error for InsufficientChunkSpace {}
+
+#[cfg(test)]
+mod tests {
+    use std::array;
+
+    use crate::{crypto::PacketBuffer, message::MessagePacket};
+
+    use super::MessageChunk;
+
+    #[test]
+    fn read_chunk_idx() {
+        let mut pb = PacketBuffer::new();
+        pb.buffer_mut()[12..20].copy_from_slice(&[0, 0, 0, 0, 0, 0, 100, 73]);
+
+        let ms = MessageChunk::new(MessagePacket::new(pb));
+
+        assert_eq!(ms.chunk_idx(), 25_673);
+    }
+
+    #[test]
+    fn write_chunk_idx() {
+        let mut ms = MessageChunk::new(MessagePacket::new(PacketBuffer::new()));
+
+        ms.set_chunk_idx(723);
+
+        // Since we don't work with packet buffer we don't have to account for the message packet
+        // header.
+        assert_eq!(&ms.buffer.buffer()[..8], &[0, 0, 0, 0, 0, 0, 2, 211]);
+        assert_eq!(ms.chunk_idx(), 723);
+    }
+
+    #[test]
+    fn read_chunk_offset() {
+        let mut pb = PacketBuffer::new();
+        pb.buffer_mut()[20..28].copy_from_slice(&[0, 0, 0, 0, 0, 20, 40, 60]);
+
+        let ms = MessageChunk::new(MessagePacket::new(pb));
+
+        assert_eq!(ms.chunk_offset(), 1_321_020);
+    }
+
+    #[test]
+    fn write_chunk_offset() {
+        let mut ms = MessageChunk::new(MessagePacket::new(PacketBuffer::new()));
+
+        ms.set_chunk_offset(1_000_000);
+
+        // Since we don't work with packet buffer we don't have to account for the message packet
+        // header.
+        assert_eq!(&ms.buffer.buffer()[8..16], &[0, 0, 0, 0, 0, 15, 66, 64]);
+        assert_eq!(ms.chunk_offset(), 1_000_000);
+    }
+
+    #[test]
+    fn read_chunk_size() {
+        let mut pb = PacketBuffer::new();
+        pb.buffer_mut()[28..36].copy_from_slice(&[0, 0, 0, 0, 0, 0, 3, 232]);
+
+        let ms = MessageChunk::new(MessagePacket::new(pb));
+
+        assert_eq!(ms.chunk_size(), 1_000);
+    }
+
+    #[test]
+    fn write_chunk_size() {
+        let mut ms = MessageChunk::new(MessagePacket::new(PacketBuffer::new()));
+
+        ms.set_chunk_size(1_300);
+
+        // Since we don't work with packet buffer we don't have to account for the message packet
+        // header.
+        assert_eq!(&ms.buffer.buffer()[16..24], &[0, 0, 0, 0, 0, 0, 5, 20]);
+        assert_eq!(ms.chunk_size(), 1_300);
+    }
+
+    #[test]
+    fn read_chunk_data() {
+        const CHUNK_DATA: &[u8] = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
+        let mut pb = PacketBuffer::new();
+        // Set data len
+        pb.buffer_mut()[28..36].copy_from_slice(&CHUNK_DATA.len().to_be_bytes());
+        pb.buffer_mut()[36..36 + CHUNK_DATA.len()].copy_from_slice(CHUNK_DATA);
+
+        let ms = MessageChunk::new(MessagePacket::new(pb));
+
+        assert_eq!(ms.chunk_size(), 16);
+        assert_eq!(ms.data(), CHUNK_DATA);
+    }
+
+    #[test]
+    fn write_chunk_data() {
+        const CHUNK_DATA: &[u8] = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
+        let mut ms = MessageChunk::new(MessagePacket::new(PacketBuffer::new()));
+
+        let res = ms.set_chunk_data(CHUNK_DATA);
+        assert!(res.is_ok());
+
+        // Since we don't work with packet buffer we don't have to account for the message packet
+        // header.
+        // Check and make sure size is properly set.
+        assert_eq!(&ms.buffer.buffer()[16..24], &[0, 0, 0, 0, 0, 0, 0, 16]);
+        assert_eq!(ms.chunk_size(), 16);
+        assert_eq!(ms.data(), CHUNK_DATA);
+    }
+
+    #[test]
+    fn write_chunk_data_oversized() {
+        let data: [u8; 1500] = array::from_fn(|_| 0xFF);
+        let mut ms = MessageChunk::new(MessagePacket::new(PacketBuffer::new()));
+
+        let res = ms.set_chunk_data(&data);
+        assert!(res.is_err());
+    }
+}

src/message/done.rs

@@ -1 +1,107 @@
-pub struct MessageDone;
+use super::{Checksum, MessagePacket, MESSAGE_CHECKSUM_LENGTH};
+
+/// A message representing a "done" message.
+///
+/// The body of a done message has the following structure:
+///   - 8 bytes: chunks transmitted
+///   - 32 bytes: checksum of the transmitted data
+pub struct MessageDone {
+    buffer: MessagePacket,
+}
+
+impl MessageDone {
+    /// Create a new `MessageDone` in the provided [`MessagePacket`].
+    pub fn new(buffer: MessagePacket) -> Self {
+        Self { buffer }
+    }
+
+    /// Return the amount of chunks in the message, as written in the body.
+    pub fn chunk_count(&self) -> u64 {
+        u64::from_be_bytes(
+            self.buffer.buffer()[..8]
+                .try_into()
+                .expect("Buffer contains a size field of valid length; qed"),
+        )
+    }
+
+    /// Set the amount of chunks field of the message body.
+    pub fn set_chunk_count(&mut self, chunk_count: u64) {
+        self.buffer.buffer_mut()[..8].copy_from_slice(&chunk_count.to_be_bytes())
+    }
+
+    /// Get the checksum of the message from the body.
+    pub fn checksum(&self) -> Checksum {
+        self.buffer.buffer()[8..8 + MESSAGE_CHECKSUM_LENGTH]
+            .try_into()
+            .expect("Buffer contains enough data for a checksum; qed")
+    }
+
+    /// Set the checksum of the message in the body.
+    pub fn set_checksum(&mut self, checksum: Checksum) {
+        self.buffer.buffer_mut()[8..8 + MESSAGE_CHECKSUM_LENGTH].copy_from_slice(&checksum)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        crypto::PacketBuffer,
+        message::{Checksum, MessagePacket},
+    };
+
+    use super::MessageDone;
+
+    #[test]
+    fn read_chunk_count() {
+        let mut pb = PacketBuffer::new();
+        pb.buffer_mut()[12..20].copy_from_slice(&[0, 0, 0, 0, 0, 0, 73, 55]);
+
+        let ms = MessageDone::new(MessagePacket::new(pb));
+
+        assert_eq!(ms.chunk_count(), 18_743);
+    }
+
+    #[test]
+    fn write_chunk_count() {
+        let mut ms = MessageDone::new(MessagePacket::new(PacketBuffer::new()));
+
+        ms.set_chunk_count(10_000);
+
+        // Since we don't work with packet buffer we don't have to account for the message packet
+        // header.
+        assert_eq!(&ms.buffer.buffer()[..8], &[0, 0, 0, 0, 0, 0, 39, 16]);
+        assert_eq!(ms.chunk_count(), 10_000);
+    }
+
+    #[test]
+    fn read_checksum() {
+        const CHECKSUM: Checksum = [
+            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D,
+            0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B,
+            0x1C, 0x1D, 0x1E, 0x1F,
+        ];
+        let mut pb = PacketBuffer::new();
+        pb.buffer_mut()[20..52].copy_from_slice(&CHECKSUM);
+
+        let ms = MessageDone::new(MessagePacket::new(pb));
+
+        assert_eq!(ms.checksum(), CHECKSUM);
+    }
+
+    #[test]
+    fn write_checksum() {
+        const CHECKSUM: Checksum = [
+            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D,
+            0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B,
+            0x1C, 0x1D, 0x1E, 0x1F,
+        ];
+        let mut ms = MessageDone::new(MessagePacket::new(PacketBuffer::new()));
+
+        ms.set_checksum(CHECKSUM);
+
+        // Since we don't work with packet buffer we don't have to account for the message packet
+        // header.
+        assert_eq!(&ms.buffer.buffer()[8..40], &CHECKSUM);
+        assert_eq!(ms.checksum(), CHECKSUM);
+    }
+}

src/message/init.rs

@@ -14,7 +14,7 @@ impl MessageInit {
         Self { buffer }
     }
 
-    /// Return the length of the ,essage, as written in the body.
+    /// Return the length of the message, as written in the body.
     pub fn length(&self) -> u64 {
         u64::from_be_bytes(
             self.buffer.buffer()[..8]
@@ -23,8 +23,37 @@ impl MessageInit {
         )
     }
 
-    /// Set the length field of the message body
+    /// Set the length field of the message body.
     pub fn set_length(&mut self, length: u64) {
         self.buffer.buffer_mut()[..8].copy_from_slice(&length.to_be_bytes())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{crypto::PacketBuffer, message::MessagePacket};
+
+    use super::MessageInit;
+
+    #[test]
+    fn read_length() {
+        let mut pb = PacketBuffer::new();
+        pb.buffer_mut()[12..20].copy_from_slice(&[0, 0, 0, 0, 2, 3, 4, 5]);
+
+        let ms = MessageInit::new(MessagePacket::new(pb));
+
+        assert_eq!(ms.length(), 33_752_069);
+    }
+
+    #[test]
+    fn write_length() {
+        let mut ms = MessageInit::new(MessagePacket::new(PacketBuffer::new()));
+
+        ms.set_length(3_432_634_632);
+
+        // Since we don't work with packet buffer we don't have to account for the message packet
+        // header.
+        assert_eq!(&ms.buffer.buffer()[..8], &[0, 0, 0, 0, 204, 153, 217, 8]);
+        assert_eq!(ms.length(), 3_432_634_632);
+    }
+}