Implement message reassembly and checksum

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

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<u8>, try_duration: Duration) -> MessageId {
+    pub fn push_message(&self, dst: IpAddr, data: Vec<u8>, 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<u8>,
 }
@@ -717,6 +764,18 @@ pub struct OutboundMessageInfo {
     chunks: Vec<ChunkState>,
 }
 
+/// 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 {
 

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);
     }
 }