From 74b80a09a928fbf034a02cbdf9ffd1472b4c3cef Mon Sep 17 00:00:00 2001 From: Lee Smet Date: Fri, 6 Oct 2023 17:27:09 +0200 Subject: [PATCH] Flesh out individual message format Signed-off-by: Lee Smet --- src/message.rs | 11 ++- src/message/chunk.rs | 225 ++++++++++++++++++++++++++++++++++++++++++- src/message/done.rs | 108 ++++++++++++++++++++- src/message/init.rs | 33 ++++++- 4 files changed, 369 insertions(+), 8 deletions(-) diff --git a/src/message.rs b/src/message.rs index 3caef3b..b379221 100644 --- a/src/message.rs +++ b/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, inbox: Arc>, @@ -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], ) diff --git a/src/message/chunk.rs b/src/message/chunk.rs index 7b8d3ee..5455d2f 100644 --- a/src/message/chunk.rs +++ b/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()); + } +} diff --git a/src/message/done.rs b/src/message/done.rs index f24ccea..9976ca8 100644 --- a/src/message/done.rs +++ b/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); + } +} diff --git a/src/message/init.rs b/src/message/init.rs index 2e1589d..a474b60 100644 --- a/src/message/init.rs +++ b/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); + } +}