1 files changed, 68 insertions, 13 deletions

diff --git a/src/crypto.rs b/src/crypto.rs
index 726a420..07c6cef 100644
--- a/src/crypto.rs
+++ b/src/crypto.rs
@@ -1,6 +1,7 @@
 
 use std::{u8, u32};
 use std::io;
+use std::cmp::min;
 use std::io::{Read,Write, ErrorKind};
 use sodiumoxide::crypto::secretbox;
 use sodiumoxide::crypto::secretbox::{Key, Nonce};
@@ -10,10 +11,13 @@ use std::mem::transmute;
 // TODO: handle case of splitting up writes > 2^32 bytes into multiple small writes
 
 pub struct SecretStream<S: Read+Write> {
-    read_nonce: Nonce,
-    write_nonce: Nonce,
+    pub read_nonce: Nonce,
+    pub write_nonce: Nonce,
     pub key: Key,
     inner: S,
+    read_buf: [u8; 4096+1024],
+    read_buf_offset: usize,
+    read_buf_len: usize,
 }
 
 impl<S: Read+Write> SecretStream<S> {
@@ -23,43 +27,88 @@ impl<S: Read+Write> SecretStream<S> {
             read_nonce: secretbox::gen_nonce(),
             write_nonce: secretbox::gen_nonce(),
             key: secretbox::gen_key(),
+            read_buf: [0; 4096+1024],
+            read_buf_offset: 0,
+            read_buf_len: 0,
         }
     }
 }
 
 impl<S: Read+Write> Read for SecretStream<S> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+
+        // First try to return any extra older decrypted data
+        if self.read_buf_len > 0 {
+            println!("crypto: Returning existing data");
+            let rlen = min(self.read_buf_len, buf.len());
+            buf[..rlen].clone_from_slice(
+                &self.read_buf[self.read_buf_offset..(self.read_buf_offset+rlen)]);
+            self.read_buf_offset += rlen;
+            self.read_buf_len -= rlen;
+            return Ok(rlen);
+        }
+
         let mut header_buf = [0; 4];
         try!(self.inner.read_exact(&mut header_buf));
         let len: u32 = unsafe { transmute(header_buf) };
         let len = len.to_be();
-        if len as usize > buf.len() {
+        let len = len as usize;
+        if len as usize > self.read_buf.len() {
             return Err(io::Error::new(ErrorKind::Other,
-                format!("Buffer not big enough ({} < {})", buf.len(), len)));
+                format!("Message too big ({})", len)));
         }
-        try!(self.inner.read_exact(buf));
-        let cleartext = match secretbox::open(buf, &self.read_nonce, &self.key) {
+        try!(self.inner.read_exact(&mut self.read_buf[..len]));
+        println!("DECRYPT:");
+        println!("\tlen: {}", len);
+        println!("\tmsg: {:?}", &self.read_buf[..len]);
+        println!("\tnonce: {}", nonce2string(&self.write_nonce));
+        println!("\tkey: {}", key2string(&self.key));
+        let cleartext = match secretbox::open(&self.read_buf[..len], &self.read_nonce, &self.key) {
             Ok(cleartext) => cleartext,
             Err(_) => { return Err(io::Error::new(ErrorKind::InvalidData,
                 "Failed to decrypt message (could mean corruption or malicious attack"))},
         };
+        println!("crypto: Successfully decrypted message: {:?}", cleartext);
         self.read_nonce.increment_le_inplace();
-        let len = len as usize;
-        buf.clone_from_slice(&cleartext[..len]);
-        return Ok(len as usize);
+        let clen = cleartext.len() as usize;
+
+        // Do we have more data than we can return this type? If so buffer it
+        if clen > buf.len() {
+            let buf_len = buf.len();
+            buf.clone_from_slice(&cleartext[..buf_len]);
+            println!("copying extra: {} {} {}", self.read_buf.len(), buf_len, clen);
+            self.read_buf[..(clen-buf_len)].clone_from_slice(&cleartext[buf_len..]);
+            self.read_buf_offset = 0;
+            self.read_buf_len = clen - buf_len;
+            return Ok(buf_len);
+        } else {
+            buf.clone_from_slice(&cleartext[..clen]);
+            return Ok(clen as usize);
+        }
     }
 }
 
 impl<S: Read+Write> Write for SecretStream<S> {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
         assert!(buf.len() < u32::MAX as usize);
-        let len = buf.len() as u32;
+        let raw_len = buf.len() as u32;
+        let ciphertext = secretbox::seal(buf, &self.write_nonce, &self.key);
+
+        let len = ciphertext.len() as u32;
         let header_buf: [u8; 4] = unsafe { transmute(len.to_be()) };
         try!(self.inner.write_all(&header_buf));
-        let ciphertext = secretbox::seal(buf, &self.write_nonce, &self.key);
+
+        println!("DECRYPT:");
+        println!("\tlen: {}", len);
+        println!("\tmsg: {:?}", ciphertext);
+        println!("\tnonce: {}", nonce2string(&self.write_nonce));
+        println!("\tkey: {}", key2string(&self.key));
+        let check = secretbox::open(&ciphertext, &self.write_nonce, &self.key).unwrap();
+        //assert!(buf == check);
+
         self.write_nonce.increment_le_inplace();
         try!(self.inner.write_all(&ciphertext[..]));
-        return Ok(len as usize);
+        return Ok(raw_len as usize);
     }
 
     fn flush(&mut self) -> io::Result<()> {
@@ -72,7 +121,13 @@ pub fn key2string(key: &Key) -> String {
 }
 
 pub fn string2key(s: &str) -> Result<Key, String> {
-    println!("KEYBYTES: {}", secretbox::KEYBYTES);
     return Ok(Key::from_slice(&s.as_bytes().from_base64().unwrap()).unwrap());
 }
 
+pub fn nonce2string(nonce: &Nonce) -> String {
+    return (&(nonce[..])).to_base64(STANDARD);
+}
+
+pub fn string2nonce(s: &str) -> Result<Nonce, String> {
+    return Ok(Nonce::from_slice(&s.as_bytes().from_base64().unwrap()).unwrap());
+}