diff --git a/src/message.rs b/src/message.rs index 78185df..be87662 100644 --- a/src/message.rs +++ b/src/message.rs @@ -8,7 +8,7 @@ use std::{ collections::HashMap, marker::PhantomData, - net::{IpAddr, Ipv6Addr}, + net::IpAddr, ops::{Deref, DerefMut}, sync::{Arc, Mutex}, time::{self, Duration}, @@ -21,7 +21,7 @@ use rand::Fill; use crate::{ crypto::PacketBuffer, data::DataPlane, - message::{chunk::MessageChunk, init::MessageInit}, + message::{chunk::MessageChunk, done::MessageDone, init::MessageInit}, }; mod chunk; @@ -63,7 +63,6 @@ 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]; #[derive(Clone)] @@ -351,7 +350,54 @@ impl MessageStack { }); } } else if flags.done() { - todo!("Checksum, reassemble and move message"); + let mut inbox = self.inbox.lock().unwrap(); + let md = MessageDone::new(mp); + // At this point, we should have all message chunks. Verify length and reassemble them. + if let Some(inbound_message) = inbox.pending_msges.get_mut(&message_id) { + // Track total size of data we have allocated. + let mut chunk_size = 0; + let mut message_data = Vec::with_capacity(inbound_message.len as usize); + + // Chunks are inserted in order. + for chunk in &inbound_message.chunks { + if let Some(chunk) = chunk { + message_data.extend_from_slice(&chunk.data); + chunk_size += chunk.data.len(); + } else { + // A none chunk is not possible, we should have all chunks + debug!("DONE received for incomplete message"); + return; + } + } + + // TODO: report back here if there is an error. + if chunk_size as u64 != inbound_message.len { + debug!("Message has invalid size"); + return; + } + + let message = Message { + id: inbound_message.id, + src: inbound_message.src, + dst: inbound_message.dst, + data: message_data, + }; + + let checksum = message.checksum(); + + if checksum != md.checksum() { + debug!( + "Message has wrong checksum, got {} expected {}", + md.checksum().to_hex(), + checksum.to_hex() + ); + return; + } + + // Move message to be read. + inbox.complete_msges.insert(message_id, message); + inbox.pending_msges.remove(&message_id); + } } else if flags.read() { let mut outbox = self.outbox.lock().unwrap(); if let Some(message) = outbox.msges.get_mut(&message_id) { @@ -377,14 +423,15 @@ impl MessageStack { impl MessageStack { /// Push a new message to be transmitted, which will be tried for the given duration. - pub fn push_message(&self, dst: Ipv6Addr, data: Vec, try_duration: Duration) -> MessageId { + pub fn push_message(&self, dst: IpAddr, data: Vec, try_duration: Duration) -> MessageId { let src = self .data_plane .lock() .unwrap() .router() .node_public_key() - .address(); + .address() + .into(); let id = MessageId::new(); @@ -695,9 +742,9 @@ pub struct Message { /// Generated ID, used to identify the message on the wire id: MessageId, /// Source IP (ours) - src: Ipv6Addr, + src: IpAddr, /// Destination IP - dst: Ipv6Addr, + dst: IpAddr, /// Data data: Vec, } @@ -717,6 +764,18 @@ pub struct OutboundMessageInfo { chunks: Vec, } +/// A message checksum. In practice this is a 32 byte blake3 digest of the entire message. +pub type MessageChecksum = blake3::Hash; + +impl Message { + /// Calculates the [`MessageChechsum`] of the message. + /// + /// Currently this is a 32 byte blake3 hash. + pub fn checksum(&self) -> MessageChecksum { + blake3::hash(&self.data) + } +} + #[cfg(test)] mod tests { diff --git a/src/message/done.rs b/src/message/done.rs index cefcd7d..d135681 100644 --- a/src/message/done.rs +++ b/src/message/done.rs @@ -1,4 +1,4 @@ -use super::{Checksum, MessagePacket, MESSAGE_CHECKSUM_LENGTH}; +use super::{MessageChecksum, MessagePacket, MESSAGE_CHECKSUM_LENGTH}; /// A message representing a "done" message. /// @@ -31,15 +31,18 @@ impl MessageDone { } /// 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") + pub fn checksum(&self) -> MessageChecksum { + MessageChecksum::from_bytes( + 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) + pub fn set_checksum(&mut self, checksum: MessageChecksum) { + self.buffer.buffer_mut()[8..8 + MESSAGE_CHECKSUM_LENGTH] + .copy_from_slice(checksum.as_bytes()) } /// Consumes this `MessageDone`, returning the underlying [`MessagePacket`]. @@ -52,7 +55,7 @@ impl MessageDone { mod tests { use crate::{ crypto::PacketBuffer, - message::{Checksum, MessagePacket}, + message::{MessageChecksum, MessagePacket}, }; use super::MessageDone; @@ -89,13 +92,13 @@ mod tests { #[test] fn read_checksum() { - const CHECKSUM: Checksum = [ + const CHECKSUM: MessageChecksum = MessageChecksum::from_bytes([ 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); + pb.buffer_mut()[20..52].copy_from_slice(CHECKSUM.as_bytes()); let ms = MessageDone::new(MessagePacket::new(pb)); @@ -104,18 +107,18 @@ mod tests { #[test] fn write_checksum() { - const CHECKSUM: Checksum = [ + const CHECKSUM: MessageChecksum = MessageChecksum::from_bytes([ 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.buffer.buffer()[8..40], CHECKSUM.as_bytes()); assert_eq!(ms.checksum(), CHECKSUM); } }