better error handling + checksum added to the protocol

Cargo.toml

@@ -12,3 +12,5 @@ futures-util = "0.3.5"
 uuid = { version  = "0.8.1", features = ["v4"] }
 chrono = "0.4.11"
 ring = "0.16.13"
+anyhow = "1.0.30"
+thiserror = "1.0.17"

src/lib.rs

@@ -1,36 +1,67 @@
 //! # Isabelle's Lazy Message Protocol
+//!
+//! ### network packet protocol
+//!
+//! I don't know whether or not this is a super practical way of doing things
+//! but i'm lazy and it seems to work so gonna roll with it lol
+//!
+//! | segment size | usage                               |
+//! |--------------|-------------------------------------|
+//! | 1 byte       | u8 signifies the type of packet     |
+//! | 8 byte       | u64 length of the packet contents   |
+//! | 32 byte      | SHA256 packet contents checksum     |
+//! | `u64::MAX`   | packet contents                     |
+//!
+
 #![allow(dead_code)]
 
+use anyhow::Result;
 use futures_util::io::{AsyncReadExt, AsyncWriteExt};
+use ring::digest;
 use std::convert::TryInto;
 use std::marker::Unpin;
+use thiserror::Error;
 
 mod message;
 pub use message::Message;
 
-/// lazy error
-pub type Error = Box<dyn std::error::Error>;
-/// lazy result
-pub type Result<T> = std::result::Result<T, Error>;
-
 struct NetworkPacket(Vec<u8>);
 
-/// A type of data that can be sent
+/// a type of data that can be sent
 pub trait Sendable: Sized {
     fn to_packet(self) -> Result<Packet>;
     fn from_packet(packet: Packet) -> Result<Self>;
 }
 
-/// Data to be sent
+/// data to be sent
 pub struct Packet {
     kind: PacketKind,
+    checksum: Vec<u8>,
     contents: Vec<u8>,
 }
 
 impl Packet {
-    /// Create a new `Packet`
+    /// create a new `Packet`
+    ///
+    /// should not be used when receiving data over a stream
     pub fn new(kind: PacketKind, contents: Vec<u8>) -> Packet {
-        Packet { kind, contents }
+        let checksum = digest::digest(&digest::SHA256, &contents).as_ref().to_vec();
+        Packet {
+            kind,
+            checksum,
+            contents,
+        }
+    }
+
+    /// creates a new `Packet` from data received over a stream
+    ///
+    /// should not be used when creating a fresh `Packet`
+    pub fn from_network(kind: PacketKind, contents: Vec<u8>, checksum: Vec<u8>) -> Packet {
+        Packet {
+            kind,
+            contents,
+            checksum,
+        }
     }
 
     fn to_network_packet(self) -> NetworkPacket {
@@ -38,14 +69,60 @@ impl Packet {
 
         // write packet kind byte
         contents.push(self.kind as u8);
+
         // write the packet length
-        let contents_length = self.contents.len() as u32;
+        let contents_length = self.contents.len() as u64;
         contents.extend_from_slice(&contents_length.to_le_bytes());
+
+        // write checksum
+        contents.extend_from_slice(&self.checksum.as_ref());
+
         // write contents
         contents.extend_from_slice(&self.contents);
 
         NetworkPacket(contents)
     }
+
+    /// verifies SHA256 checksum
+    pub fn verify_integrity(&self) -> Result<()> {
+        let found = digest::digest(&digest::SHA256, &self.contents)
+            .as_ref()
+            .to_vec();
+        if found == self.checksum {
+            Ok(())
+        } else {
+            Err(IlmpError::BadChecksumIntegrity {
+                expected: self.checksum.clone(),
+                found,
+            }
+            .into())
+        }
+    }
+}
+
+/// kinds of packets that can be sent
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[repr(u8)]
+pub enum PacketKind {
+    Message = 0,
+    PublicKey = 1,
+}
+
+impl PacketKind {
+    /// returns `Option<PacketKind> given valid matching variant
+    pub fn from_u8(kind: u8) -> Option<PacketKind> {
+        match kind {
+            0 => Some(PacketKind::Message),
+            _ => None,
+        }
+    }
+}
+
+/// ilmp's error type
+#[derive(Error, Debug)]
+pub enum IlmpError {
+    #[error("checksum integrity check failed: (expected {expected:?} found {found:?})")]
+    BadChecksumIntegrity { expected: Vec<u8>, found: Vec<u8> },
 }
 
 /// reads a `Packet` from a stream
@@ -55,24 +132,27 @@ pub async fn read<S>(stream: &mut S) -> Result<Option<Packet>>
 where
     S: AsyncReadExt + Unpin,
 {
-    let mut info_buf = [0u8; 5];
+    let mut info_buf = [0u8; 9];
     let check = stream.read(&mut info_buf).await?;
     if check == 0 {
         return Ok(None);
     }
 
     let packet_kind = PacketKind::from_u8(info_buf[0]).unwrap();
-    let length = u32::from_le_bytes(info_buf[1..5].try_into().unwrap()) as usize;
+    let length = u32::from_le_bytes(info_buf[1..9].try_into().unwrap()) as usize;
+
+    let mut checksum: Vec<u8> = vec![0; 32];
+    stream.read(&mut checksum).await?;
 
     let mut contents: Vec<u8> = vec![0; length];
     stream.read(&mut contents).await?;
 
-    let packet = Packet::new(packet_kind, contents);
+    let packet = Packet::from_network(packet_kind, contents, checksum);
 
     Ok(Some(packet))
 }
 
-/// Writes a `Sendable` packet to a stream
+/// writes a `Sendable` packet to a stream
 pub async fn write<S, P>(stream: &mut S, packet: P) -> Result<()>
 where
     S: AsyncWriteExt + Unpin,
@@ -82,21 +162,3 @@ where
     stream.write(&network_packet.0).await?;
     Ok(())
 }
-
-/// Kinds of packets that can be sent
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-#[repr(u8)]
-pub enum PacketKind {
-    Message = 0,
-    PublicKey = 1,
-}
-
-impl PacketKind {
-    /// returns `Option<PacketKind> given valid matching variant
-    pub fn from_u8(kind: u8) -> Option<PacketKind> {
-        match kind {
-            0 => Some(PacketKind::Message),
-            _ => None,
-        }
-    }
-}

src/message.rs

@@ -1,5 +1,6 @@
 use crate::{Packet, PacketKind, Result};
 use chrono::prelude::*;
+use ring::digest;
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
@@ -30,9 +31,14 @@ impl Message {
 impl crate::Sendable for Message {
     fn to_packet(self) -> Result<Packet> {
         let contents: Vec<u8> = serde_json::to_string(&self)?.into_bytes();
+        let checksum = digest::digest(&digest::SHA256, &contents).as_ref().to_vec();
         let kind = PacketKind::Message;
 
-        Ok(Packet { kind, contents })
+        Ok(Packet {
+            kind,
+            checksum,
+            contents,
+        })
     }
     fn from_packet(packet: Packet) -> Result<Self> {
         let contents = &String::from_utf8(packet.contents)?;