refactor errors and test cli

src/lib.rs

@@ -1,4 +1,4 @@
-use crate::EncryptionError::IO;
+use crate::EncryptionError::{ShortData, IO};
 use chacha::Block;
 use diffie_hellman::{Public, Secret};
 use poly1305::oneoff_authenticate;
@@ -25,7 +25,14 @@ const KEY_DERIVATION_CONTEXT: &str =
 #[derive(Debug)]
 pub enum EncryptionError {
     InvalidMAC,
-    IO(Error),
+    IO {
+        err: Error,
+        at: &'static str,
+    },
+    ShortData {
+        got: usize,
+        expected: usize,
+    },
 }
 
 impl EncryptedStream {
@@ -75,7 +82,7 @@ impl EncryptedStream {
     pub fn read_packet(&mut self) -> Result<Vec<u8>, EncryptionError> {
         // the first bytes should be the nonce
         let mut nonce_bytes = [0u8; 12];
-        self.socket.read(&mut nonce_bytes).map_err(|err| IO(err))?;
+        self.socket.read(&mut nonce_bytes).map_err(|err| IO { err, at: "reading nonce" })?;
 
         self.reset(nonce_bytes);
 
@@ -83,15 +90,21 @@ impl EncryptedStream {
         self.block.advance();
 
         let mut mac = [0u8; 16];
-        self.socket.read(&mut mac).map_err(|err| IO(err))?;
+        self.socket.read(&mut mac).map_err(|err| IO { err, at: "reading mac" })?;
 
         let mut len_bytes = [0u8; 4];
-        self.socket.read(&mut len_bytes).map_err(|err| IO(err))?;
+        self.socket.read(&mut len_bytes).map_err(|err| IO { err, at: "reading data length" })?;
 
-        let len = u32::from_le_bytes(len_bytes);
+        let len = u32::from_le_bytes(len_bytes) as usize;
 
-        let mut buf = vec![0u8; len as usize];
-        self.read(&mut buf).map_err(|err| IO(err))?;
+        let mut buf = vec![0u8; len];
+        let read_bytes = self.read(&mut buf).map_err(|err| IO { err, at: "reading data" })?;
+        if read_bytes != len {
+            return Err(ShortData {
+                got: read_bytes,
+                expected: len,
+            })
+        }
 
         let my_mac = oneoff_authenticate(&buf, &poly_key[0..32].try_into().unwrap());
 

src/main.rs

@@ -12,11 +12,12 @@ fn main() {
 
     match args[0].as_ref() {
         "s" => {
-            let server = TcpListener::bind("127.0.0.1:2007").unwrap();
+            let server = TcpListener::bind(&*args[1]).unwrap();
             println!("listening");
 
             for incoming in server.incoming() {
                 println!("got incoming connection");
+
                 let mut encrypted = EncryptedStream::try_from(incoming.unwrap()).unwrap();
                 println!("successfully encrypted stream");
 
@@ -29,9 +30,9 @@ fn main() {
             }
         }
         "c" => {
-            let mut client = EncryptedStream::connect("127.0.0.1:2007").expect("Failed to connect");
+            let mut client = EncryptedStream::connect(&*args[1]).expect("Failed to connect");
 
-            let buf = client.read_packet().unwrap();
+            let buf = client.read_packet().expect("Failed to read packet");
 
             println!("success!");
             println!("data: {}", String::from_utf8_lossy(&buf));