From de845283519e1a8a14da15aed1ec508b4483a6bf Mon Sep 17 00:00:00 2001 From: Lee Smet Date: Fri, 13 Oct 2023 17:02:45 +0200 Subject: [PATCH] Implement message handling Signed-off-by: Lee Smet --- src/message.rs | 191 +++++++++++++++++++++++++++++++++++++++---- src/message/chunk.rs | 12 ++- 2 files changed, 183 insertions(+), 20 deletions(-) diff --git a/src/message.rs b/src/message.rs index 5edfdff..78185df 100644 --- a/src/message.rs +++ b/src/message.rs @@ -8,7 +8,7 @@ use std::{ collections::HashMap, marker::PhantomData, - net::Ipv6Addr, + net::{IpAddr, Ipv6Addr}, ops::{Deref, DerefMut}, sync::{Arc, Mutex}, time::{self, Duration}, @@ -18,12 +18,23 @@ use futures::{Stream, StreamExt}; use log::{debug, error, warn}; use rand::Fill; -use crate::{crypto::PacketBuffer, data::DataPlane}; +use crate::{ + crypto::PacketBuffer, + data::DataPlane, + message::{chunk::MessageChunk, init::MessageInit}, +}; mod chunk; mod done; mod init; +/// The average size of a single chunk. This is mainly intended to preallocate the chunk array on +/// the receiver size. This value should allow reasonable overhead for standard MTU. +const AVERAGE_CHUNK_SIZE: usize = 1_300; +/// The minimum size of a data chunk. Chunks which have a size smaller than this are rejected. An +/// exception is made for the last chunk. +const MINIMUM_CHUNK_SIZE: u64 = 250; + /// The size in bytes of the message header which starts each user message packet. const MESSAGE_HEADER_SIZE: usize = 12; /// The size in bytes of a message ID. @@ -69,21 +80,26 @@ struct MessageOutbox { struct MessageInbox { /// Messages which are still being transmitted. + // TODO: MessageID is part of ReceivedMessageInfo, rework this into HashSet? pending_msges: HashMap, /// Messages which have been completed. + // TODO: MessageID is part of Message, rework this into HashSet? complete_msges: HashMap, } struct ReceivedMessageInfo { id: MessageId, - src: Ipv6Addr, - dst: Ipv6Addr, - message: Vec, + src: IpAddr, + dst: IpAddr, + /// Length of the finished message. + len: u64, + chunks: Vec>, } /// A chunk of a message. This represents individual data pieces on the receiver side. +#[derive(Clone)] struct Chunk { - chunk_idx: usize, + chunk_idx: u64, data: Vec, } @@ -154,7 +170,10 @@ impl MessageStack { /// [`Stream`]. pub fn new(data_plane: DataPlane, message_packet_stream: S) -> Self where - S: Stream> + Send + Unpin + 'static, + S: Stream> + + Send + + Unpin + + 'static, { let ms = Self { data_plane: Arc::new(Mutex::new(data_plane)), @@ -173,11 +192,14 @@ impl MessageStack { /// Handle incoming messages from the [`DataPlane`]. async fn handle_incoming_message_packets(self, mut message_packet_stream: S) where - S: Stream> + Send + Unpin + 'static, + S: Stream> + + Send + + Unpin + + 'static, { while let Some(maybe_message) = message_packet_stream.next().await { - let packet = match maybe_message { - Ok(packet) => packet, + let (packet, src, dst) = match maybe_message { + Ok((packet, src, dst)) => (packet, src, dst), Err(e) => { error!("Error reading message packet: {e}"); // TODO: is this fatal? @@ -190,7 +212,7 @@ impl MessageStack { if mp.header().flags().ack() { self.handle_message_reply(mp); } else { - self.handle_message(mp); + self.handle_message(mp, src, dst); } } @@ -199,9 +221,12 @@ impl MessageStack { /// Handle an incoming message packet which is a reply to a message we previously sent. fn handle_message_reply(&self, mp: MessagePacket) { - if mp.header().flags().init() { + let header = mp.header(); + let message_id = header.message_id(); + let flags = header.flags(); + if flags.init() { let mut outbox = self.outbox.lock().unwrap(); - if let Some(message) = outbox.msges.get_mut(&mp.header().message_id()) { + if let Some(message) = outbox.msges.get_mut(&message_id) { if message.state != TransmissionState::Init { debug!("Dropping INIT ACK for message not in init state"); return; @@ -209,13 +234,144 @@ impl MessageStack { message.state = TransmissionState::InProgress; todo!("start sending chunks"); } + } else if flags.chunk() { + // ACK for a chunk, mark chunk as received so it is not retried again. + let mut outbox = self.outbox.lock().unwrap(); + if let Some(message) = outbox.msges.get_mut(&message_id) { + if message.state != TransmissionState::InProgress { + debug!("Dropping CHUNK ACK for message not being transmitted"); + return; + } + let mc = MessageChunk::new(mp); + // Sanity checks. This is just to protect ourselves, if the other party is + // malicious it can return any data it wants here. + if mc.chunk_idx() > message.chunks.len() as u64 { + debug!("Dropping CHUNK ACK for message because ACK'ed chunk is out of bounds"); + return; + } + // Don't check data size. It is the repsonsiblity of the other party to ensure he + // ACKs the right chunk. Additionally a malicious node could return a crafted input + // here anyway. + + message.chunks[mc.chunk_idx() as usize].chunk_transmit_state = + ChunkTransmitState::Acked; + } + } else if flags.done() { + // ACK for full message. + let mut outbox = self.outbox.lock().unwrap(); + if let Some(message) = outbox.msges.get_mut(&message_id) { + if message.state != TransmissionState::InProgress { + debug!("Dropping DONE ACK for message which is not being transmitted"); + return; + } + message.state = TransmissionState::Received; + } + } else if flags.read() { + // Ack for a read flag. Since the original read flag is sent by the receiver, this + // means the sender indicates he has successfully received the notification that a + // userspace process has read the message. Note that read flags are only sent once, and + // this ack is only sent at most once, even if it gets lost. As a result, there is + // nothing to really do here, and this behavior (ACK READ) might be dropped in the + // future. + debug!("Received READ ACK"); + } else { + debug!("Received unknown ACK message flags {:x}", flags.flags); } - todo!(); } /// Handle an incoming message packet which is **not** a reply to a packet we previously sent. - fn handle_message(&self, mp: MessagePacket) { - todo!(); + fn handle_message(&self, mp: MessagePacket, src: IpAddr, dst: IpAddr) { + let header = mp.header(); + let message_id = header.message_id(); + let flags = header.flags(); + if flags.init() { + // We receive a new message with an ID. If we already have a complete message, ignore + // it. + let mut inbox = self.inbox.lock().unwrap(); + if inbox.complete_msges.contains_key(&message_id) { + debug!("Dropping INIT message as we already have a complete message with this ID"); + return; + } + // Otherwise unilaterally reset the state. The message id space is large enough to + // avoid accidental collisions. + let mi = MessageInit::new(mp); + let expected_chunks = + (mi.length() as usize + AVERAGE_CHUNK_SIZE - 1) / AVERAGE_CHUNK_SIZE; + let chunks = vec![None; expected_chunks]; + let message = ReceivedMessageInfo { + id: message_id, + src, + dst, + len: mi.length(), + chunks, + }; + + if inbox.pending_msges.insert(message_id, message).is_some() { + debug!("Dropped current pending message because we received a new message with INIT flag set for the same ID"); + } + } else if flags.chunk() { + // A chunk can only be received for incomplete messages. We don't have to check the + // completed messages. Either there is none, so no problem, or there is one, in which + // case we consider this to be a lingering chunk which was already accepted in the + // meantime (as the message is complete). + // + // SAFETY: a malcious node could send a lot of empty chunks, which trigger allocations + // to hold the chunk array, effectively exhausting memory. As such, we first need to + // determine if the chunk is feasible. + let mut inbox = self.inbox.lock().unwrap(); + if let Some(message) = inbox.pending_msges.get_mut(&message_id) { + let mc = MessageChunk::new(mp); + // Make sure the data is within bounds of the message being sent. + if message.len < mc.chunk_offset() + mc.chunk_size() { + debug!("Dropping invalid message CHUNK for being out of bounds"); + return; + } + // Check max chunk idx. + let max_chunk_idx = (message.len + MINIMUM_CHUNK_SIZE - 1) / MINIMUM_CHUNK_SIZE; + if mc.chunk_idx() > max_chunk_idx { + debug!("Dropping CHUNK because index is too high"); + return; + } + // Check chunk size, allow exception on last chunk. + if mc.chunk_size() < MINIMUM_CHUNK_SIZE && mc.chunk_idx() != max_chunk_idx { + debug!("Dropping CHUNK which is too small"); + return; + } + // Finally check if we have sufficient space for our chunks. + if message.chunks.len() as u64 <= mc.chunk_idx() { + // TODO: optimize + let chunks = + vec![None; (mc.chunk_idx() + 1 - message.chunks.len() as u64) as usize]; + message.chunks.extend_from_slice(&chunks); + } + // Now insert the chunk. Overwrite any previous chunk. + message.chunks[mc.chunk_idx() as usize] = Some(Chunk { + chunk_idx: mc.chunk_idx(), + data: mc.data().to_vec(), + }); + } + } else if flags.done() { + todo!("Checksum, reassemble and move message"); + } else if flags.read() { + let mut outbox = self.outbox.lock().unwrap(); + if let Some(message) = outbox.msges.get_mut(&message_id) { + if message.state != TransmissionState::Received { + debug!("Got READ for message which is not in received state"); + return; + } + message.state = TransmissionState::Read; + } + } else if flags.aborted() { + // If the message is not finished yet, discard it completely. + // But if it is finished, ignore this, i.e, nothing to do. + let mut inbox = self.inbox.lock().unwrap(); + if inbox.pending_msges.remove(&message_id).is_some() { + debug!("Dropping pending message because we received an ABORT"); + } + } else { + debug!("Received unknown ACK message flags {:x}", flags.flags); + } + todo!("reply"); } } @@ -241,6 +397,7 @@ impl MessageStack { state: TransmissionState::Init, created, deadline, + len: msg.data.len(), msg, chunks: vec![], // leave Vec empty at start }; @@ -552,6 +709,8 @@ pub struct OutboundMessageInfo { created: time::SystemTime, /// Timestamp indicating when we stop trying to send the message. deadline: time::SystemTime, + /// Length of the message. + len: usize, /// The message to send. msg: Message, /// Chunks of the message. diff --git a/src/message/chunk.rs b/src/message/chunk.rs index 74bef54..2e0af32 100644 --- a/src/message/chunk.rs +++ b/src/message/chunk.rs @@ -50,10 +50,14 @@ impl MessageChunk { /// 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"), + // Shield against a corrupt value. + u64::min( + u64::from_be_bytes( + self.buffer.buffer()[16..24] + .try_into() + .expect("Buffer contains a size field of valid length; qed"), + ), + self.buffer.buffer().len() as u64 - 24, ) }