WIP: partially working buffered crypto

3 files changed, 110 insertions, 19 deletions

diff --git a/src/client.rs b/src/client.rs
index 1458998..165a898 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -8,7 +8,7 @@ use std::process::exit;
 use std::process::Command;
 use getopts::Options;
 use utp::{UtpSocket, UtpStream};
-use crypto::{SecretStream, key2string, string2key};
+use crypto::{SecretStream, key2string, string2key, nonce2string, string2nonce};
 use sodiumoxide::crypto::secretbox;
 
 pub fn run_client(host: &str, local_file: &str, remote_file: &str, remote_is_dir: bool, is_recv: bool) {
@@ -38,22 +38,28 @@ pub fn run_client(host: &str, local_file: &str, remote_file: &str, remote_is_dir
     let reply = String::from_utf8_lossy(&ssh_output.stdout);
     println!("SSH reply: {}", reply);
     let words: Vec<&str> = reply.split_whitespace().collect();
-    if words.len() != 5 || words[0] != "UCP" || words[1] != "CONNECT" {
+    if words.len() != 7 || words[0] != "UCP" || words[1] != "CONNECT" {
         panic!("Unexpected data via SSH pipe (TCP)");
     }
     let remote_host = words[2];
     let remote_port = words[3].parse::<u16>().expect("failed to parse remote port number");
     let remote_secret = words[4];
+    let remote_read_nonce = words[5];
+    let remote_write_nonce = words[6];
 
     println!("Got remote details:");
     println!("\tport: {}", remote_port);
     println!("\thost: {}", remote_host);
-    println!("\tsecret: {}", remote_secret);
+    println!("\tsecret key: {}", remote_secret);
+    println!("\tsecret read key: {}", remote_read_nonce);
+    println!("\tsecret write key: {}", remote_write_nonce);
 
     let mut socket = UtpSocket::connect((remote_host, remote_port)).unwrap();;
     let mut stream: UtpStream = socket.into();
     let mut stream = SecretStream::new(stream);
     stream.key = string2key(remote_secret).unwrap();
+    stream.read_nonce = string2nonce(remote_write_nonce).unwrap();
+    stream.write_nonce = string2nonce(remote_read_nonce).unwrap();
 
     if is_recv {
         common::sink_files(&mut stream, local_file, remote_is_dir);
@@ -83,6 +89,9 @@ pub fn main_client() {
     opts.optopt("t", "to", "file or dir to write to (client side)", "FILE");
     opts.reqopt("", "host", "remote hostname to connect to", "HOSTNAME");
     opts.reqopt("", "port", "remote port to connect to", "PORT");
+    opts.reqopt("", "read-nonce", "secret read nonce", "NONCE");
+    opts.reqopt("", "write-nonce", "secret write nonce", "NONCE");
+    opts.reqopt("", "key", "secret key", "NONCE");
 
     assert!(args.len() >= 2 && args[1] == "client");
     let matches = match opts.parse(&args[2..]) {
@@ -110,11 +119,22 @@ pub fn main_client() {
     let mut socket = UtpSocket::connect((matches.opt_str("host").unwrap().as_str(), port)).unwrap();
     let mut stream: UtpStream = socket.into();
     println!("opened socket");
+
+    let mut stream = SecretStream::new(stream);
+    stream.key = string2key(&matches.opt_str("key").unwrap()).unwrap();
+    stream.read_nonce = string2nonce(&matches.opt_str("read-nonce").unwrap()).unwrap();
+    stream.write_nonce = string2nonce(&matches.opt_str("write-nonce").unwrap()).unwrap();
+
+    // XXX:
+    stream.read_nonce = secretbox::Nonce::from_slice(&[0; secretbox::NONCEBYTES]).unwrap();
+    stream.write_nonce = secretbox::Nonce::from_slice(&[0; secretbox::NONCEBYTES]).unwrap();
+
     if matches.opt_present("f") {
         common::source_files(&mut stream, &matches.opt_str("f").unwrap(), dir_mode);
     }
     if matches.opt_present("t") {
         common::sink_files(&mut stream, &matches.opt_str("t").unwrap(), dir_mode);
     }
-    stream.close().unwrap();
+    // XXX: does Drop do this well enough?
+    //stream.close().unwrap();
 }
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());
+}
diff --git a/src/server.rs b/src/server.rs
index 62d4b19..508bb36 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -9,7 +9,7 @@ use std::env::home_dir;
 use std::process::exit;
 use getopts::Options;
 use utp::{UtpSocket, UtpStream, UtpListener};
-use crypto::{SecretStream, key2string, string2key};
+use crypto::{SecretStream, key2string, string2key, nonce2string, string2nonce};
 use sodiumoxide::crypto::secretbox;
 
 fn run_server(path: &str, is_recv: bool, recursive: bool, daemonize: bool) {
@@ -27,9 +27,23 @@ fn run_server(path: &str, is_recv: bool, recursive: bool, daemonize: bool) {
     let listen_addr = listener.local_addr().unwrap().ip();
 
     let secret_key = secretbox::gen_key();
+    let read_nonce = secretbox::gen_nonce();
+    let write_nonce = secretbox::gen_nonce();
+
+    // XXX:
+    assert!(secret_key == string2key(&key2string(&secret_key)).unwrap());
+    assert!(read_nonce == string2nonce(&nonce2string(&read_nonce)).unwrap());
+    let read_nonce = secretbox::Nonce::from_slice(&[0; secretbox::NONCEBYTES]).unwrap();
+    let write_nonce = secretbox::Nonce::from_slice(&[0; secretbox::NONCEBYTES]).unwrap();
 
     // Send back details so client can connect
-    println!("UCP CONNECT {} {} {}", listen_addr, listen_port, key2string(&secret_key));
+    println!("UCP CONNECT {} {} {} {} {}",
+        listen_addr,
+        listen_port,
+        key2string(&secret_key),
+        nonce2string(&read_nonce),
+        nonce2string(&write_nonce));
+
 
     // TODO: maybe wait for an ACK of some sort here before daemonizing?
 
@@ -56,6 +70,8 @@ fn run_server(path: &str, is_recv: bool, recursive: bool, daemonize: bool) {
     let mut stream: UtpStream = socket.into();
     let mut stream = SecretStream::new(stream);
     stream.key = secret_key;
+    stream.read_nonce = read_nonce;
+    stream.write_nonce = write_nonce;
 
     if is_recv {
         common::sink_files(&mut stream, path, recursive);