summary history branches tags files
commit:9766787193cef7d74d5cc928448a2566ad29d0bd
author:Trevor Bentley
committer:Trevor Bentley
date:Mon Jan 28 00:02:17 2019 +0100
parents:7002994ea84e5f78be64d855029ebd3a3c933cd7
Move comm functions into ConnectionContext impl
diff --git a/benches/basic.rs b/benches/basic.rs
line changes: +12/-16
index e66fb52..7e9191b
--- a/benches/basic.rs
+++ b/benches/basic.rs
@@ -5,8 +5,6 @@ use std::thread;
 use std::net::{TcpListener, TcpStream};
 
 use ossuary::{ConnectionContext, ConnectionType};
-use ossuary::{crypto_send_handshake,crypto_recv_handshake, crypto_handshake_done};
-use ossuary::{crypto_send_data,crypto_recv_data,crypto_flush};
 use ossuary::OssuaryError;
 //use crate::*;
 
@@ -16,10 +14,10 @@ fn bench_test(b: &mut Bencher) {
         let listener = TcpListener::bind("127.0.0.1:9987").unwrap();
         let mut server_stream = listener.incoming().next().unwrap().unwrap();
         let mut server_conn = ConnectionContext::new(ConnectionType::UnauthenticatedServer);
-        while crypto_handshake_done(&server_conn).unwrap() == false {
-            if crypto_send_handshake(&mut server_conn, &mut server_stream).is_ok() {
+        while server_conn.crypto_handshake_done().unwrap() == false {
+            if server_conn.crypto_send_handshake(&mut server_stream).is_ok() {
                 loop {
-                    match crypto_recv_handshake(&mut server_conn, &mut server_stream) {
+                    match server_conn.crypto_recv_handshake(&mut server_stream) {
                         Ok(_) => break,
                         Err(OssuaryError::WouldBlock(_)) => {},
                         _ => panic!("Handshake failed"),
@@ -33,9 +31,8 @@ fn bench_test(b: &mut Bencher) {
         let start = std::time::SystemTime::now();
         loop {
             //std::thread::sleep(std::time::Duration::from_millis(100));
-            match crypto_recv_data(&mut server_conn,
-                                   &mut server_stream,
-                                   &mut plaintext) {
+            match server_conn.crypto_recv_data(&mut server_stream,
+                                               &mut plaintext) {
                 Ok((read, _written)) => bytes += read as u64,
                 Err(OssuaryError::WouldBlock(_)) => continue,
                 Err(e) => {
@@ -61,10 +58,10 @@ fn bench_test(b: &mut Bencher) {
     let mut client_stream = TcpStream::connect("127.0.0.1:9987").unwrap();
     client_stream.set_nonblocking(true).unwrap();
     let mut client_conn = ConnectionContext::new(ConnectionType::Client);
-    while crypto_handshake_done(&client_conn).unwrap() == false {
-        if crypto_send_handshake(&mut client_conn, &mut client_stream).is_ok() {
+    while client_conn.crypto_handshake_done().unwrap() == false {
+        if client_conn.crypto_send_handshake(&mut client_stream).is_ok() {
             loop {
-                match crypto_recv_handshake(&mut client_conn, &mut client_stream) {
+                match client_conn.crypto_recv_handshake(&mut client_stream) {
                     Ok(_) => break,
                     Err(OssuaryError::WouldBlock(_)) => {},
                     Err(e) => {
@@ -81,9 +78,8 @@ fn bench_test(b: &mut Bencher) {
     let start = std::time::SystemTime::now();
     let mut plaintext: &[u8] = &[0xaa; 16384];
     b.iter(|| {
-        match crypto_send_data(&mut client_conn,
-                               &mut plaintext,
-                               &mut client_stream) {
+        match client_conn.crypto_send_data(&mut plaintext,
+                                           &mut client_stream) {
             Ok(b) => bytes += b as u64,
             Err(OssuaryError::WouldBlock(_)) => {},
             _ => panic!("send error"),
@@ -97,7 +93,7 @@ fn bench_test(b: &mut Bencher) {
     }
     let mut plaintext: &[u8] = &[0xde, 0xde, 0xbe, 0xbe];
     loop {
-        match crypto_send_data(&mut client_conn, &mut plaintext, &mut client_stream) {
+        match client_conn.crypto_send_data(&mut plaintext, &mut client_stream) {
             Ok(w) => {
                 println!("wrote finish: {}", w);
                 break;
@@ -107,7 +103,7 @@ fn bench_test(b: &mut Bencher) {
         }
     }
 
-    while let Ok(w) = crypto_flush(&mut client_conn, &mut client_stream) {
+    while let Ok(w) = client_conn.crypto_flush(&mut client_stream) {
         if w == 0 {
             break;
         }

diff --git a/src/clib.rs b/src/clib.rs
line changes: +12/-14
index 6702e30..e4dee8c
--- a/src/clib.rs
+++ b/src/clib.rs
@@ -1,6 +1,4 @@
-use crate::{crypto_send_data, crypto_recv_data, crypto_flush,
-            crypto_send_handshake, crypto_recv_handshake, crypto_handshake_done,
-            ConnectionContext, ConnectionType, OssuaryError};
+use crate::{ConnectionContext, ConnectionType, OssuaryError};
 
 const ERROR_WOULD_BLOCK: i32 = -64;
 
@@ -74,11 +72,11 @@ pub extern "C" fn ossuary_recv_handshake(conn: *mut ConnectionContext,
     if conn.is_null() || in_buf.is_null() || in_buf_len.is_null() {
         return -1i32;
     }
-    let mut conn = unsafe { &mut *conn };
+    let conn = unsafe { &mut *conn };
     let inlen = unsafe { *in_buf_len as usize };
     let r_in_buf: &[u8] = unsafe { std::slice::from_raw_parts(in_buf, inlen) };
     let mut slice = r_in_buf;
-    let read: i32 = match crypto_recv_handshake(&mut conn, &mut slice) {
+    let read: i32 = match conn.crypto_recv_handshake(&mut slice) {
         Ok(read) => {
             unsafe { *in_buf_len = read as u16; }
             read as i32
@@ -99,11 +97,11 @@ pub extern "C" fn ossuary_send_handshake(conn: *mut ConnectionContext,
     if conn.is_null() || out_buf.is_null() || out_buf_len.is_null() {
         return -1i32;
     }
-    let mut conn = unsafe { &mut *conn };
+    let conn = unsafe { &mut *conn };
     let outlen = unsafe { *out_buf_len as usize };
     let r_out_buf: &mut [u8] = unsafe { std::slice::from_raw_parts_mut(out_buf, outlen) };
     let mut slice = r_out_buf;
-    let wrote: i32 = match crypto_send_handshake(&mut conn, &mut slice) {
+    let wrote: i32 = match conn.crypto_send_handshake(&mut slice) {
         Ok(w) => {
             unsafe { *out_buf_len = w as u16 };
             w as i32
@@ -124,7 +122,7 @@ pub extern "C" fn ossuary_handshake_done(conn: *const ConnectionContext) -> i32 
         return -1i32;
     }
     let conn = unsafe { &*conn };
-    let done = crypto_handshake_done(&conn);
+    let done = conn.crypto_handshake_done();
     ::std::mem::forget(conn);
     match done {
         Ok(done) => done as i32,
@@ -140,14 +138,14 @@ pub extern "C" fn ossuary_send_data(conn: *mut ConnectionContext,
         out_buf.is_null() || out_buf_len.is_null() {
         return -1i32;
     }
-    let mut conn = unsafe { &mut *conn };
+    let conn = unsafe { &mut *conn };
     let r_out_buf: &mut [u8] = unsafe {
         std::slice::from_raw_parts_mut(out_buf, *out_buf_len as usize)
     };
     let r_in_buf: &[u8] = unsafe { std::slice::from_raw_parts(in_buf, in_buf_len as usize) };
     let mut out_slice = r_out_buf;
     let in_slice = r_in_buf;
-    let bytes_written = match crypto_send_data(&mut conn, &in_slice, &mut out_slice) {
+    let bytes_written = match conn.crypto_send_data(&in_slice, &mut out_slice) {
         Ok(w) => {
             unsafe { *out_buf_len = w as u16; }
             w as i32
@@ -170,14 +168,14 @@ pub extern "C" fn ossuary_recv_data(conn: *mut ConnectionContext,
         in_buf_len.is_null() || out_buf_len.is_null() {
         return -1i32;
     }
-    let mut conn = unsafe { &mut *conn };
+    let conn = unsafe { &mut *conn };
     let r_out_buf: &mut [u8] = unsafe {
         std::slice::from_raw_parts_mut(out_buf, *out_buf_len as usize)
     };
     let r_in_buf: &[u8] = unsafe { std::slice::from_raw_parts(in_buf, *in_buf_len as usize) };
     let mut out_slice = r_out_buf;
     let mut in_slice = r_in_buf;
-    let bytes_read = match crypto_recv_data(&mut conn, &mut in_slice, &mut out_slice) {
+    let bytes_read = match conn.crypto_recv_data(&mut in_slice, &mut out_slice) {
         Ok((read,written)) => {
             unsafe {
                 *in_buf_len = read as u16;
@@ -203,12 +201,12 @@ pub extern "C" fn ossuary_flush(conn: *mut ConnectionContext,
     if conn.is_null() || out_buf.is_null() {
         return -1i32;
     }
-    let mut conn = unsafe { &mut *conn };
+    let conn = unsafe { &mut *conn };
     let r_out_buf: &mut [u8] = unsafe {
         std::slice::from_raw_parts_mut(out_buf, out_buf_len as usize)
     };
     let mut out_slice = r_out_buf;
-    let bytes_written = match crypto_flush(&mut conn, &mut out_slice) {
+    let bytes_written = match conn.crypto_flush(&mut out_slice) {
         Ok(x) => x as i32,
         Err(OssuaryError::WouldBlock(_)) => ERROR_WOULD_BLOCK,
         Err(_) => -1i32,

diff --git a/src/lib.rs b/src/lib.rs
line changes: +580/-582
index 543d8a6..ff07c15
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -566,6 +566,586 @@ impl ConnectionContext {
             }
         }
     }
+
+    /// Writes the next handshake packet to the output stream.
+    ///
+    ///
+    ///
+    /// On success, returns the number of bytes written to the output buffer.
+    pub fn crypto_send_handshake<T,U>(&mut self, mut buf: T) -> Result<usize, OssuaryError>
+    where T: std::ops::DerefMut<Target = U>,
+          U: std::io::Write {
+        // Try to send any unsent buffered data
+        match write_stored_packet(self, &mut buf) {
+            Ok(w) if w == 0 => {},
+            Ok(w) => return Err(OssuaryError::WouldBlock(w)),
+            Err(e) => return Err(e),
+        }
+        let written = match self.state {
+            ConnectionState::ServerNew => {
+                // Wait for client to initiate connection
+                0
+            },
+            ConnectionState::Encrypted => {
+                // Handshake finished
+                0
+            },
+            ConnectionState::ServerWaitAck(t) |
+            ConnectionState::ServerWaitAuth(t) |
+            ConnectionState::ClientWaitKey(t) |
+            ConnectionState::ClientWaitAck(t) => {
+                let mut w: usize = 0;
+                // Wait for response, with timeout
+                if let Ok(dur) = t.elapsed() {
+                    if dur.as_secs() > MAX_HANDSHAKE_WAIT_TIME {
+                        let pkt: HandshakePacket = Default::default();
+                        w = write_packet(self, &mut buf, struct_as_slice(&pkt),
+                                         PacketType::Reset)?;
+                        self.reset_state(None);
+                    }
+                }
+                w
+            },
+            ConnectionState::ServerSendPubKey => {
+                // Send session public key and nonce to the client
+                let mut pkt: HandshakePacket = Default::default();
+                pkt.public_key.copy_from_slice(&self.local_key.public);
+                pkt.nonce.copy_from_slice(&self.local_key.nonce);
+                let w = write_packet(self, &mut buf, struct_as_slice(&pkt),
+                                     PacketType::PublicKeyNonce)?;
+                self.state = ConnectionState::ServerWaitAck(std::time::SystemTime::now());
+                w
+            },
+            ConnectionState::ServerSendChallenge => {
+                match self.conn_type {
+                    ConnectionType::AuthenticatedServer => {
+                        // Send a block of random data over the encrypted session to
+                        // the client.  The client must sign it with its key to prove
+                        // key possession.
+                        let mut rng = match OsRng::new() {
+                            Ok(rng) => rng,
+                            Err(_) => {
+                                self.reset_state(None);
+                                return Err(OssuaryError::InvalidKey);
+                            }
+                        };
+                        let aad = [];
+                        let mut challenge: [u8; CHALLENGE_LEN] = [0; CHALLENGE_LEN];
+                        rng.fill_bytes(&mut challenge);
+                        self.challenge = Some(challenge.to_vec());
+                        let mut ciphertext = Vec::with_capacity(CHALLENGE_LEN);
+                        let session_key = match self.local_key.session {
+                            Some(ref s) => s,
+                            None => {
+                                self.reset_state(None);
+                                return Err(OssuaryError::InvalidKey);
+                            }
+                        };
+                        let tag = match encrypt(session_key.as_bytes(),
+                                                &self.local_key.nonce,
+                                                &aad, &challenge, &mut ciphertext) {
+                            Ok(tag) => tag,
+                            Err(_) => {
+                                self.reset_state(None);
+                                return Err(OssuaryError::InvalidKey);
+                            }
+                        };
+                        let pkt: EncryptedPacket = EncryptedPacket {
+                            tag_len: tag.len() as u16,
+                            data_len: ciphertext.len() as u16,
+                        };
+                        let mut pkt_buf: Vec<u8>= vec![];
+                        pkt_buf.extend(struct_as_slice(&pkt));
+                        pkt_buf.extend(&ciphertext);
+                        pkt_buf.extend(&tag);
+                        let w = write_packet(self, &mut buf, &pkt_buf,
+                                             PacketType::AuthChallenge)?;
+                        self.state = ConnectionState::ServerWaitAuth(std::time::SystemTime::now());
+                        w
+                    },
+                    _ => {
+                        // For unauthenticated connections, we are done.  Already encrypted.
+                        let pkt: HandshakePacket = Default::default();
+                        let w = write_packet(self, &mut buf, struct_as_slice(&pkt),
+                                             PacketType::PubKeyAck)?;
+                        self.state = ConnectionState::Encrypted;
+                        w // handshake is finished (success)
+                    },
+                }
+            },
+            ConnectionState::ClientNew => {
+                // Send session public key and nonce to initiate connection
+                let mut pkt: HandshakePacket = Default::default();
+                pkt.public_key.copy_from_slice(&self.local_key.public);
+                pkt.nonce.copy_from_slice(&self.local_key.nonce);
+                let w = write_packet(self, &mut buf, struct_as_slice(&pkt),
+                                     PacketType::PublicKeyNonce)?;
+                self.state = ConnectionState::ClientWaitKey(std::time::SystemTime::now());
+                w
+            },
+            ConnectionState::ClientSendAck => {
+                // Acknowledge reception of server's session public key and nonce
+                let pkt: HandshakePacket = Default::default();
+                let w = write_packet(self, &mut buf, struct_as_slice(&pkt),
+                                     PacketType::PubKeyAck)?;
+                self.state = ConnectionState::ClientWaitAck(std::time::SystemTime::now());
+                w
+            },
+            ConnectionState::ClientSendAuth => {
+                // Send signature of the server's challenge back to the server,
+                // along with the public part of the authentication key.  This is
+                // done over the established encrypted channel.
+                let secret = match self.secret_key {
+                    Some(ref s) => match SecretKey::from_bytes(s.as_bytes()) {
+                        Ok(s) => s, // local copy of secret key
+                        Err(_) => {
+                            self.reset_state(Some(OssuaryError::InvalidKey));
+                            return Err(OssuaryError::InvalidKey);
+                        }
+                    },
+                    None => {
+                        self.reset_state(Some(OssuaryError::InvalidKey));
+                        return Err(OssuaryError::InvalidKey);
+                    }
+                };
+                let public = PublicKey::from(&secret);
+                let keypair = Keypair { secret: secret, public: public };
+                let sig = match self.challenge {
+                    Some(ref c) => keypair.sign(c).to_bytes(),
+                    None => {
+                        self.reset_state(None);
+                        return Err(OssuaryError::InvalidSignature);
+                    }
+                };
+                let mut pkt_data: Vec<u8> = Vec::with_capacity(CHALLENGE_LEN + 32);
+                pkt_data.extend_from_slice(public.as_bytes());
+                pkt_data.extend_from_slice(&sig);
+                self.challenge_sig = Some(sig.to_vec());
+
+                let aad = [];
+                let mut ciphertext = Vec::with_capacity(pkt_data.len());
+                let session_key = match self.local_key.session {
+                    Some(ref s) => s,
+                    None => {
+                        self.reset_state(None);
+                        return Err(OssuaryError::InvalidKey);
+                    }
+                };
+                let tag = match encrypt(session_key.as_bytes(),
+                                        &self.local_key.nonce,
+                                        &aad, &pkt_data, &mut ciphertext) {
+                    Ok(t) => t,
+                    Err(_) => {
+                        self.reset_state(None);
+                        return Err(OssuaryError::InvalidKey);
+                    }
+                };
+
+                let pkt: EncryptedPacket = EncryptedPacket {
+                    tag_len: tag.len() as u16,
+                    data_len: ciphertext.len() as u16,
+                };
+                let mut pkt_buf: Vec<u8>= vec![];
+                pkt_buf.extend(struct_as_slice(&pkt));
+                pkt_buf.extend(&ciphertext);
+                pkt_buf.extend(&tag);
+                let w = write_packet(self, &mut buf, &pkt_buf,
+                                     PacketType::AuthResponse)?;
+                self.state = ConnectionState::Encrypted;
+                w // handshake is finished (success)
+            },
+            ConnectionState::Failed(_) => {
+                // This is a permanent failure.
+                let pkt: HandshakePacket = Default::default();
+                let w = write_packet(self, &mut buf, struct_as_slice(&pkt),
+                                     PacketType::Disconnect)?;
+                self.reset_state(Some(OssuaryError::ConnectionFailed));
+                w // handshake is finished (failed)
+            },
+        };
+        Ok(written)
+    }
+
+    /// Receive the next handshake packet from the input buffer
+    ///
+    /// On success, returns the number of bytes consumed from the input buffer.
+    pub fn crypto_recv_handshake<T,U>(&mut self, buf: T) -> Result<usize, OssuaryError>
+    where T: std::ops::DerefMut<Target = U>,
+          U: std::io::Read {
+        let mut bytes_read: usize = 0;
+
+        match self.state {
+            ConnectionState::Encrypted => return Ok(0),
+            _ => {},
+        }
+
+        let pkt: NetworkPacket = match read_packet(self, buf) {
+            Ok((p, r)) => {
+                bytes_read += r;
+                p
+            },
+            Err(OssuaryError::WouldBlock(b)) => {
+                return Err(OssuaryError::WouldBlock(b));
+            }
+            Err(e) => {
+                self.reset_state(None);
+                return Err(e);
+            }
+        };
+
+        let mut error = false;
+        match pkt.kind() {
+            PacketType::Reset => {
+                self.reset_state(None);
+                return Err(OssuaryError::ConnectionReset);
+            },
+            PacketType::Disconnect => {
+                self.reset_state(Some(OssuaryError::ConnectionFailed));
+                return Err(OssuaryError::ConnectionFailed);
+            },
+            _ => {},
+        }
+
+        if pkt.header.msg_id != self.remote_msg_id {
+            println!("Message gap detected.  Restarting connection.");
+            println!("Server: {}", self.is_server());
+            self.reset_state(None);
+            return Err(OssuaryError::InvalidPacket("Message ID does not match".into()));
+        }
+        self.remote_msg_id = pkt.header.msg_id + 1;
+
+        match self.state {
+            ConnectionState::ServerNew => {
+                match pkt.kind() {
+                    PacketType::PublicKeyNonce => {
+                        let data_pkt: Result<&HandshakePacket, _> = interpret_packet(&pkt);
+                        match data_pkt {
+                            Ok(ref data_pkt) => {
+                                self.add_remote_key(&data_pkt.public_key, &data_pkt.nonce);
+                                self.state = ConnectionState::ServerSendPubKey;
+                            },
+                            Err(_) => {
+                                error = true;
+                            },
+                        }
+                    },
+                    _ => { error = true; }
+                }
+            },
+            ConnectionState::ServerWaitAck(_t) => {
+                match pkt.kind() {
+                    PacketType::PubKeyAck => {
+                        self.state = ConnectionState::ServerSendChallenge;
+                    },
+                    _ => { error = true; }
+                }
+            },
+            ConnectionState::ServerWaitAuth(_t) => {
+                match pkt.kind() {
+                    PacketType::AuthResponse => {
+                        match interpret_packet_extra::<EncryptedPacket>(&pkt) {
+                            Ok((data_pkt, rest)) => {
+                                let ciphertext = &rest[..data_pkt.data_len as usize];
+                                let tag = &rest[data_pkt.data_len as usize..];
+                                let aad = [];
+                                let mut plaintext = Vec::with_capacity(ciphertext.len());
+                                let session_key = match self.local_key.session {
+                                    Some(ref k) => k,
+                                    None => {
+                                        self.reset_state(None);
+                                        return Err(OssuaryError::InvalidKey);
+                                    }
+                                };
+                                let remote_nonce = match self.remote_key {
+                                    Some(ref rem) => rem.nonce,
+                                    None => {
+                                        self.reset_state(None);
+                                        return Err(OssuaryError::InvalidKey);
+                                    }
+                                };
+                                decrypt(session_key.as_bytes(),
+                                        &remote_nonce,
+                                        &aad, &ciphertext, &tag, &mut plaintext)?;
+                                let pubkey = &plaintext[0..32];
+                                let sig = &plaintext[32..];
+                                if self.authorized_keys.iter().filter(
+                                    |k| &pubkey == k).count() > 0 {
+                                    let public = match PublicKey::from_bytes(pubkey) {
+                                        Ok(p) => p,
+                                        Err(_) => {
+                                            self.reset_state(None);
+                                            return Err(OssuaryError::InvalidKey);
+                                        }
+                                    };
+                                    let sig = match Signature::from_bytes(sig) {
+                                        Ok(s) => s,
+                                        Err(_) => {
+                                            self.reset_state(None);
+                                            return Err(OssuaryError::InvalidKey);
+                                        }
+                                    };
+                                    let challenge = match self.challenge {
+                                        Some(ref c) => c,
+                                        None => {
+                                            self.reset_state(None);
+                                            return Err(OssuaryError::InvalidKey);
+                                        }
+                                    };
+                                    match public.verify(challenge, &sig) {
+                                        Ok(_) => {
+                                            self.state = ConnectionState::Encrypted;
+                                        },
+                                        Err(_) => {
+                                            println!("Verify bad");
+                                            self.state = ConnectionState::Failed(
+                                                OssuaryError::InvalidSignature);
+                                        },
+                                    }
+                                }
+                                else {
+                                    println!("Key not allowed");
+                                    self.state = ConnectionState::Failed(OssuaryError::InvalidKey);
+                                }
+                            },
+                            Err(_) => {
+                                self.reset_state(None);
+                                return Err(OssuaryError::InvalidPacket("Response invalid".into()));
+                            },
+                        };
+                    },
+                    _ => { error = true; }
+                }
+            },
+            ConnectionState::ClientWaitKey(_t) => {
+                match pkt.kind() {
+                    PacketType::PublicKeyNonce => {
+                        let data_pkt: Result<&HandshakePacket, _> = interpret_packet(&pkt);
+                        match data_pkt {
+                            Ok(data_pkt) => {
+                                self.add_remote_key(&data_pkt.public_key, &data_pkt.nonce);
+                                self.state = ConnectionState::ClientSendAck;
+                            },
+                            Err(_) => {
+                                error = true;
+                            },
+                        }
+                    },
+                    _ => {
+                        error = true;
+                    }
+                }
+            },
+            ConnectionState::ClientWaitAck(_t) => {
+                match pkt.kind() {
+                    PacketType::PubKeyAck => {
+                        self.state = ConnectionState::Encrypted;
+                    },
+                    PacketType::AuthChallenge => {
+                        match interpret_packet_extra::<EncryptedPacket>(&pkt) {
+                            Ok((data_pkt, rest)) => {
+                                let ciphertext = &rest[..data_pkt.data_len as usize];
+                                let tag = &rest[data_pkt.data_len as usize..];
+                                let aad = [];
+                                let mut plaintext = Vec::with_capacity(ciphertext.len());
+
+                                let session_key = match self.local_key.session {
+                                    Some(ref k) => k,
+                                    None => {
+                                        self.reset_state(None);
+                                        return Err(OssuaryError::InvalidKey);
+                                    }
+                                };
+                                let remote_nonce = match self.remote_key {
+                                    Some(ref rem) => rem.nonce,
+                                    None => {
+                                        self.reset_state(None);
+                                        return Err(OssuaryError::InvalidKey);
+                                    }
+                                };
+                                decrypt(session_key.as_bytes(),
+                                        &remote_nonce,
+                                        &aad, &ciphertext, &tag, &mut plaintext)?;
+                                self.challenge = Some(plaintext);
+                                self.state = ConnectionState::ClientSendAuth;
+                            },
+                            Err(_) => {
+                                error = true;
+                            },
+                        }
+                    },
+                    _ => {
+                        error = true;
+                    },
+                }
+            },
+            ConnectionState::ServerSendPubKey |
+            ConnectionState::ServerSendChallenge |
+            ConnectionState::ClientNew |
+            ConnectionState::ClientSendAck |
+            ConnectionState::ClientSendAuth |
+            ConnectionState::Encrypted => {
+                // no-op
+            },
+            ConnectionState::Failed(_) => {
+                error = true;
+            },
+        }
+        if error {
+            self.reset_state(None);
+            return Err(OssuaryError::ConnectionReset);
+        }
+        Ok(bytes_read)
+    }
+
+    /// Returns whether the handshake process is complete.
+    ///
+    ///
+    pub fn crypto_handshake_done(&self) -> Result<bool, &OssuaryError> {
+        match self.state {
+            ConnectionState::Encrypted => Ok(true),
+            ConnectionState::Failed(ref e) => Err(e),
+            _ => Ok(false),
+        }
+    }
+
+    pub fn crypto_send_data<T,U>(&mut self,
+                                 in_buf: &[u8],
+                                 mut out_buf: T) -> Result<usize, OssuaryError>
+    where T: std::ops::DerefMut<Target = U>,
+          U: std::io::Write {
+        // Try to send any unsent buffered data
+        match write_stored_packet(self, &mut out_buf) {
+            Ok(w) if w == 0 => {},
+            Ok(w) => return Err(OssuaryError::WouldBlock(w)),
+            Err(e) => return Err(e),
+        }
+        match self.state {
+            ConnectionState::Encrypted => {},
+            _ => {
+                return Err(OssuaryError::InvalidPacket(
+                    "Encrypted channel not established.".into()));
+            }
+        }
+        let aad = [];
+        let mut ciphertext = Vec::with_capacity(in_buf.len());
+        let session_key = match self.local_key.session {
+            Some(ref k) => k,
+            None => {
+                self.reset_state(None);
+                return Err(OssuaryError::InvalidKey);;
+            }
+        };
+        let tag = match encrypt(session_key.as_bytes(),
+                                &self.local_key.nonce, &aad, in_buf, &mut ciphertext) {
+            Ok(t) => t,
+            Err(_) => {
+                self.reset_state(None);
+                return Err(OssuaryError::InvalidKey);;
+            }
+        };
+
+        let pkt: EncryptedPacket = EncryptedPacket {
+            tag_len: tag.len() as u16,
+            data_len: ciphertext.len() as u16,
+        };
+        let mut buf: Vec<u8>= vec![];
+        buf.extend(struct_as_slice(&pkt));
+        buf.extend(&ciphertext);
+        buf.extend(&tag);
+        let written = write_packet(self, &mut out_buf, &buf,
+                                   PacketType::EncryptedData)?;
+        Ok(written)
+    }
+
+    pub fn crypto_recv_data<T,U,R,V>(&mut self,
+                                     in_buf: T,
+                                     mut out_buf: R) -> Result<(usize, usize), OssuaryError>
+    where T: std::ops::DerefMut<Target = U>,
+          U: std::io::Read,
+          R: std::ops::DerefMut<Target = V>,
+          V: std::io::Write {
+        let bytes_written: usize;
+        let mut bytes_read: usize = 0;
+        match self.state {
+            ConnectionState::Encrypted => {},
+            _ => {
+                return Err(OssuaryError::InvalidPacket(
+                    "Encrypted channel not established.".into()));
+            }
+        }
+
+        match read_packet(self, in_buf) {
+            Ok((pkt, bytes)) => {
+                bytes_read += bytes;
+                if pkt.header.msg_id != self.remote_msg_id {
+                    let msg_id = pkt.header.msg_id;
+                    println!("Message gap detected.  Restarting connection. ({} != {})",
+                             msg_id, self.remote_msg_id);
+                    println!("Server: {}", self.is_server());
+                    self.reset_state(None);
+                    return Err(OssuaryError::InvalidPacket("Message ID mismatch".into()))
+                }
+                self.remote_msg_id = pkt.header.msg_id + 1;
+
+                match pkt.kind() {
+                    PacketType::EncryptedData => {
+                        match interpret_packet_extra::<EncryptedPacket>(&pkt) {
+                            Ok((data_pkt, rest)) => {
+                                let ciphertext = &rest[..data_pkt.data_len as usize];
+                                let tag = &rest[data_pkt.data_len as usize..];
+                                let aad = [];
+                                let mut plaintext = Vec::with_capacity(ciphertext.len());
+                                let session_key = match self.local_key.session {
+                                    Some(ref k) => k,
+                                    None => {
+                                        self.reset_state(None);
+                                        return Err(OssuaryError::InvalidKey);
+                                    }
+                                };
+                                let remote_nonce = match self.remote_key {
+                                    Some(ref rem) => rem.nonce,
+                                    None => {
+                                        self.reset_state(None);
+                                        return Err(OssuaryError::InvalidKey);
+                                    }
+                                };
+                                decrypt(session_key.as_bytes(),
+                                        &remote_nonce,
+                                        &aad, &ciphertext, &tag, &mut plaintext)?;
+                                bytes_written = match out_buf.write(&plaintext) {
+                                    Ok(w) => w,
+                                    Err(e) => return Err(e.into()),
+                                };
+                            },
+                            Err(_) => {
+                                self.reset_state(None);
+                                return Err(OssuaryError::InvalidKey);
+                            },
+                        }
+                    },
+                    _ => {
+                        return Err(OssuaryError::InvalidPacket(
+                            "Received non-encrypted data on encrypted channel.".into()));
+                    },
+                }
+            },
+            Err(OssuaryError::WouldBlock(b)) => {
+                return Err(OssuaryError::WouldBlock(b));
+            },
+            Err(_e) => {
+                self.reset_state(None);
+                return Err(OssuaryError::InvalidPacket("Packet header did not parse.".into()));
+            },
+        }
+        Ok((bytes_read, bytes_written))
+    }
+
+    pub fn crypto_flush<R,V>(&mut self,
+                             mut out_buf: R) -> Result<usize, OssuaryError>
+    where R: std::ops::DerefMut<Target = V>,
+          V: std::io::Write {
+        return write_stored_packet(self, &mut out_buf);
+    }
 }
 
 /// Cast the data bytes in a NetworkPacket into a struct
@@ -691,588 +1271,6 @@ where T: std::ops::DerefMut<Target = U>,
     Ok(written)
 }
 
-/// Writes the next handshake packet to the output stream.
-///
-///
-///
-/// On success, returns the number of bytes written to the output buffer.
-pub fn crypto_send_handshake<T,U>(conn: &mut ConnectionContext,
-                                  mut buf: T) -> Result<usize, OssuaryError>
-where T: std::ops::DerefMut<Target = U>,
-      U: std::io::Write {
-    // Try to send any unsent buffered data
-    match write_stored_packet(conn, &mut buf) {
-        Ok(w) if w == 0 => {},
-        Ok(w) => return Err(OssuaryError::WouldBlock(w)),
-        Err(e) => return Err(e),
-    }
-    let written = match conn.state {
-        ConnectionState::ServerNew => {
-            // Wait for client to initiate connection
-            0
-        },
-        ConnectionState::Encrypted => {
-            // Handshake finished
-            0
-        },
-        ConnectionState::ServerWaitAck(t) |
-        ConnectionState::ServerWaitAuth(t) |
-        ConnectionState::ClientWaitKey(t) |
-        ConnectionState::ClientWaitAck(t) => {
-            let mut w: usize = 0;
-            // Wait for response, with timeout
-            if let Ok(dur) = t.elapsed() {
-                if dur.as_secs() > MAX_HANDSHAKE_WAIT_TIME {
-                    let pkt: HandshakePacket = Default::default();
-                    w = write_packet(conn, &mut buf, struct_as_slice(&pkt),
-                                     PacketType::Reset)?;
-                    conn.reset_state(None);
-                }
-            }
-            w
-        },
-        ConnectionState::ServerSendPubKey => {
-            // Send session public key and nonce to the client
-            let mut pkt: HandshakePacket = Default::default();
-            pkt.public_key.copy_from_slice(&conn.local_key.public);
-            pkt.nonce.copy_from_slice(&conn.local_key.nonce);
-            let w = write_packet(conn, &mut buf, struct_as_slice(&pkt),
-                                 PacketType::PublicKeyNonce)?;
-            conn.state = ConnectionState::ServerWaitAck(std::time::SystemTime::now());
-            w
-        },
-        ConnectionState::ServerSendChallenge => {
-            match conn.conn_type {
-                ConnectionType::AuthenticatedServer => {
-                    // Send a block of random data over the encrypted session to
-                    // the client.  The client must sign it with its key to prove
-                    // key possession.
-                    let mut rng = match OsRng::new() {
-                        Ok(rng) => rng,
-                        Err(_) => {
-                            conn.reset_state(None);
-                            return Err(OssuaryError::InvalidKey);
-                        }
-                    };
-                    let aad = [];
-                    let mut challenge: [u8; CHALLENGE_LEN] = [0; CHALLENGE_LEN];
-                    rng.fill_bytes(&mut challenge);
-                    conn.challenge = Some(challenge.to_vec());
-                    let mut ciphertext = Vec::with_capacity(CHALLENGE_LEN);
-                    let session_key = match conn.local_key.session {
-                        Some(ref s) => s,
-                        None => {
-                            conn.reset_state(None);
-                            return Err(OssuaryError::InvalidKey);
-                        }
-                    };
-                    let tag = match encrypt(session_key.as_bytes(),
-                                            &conn.local_key.nonce,
-                                            &aad, &challenge, &mut ciphertext) {
-                        Ok(tag) => tag,
-                        Err(_) => {
-                            conn.reset_state(None);
-                            return Err(OssuaryError::InvalidKey);
-                        }
-                    };
-                    let pkt: EncryptedPacket = EncryptedPacket {
-                        tag_len: tag.len() as u16,
-                        data_len: ciphertext.len() as u16,
-                    };
-                    let mut pkt_buf: Vec<u8>= vec![];
-                    pkt_buf.extend(struct_as_slice(&pkt));
-                    pkt_buf.extend(&ciphertext);
-                    pkt_buf.extend(&tag);
-                    let w = write_packet(conn, &mut buf, &pkt_buf,
-                                         PacketType::AuthChallenge)?;
-                    conn.state = ConnectionState::ServerWaitAuth(std::time::SystemTime::now());
-                    w
-                },
-                _ => {
-                    // For unauthenticated connections, we are done.  Already encrypted.
-                    let pkt: HandshakePacket = Default::default();
-                    let w = write_packet(conn, &mut buf, struct_as_slice(&pkt),
-                                         PacketType::PubKeyAck)?;
-                    conn.state = ConnectionState::Encrypted;
-                    w // handshake is finished (success)
-                },
-            }
-        },
-        ConnectionState::ClientNew => {
-            // Send session public key and nonce to initiate connection
-            let mut pkt: HandshakePacket = Default::default();
-            pkt.public_key.copy_from_slice(&conn.local_key.public);
-            pkt.nonce.copy_from_slice(&conn.local_key.nonce);
-            let w = write_packet(conn, &mut buf, struct_as_slice(&pkt),
-                                 PacketType::PublicKeyNonce)?;
-            conn.state = ConnectionState::ClientWaitKey(std::time::SystemTime::now());
-            w
-        },
-        ConnectionState::ClientSendAck => {
-            // Acknowledge reception of server's session public key and nonce
-            let pkt: HandshakePacket = Default::default();
-            let w = write_packet(conn, &mut buf, struct_as_slice(&pkt),
-                                 PacketType::PubKeyAck)?;
-            conn.state = ConnectionState::ClientWaitAck(std::time::SystemTime::now());
-            w
-        },
-        ConnectionState::ClientSendAuth => {
-            // Send signature of the server's challenge back to the server,
-            // along with the public part of the authentication key.  This is
-            // done over the established encrypted channel.
-            let secret = match conn.secret_key {
-                Some(ref s) => match SecretKey::from_bytes(s.as_bytes()) {
-                    Ok(s) => s, // local copy of secret key
-                    Err(_) => {
-                        conn.reset_state(Some(OssuaryError::InvalidKey));
-                        return Err(OssuaryError::InvalidKey);
-                    }
-                },
-                None => {
-                    conn.reset_state(Some(OssuaryError::InvalidKey));
-                    return Err(OssuaryError::InvalidKey);
-                }
-            };
-            let public = PublicKey::from(&secret);
-            let keypair = Keypair { secret: secret, public: public };
-            let sig = match conn.challenge {
-                Some(ref c) => keypair.sign(c).to_bytes(),
-                None => {
-                    conn.reset_state(None);
-                    return Err(OssuaryError::InvalidSignature);
-                }
-            };
-            let mut pkt_data: Vec<u8> = Vec::with_capacity(CHALLENGE_LEN + 32);
-            pkt_data.extend_from_slice(public.as_bytes());
-            pkt_data.extend_from_slice(&sig);
-            conn.challenge_sig = Some(sig.to_vec());
-
-            let aad = [];
-            let mut ciphertext = Vec::with_capacity(pkt_data.len());
-            let session_key = match conn.local_key.session {
-                Some(ref s) => s,
-                None => {
-                    conn.reset_state(None);
-                    return Err(OssuaryError::InvalidKey);
-                }
-            };
-            let tag = match encrypt(session_key.as_bytes(),
-                                    &conn.local_key.nonce,
-                                    &aad, &pkt_data, &mut ciphertext) {
-                Ok(t) => t,
-                Err(_) => {
-                    conn.reset_state(None);
-                    return Err(OssuaryError::InvalidKey);
-                }
-            };
-
-            let pkt: EncryptedPacket = EncryptedPacket {
-                tag_len: tag.len() as u16,
-                data_len: ciphertext.len() as u16,
-            };
-            let mut pkt_buf: Vec<u8>= vec![];
-            pkt_buf.extend(struct_as_slice(&pkt));
-            pkt_buf.extend(&ciphertext);
-            pkt_buf.extend(&tag);
-            let w = write_packet(conn, &mut buf, &pkt_buf,
-                                 PacketType::AuthResponse)?;
-            conn.state = ConnectionState::Encrypted;
-            w // handshake is finished (success)
-        },
-        ConnectionState::Failed(_) => {
-            // This is a permanent failure.
-            let pkt: HandshakePacket = Default::default();
-            let w = write_packet(conn, &mut buf, struct_as_slice(&pkt),
-                                 PacketType::Disconnect)?;
-            conn.reset_state(Some(OssuaryError::ConnectionFailed));
-            w // handshake is finished (failed)
-        },
-    };
-    Ok(written)
-}
-
-/// Receive the next handshake packet from the input buffer
-///
-/// On success, returns the number of bytes consumed from the input buffer.
-pub fn crypto_recv_handshake<T,U>(conn: &mut ConnectionContext,
-                                  buf: T) -> Result<usize, OssuaryError>
-where T: std::ops::DerefMut<Target = U>,
-      U: std::io::Read {
-    let mut bytes_read: usize = 0;
-
-    match conn.state {
-        ConnectionState::Encrypted => return Ok(0),
-        _ => {},
-    }
-
-    let pkt: NetworkPacket = match read_packet(conn, buf) {
-        Ok((p, r)) => {
-            bytes_read += r;
-            p
-        },
-        Err(OssuaryError::WouldBlock(b)) => {
-            return Err(OssuaryError::WouldBlock(b));
-        }
-        Err(e) => {
-            conn.reset_state(None);
-            return Err(e);
-        }
-    };
-
-    let mut error = false;
-    match pkt.kind() {
-        PacketType::Reset => {
-            conn.reset_state(None);
-            return Err(OssuaryError::ConnectionReset);
-        },
-        PacketType::Disconnect => {
-            conn.reset_state(Some(OssuaryError::ConnectionFailed));
-            return Err(OssuaryError::ConnectionFailed);
-        },
-        _ => {},
-    }
-
-    if pkt.header.msg_id != conn.remote_msg_id {
-        println!("Message gap detected.  Restarting connection.");
-        println!("Server: {}", conn.is_server());
-        conn.reset_state(None);
-        return Err(OssuaryError::InvalidPacket("Message ID does not match".into()));
-    }
-    conn.remote_msg_id = pkt.header.msg_id + 1;
-
-    match conn.state {
-        ConnectionState::ServerNew => {
-            match pkt.kind() {
-                PacketType::PublicKeyNonce => {
-                    let data_pkt: Result<&HandshakePacket, _> = interpret_packet(&pkt);
-                    match data_pkt {
-                        Ok(ref data_pkt) => {
-                            conn.add_remote_key(&data_pkt.public_key, &data_pkt.nonce);
-                            conn.state = ConnectionState::ServerSendPubKey;
-                        },
-                        Err(_) => {
-                            error = true;
-                        },
-                    }
-                },
-                _ => { error = true; }
-            }
-        },
-        ConnectionState::ServerWaitAck(_t) => {
-            match pkt.kind() {
-                PacketType::PubKeyAck => {
-                    conn.state = ConnectionState::ServerSendChallenge;
-                },
-                _ => { error = true; }
-            }
-        },
-        ConnectionState::ServerWaitAuth(_t) => {
-            match pkt.kind() {
-                PacketType::AuthResponse => {
-                    match interpret_packet_extra::<EncryptedPacket>(&pkt) {
-                        Ok((data_pkt, rest)) => {
-                            let ciphertext = &rest[..data_pkt.data_len as usize];
-                            let tag = &rest[data_pkt.data_len as usize..];
-                            let aad = [];
-                            let mut plaintext = Vec::with_capacity(ciphertext.len());
-                            let session_key = match conn.local_key.session {
-                                Some(ref k) => k,
-                                None => {
-                                    conn.reset_state(None);
-                                    return Err(OssuaryError::InvalidKey);
-                                }
-                            };
-                            let remote_nonce = match conn.remote_key {
-                                Some(ref rem) => rem.nonce,
-                                None => {
-                                    conn.reset_state(None);
-                                    return Err(OssuaryError::InvalidKey);
-                                }
-                            };
-                            decrypt(session_key.as_bytes(),
-                                    &remote_nonce,
-                                    &aad, &ciphertext, &tag, &mut plaintext)?;
-                            let pubkey = &plaintext[0..32];
-                            let sig = &plaintext[32..];
-                            if conn.authorized_keys.iter().filter(
-                                |k| &pubkey == k).count() > 0 {
-                                let public = match PublicKey::from_bytes(pubkey) {
-                                    Ok(p) => p,
-                                    Err(_) => {
-                                        conn.reset_state(None);
-                                        return Err(OssuaryError::InvalidKey);
-                                    }
-                                };
-                                let sig = match Signature::from_bytes(sig) {
-                                    Ok(s) => s,
-                                    Err(_) => {
-                                        conn.reset_state(None);
-                                        return Err(OssuaryError::InvalidKey);
-                                    }
-                                };
-                                let challenge = match conn.challenge {
-                                    Some(ref c) => c,
-                                    None => {
-                                        conn.reset_state(None);
-                                        return Err(OssuaryError::InvalidKey);
-                                    }
-                                };
-                                match public.verify(challenge, &sig) {
-                                    Ok(_) => {
-                                        conn.state = ConnectionState::Encrypted;
-                                    },
-                                    Err(_) => {
-                                        println!("Verify bad");
-                                        conn.state = ConnectionState::Failed(
-                                            OssuaryError::InvalidSignature);
-                                    },
-                                }
-                            }
-                            else {
-                                println!("Key not allowed");
-                                conn.state = ConnectionState::Failed(OssuaryError::InvalidKey);
-                            }
-                        },
-                        Err(_) => {
-                            conn.reset_state(None);
-                            return Err(OssuaryError::InvalidPacket("Response invalid".into()));
-                        },
-                    };
-                },
-                _ => { error = true; }
-            }
-        },
-        ConnectionState::ClientWaitKey(_t) => {
-            match pkt.kind() {
-                PacketType::PublicKeyNonce => {
-                    let data_pkt: Result<&HandshakePacket, _> = interpret_packet(&pkt);
-                    match data_pkt {
-                        Ok(data_pkt) => {
-                            conn.add_remote_key(&data_pkt.public_key, &data_pkt.nonce);
-                            conn.state = ConnectionState::ClientSendAck;
-                        },
-                        Err(_) => {
-                            error = true;
-                        },
-                    }
-                },
-                _ => {
-                    error = true;
-                }
-            }
-        },
-        ConnectionState::ClientWaitAck(_t) => {
-            match pkt.kind() {
-                PacketType::PubKeyAck => {
-                    conn.state = ConnectionState::Encrypted;
-                },
-                PacketType::AuthChallenge => {
-                    match interpret_packet_extra::<EncryptedPacket>(&pkt) {
-                        Ok((data_pkt, rest)) => {
-                            let ciphertext = &rest[..data_pkt.data_len as usize];
-                            let tag = &rest[data_pkt.data_len as usize..];
-                            let aad = [];
-                            let mut plaintext = Vec::with_capacity(ciphertext.len());
-
-                            let session_key = match conn.local_key.session {
-                                Some(ref k) => k,
-                                None => {
-                                    conn.reset_state(None);
-                                    return Err(OssuaryError::InvalidKey);
-                                }
-                            };
-                            let remote_nonce = match conn.remote_key {
-                                Some(ref rem) => rem.nonce,
-                                None => {
-                                    conn.reset_state(None);
-                                    return Err(OssuaryError::InvalidKey);
-                                }
-                            };
-                            decrypt(session_key.as_bytes(),
-                                    &remote_nonce,
-                                    &aad, &ciphertext, &tag, &mut plaintext)?;
-                            conn.challenge = Some(plaintext);
-                            conn.state = ConnectionState::ClientSendAuth;
-                        },
-                        Err(_) => {
-                            error = true;
-                        },
-                    }
-                },
-                _ => {
-                    error = true;
-                },
-            }
-        },
-        ConnectionState::ServerSendPubKey |
-        ConnectionState::ServerSendChallenge |
-        ConnectionState::ClientNew |
-        ConnectionState::ClientSendAck |
-        ConnectionState::ClientSendAuth |
-        ConnectionState::Encrypted => {
-            // no-op
-        },
-        ConnectionState::Failed(_) => {
-            error = true;
-        },
-    }
-    if error {
-        conn.reset_state(None);
-        return Err(OssuaryError::ConnectionReset);
-    }
-    Ok(bytes_read)
-}
-
-/// Returns whether the handshake process is complete.
-///
-///
-pub fn crypto_handshake_done(conn: &ConnectionContext) -> Result<bool, &OssuaryError> {
-    match conn.state {
-        ConnectionState::Encrypted => Ok(true),
-        ConnectionState::Failed(ref e) => Err(e),
-        _ => Ok(false),
-    }
-}
-
-pub fn crypto_send_data<T,U>(conn: &mut ConnectionContext,
-                             in_buf: &[u8],
-                             mut out_buf: T) -> Result<usize, OssuaryError>
-where T: std::ops::DerefMut<Target = U>,
-      U: std::io::Write {
-    // Try to send any unsent buffered data
-    match write_stored_packet(conn, &mut out_buf) {
-        Ok(w) if w == 0 => {},
-        Ok(w) => return Err(OssuaryError::WouldBlock(w)),
-        Err(e) => return Err(e),
-    }
-    match conn.state {
-        ConnectionState::Encrypted => {},
-        _ => {
-            return Err(OssuaryError::InvalidPacket(
-                "Encrypted channel not established.".into()));
-        }
-    }
-    let aad = [];
-    let mut ciphertext = Vec::with_capacity(in_buf.len());
-    let session_key = match conn.local_key.session {
-        Some(ref k) => k,
-        None => {
-            conn.reset_state(None);
-            return Err(OssuaryError::InvalidKey);;
-        }
-    };
-    let tag = match encrypt(session_key.as_bytes(),
-                            &conn.local_key.nonce, &aad, in_buf, &mut ciphertext) {
-        Ok(t) => t,
-        Err(_) => {
-            conn.reset_state(None);
-            return Err(OssuaryError::InvalidKey);;
-        }
-    };
-
-    let pkt: EncryptedPacket = EncryptedPacket {
-        tag_len: tag.len() as u16,
-        data_len: ciphertext.len() as u16,
-    };
-    let mut buf: Vec<u8>= vec![];
-    buf.extend(struct_as_slice(&pkt));
-    buf.extend(&ciphertext);
-    buf.extend(&tag);
-    let written = write_packet(conn, &mut out_buf, &buf,
-                               PacketType::EncryptedData)?;
-    Ok(written)
-}
-
-pub fn crypto_recv_data<T,U,R,V>(conn: &mut ConnectionContext,
-                                 in_buf: T,
-                                 mut out_buf: R) -> Result<(usize, usize), OssuaryError>
-where T: std::ops::DerefMut<Target = U>,
-      U: std::io::Read,
-      R: std::ops::DerefMut<Target = V>,
-      V: std::io::Write {
-    let bytes_written: usize;
-    let mut bytes_read: usize = 0;
-    match conn.state {
-        ConnectionState::Encrypted => {},
-        _ => {
-            return Err(OssuaryError::InvalidPacket(
-                "Encrypted channel not established.".into()));
-        }
-    }
-
-    match read_packet(conn, in_buf) {
-        Ok((pkt, bytes)) => {
-            bytes_read += bytes;
-            if pkt.header.msg_id != conn.remote_msg_id {
-                let msg_id = pkt.header.msg_id;
-                println!("Message gap detected.  Restarting connection. ({} != {})",
-                         msg_id, conn.remote_msg_id);
-                println!("Server: {}", conn.is_server());
-                conn.reset_state(None);
-                return Err(OssuaryError::InvalidPacket("Message ID mismatch".into()))
-            }
-            conn.remote_msg_id = pkt.header.msg_id + 1;
-
-            match pkt.kind() {
-                PacketType::EncryptedData => {
-                    match interpret_packet_extra::<EncryptedPacket>(&pkt) {
-                        Ok((data_pkt, rest)) => {
-                            let ciphertext = &rest[..data_pkt.data_len as usize];
-                            let tag = &rest[data_pkt.data_len as usize..];
-                            let aad = [];
-                            let mut plaintext = Vec::with_capacity(ciphertext.len());
-                            let session_key = match conn.local_key.session {
-                                Some(ref k) => k,
-                                None => {
-                                    conn.reset_state(None);
-                                    return Err(OssuaryError::InvalidKey);
-                                }
-                            };
-                            let remote_nonce = match conn.remote_key {
-                                Some(ref rem) => rem.nonce,
-                                None => {
-                                    conn.reset_state(None);
-                                    return Err(OssuaryError::InvalidKey);
-                                }
-                            };
-                            decrypt(session_key.as_bytes(),
-                                    &remote_nonce,
-                                    &aad, &ciphertext, &tag, &mut plaintext)?;
-                            bytes_written = match out_buf.write(&plaintext) {
-                                Ok(w) => w,
-                                Err(e) => return Err(e.into()),
-                            };
-                        },
-                        Err(_) => {
-                            conn.reset_state(None);
-                            return Err(OssuaryError::InvalidKey);
-                        },
-                    }
-                },
-                _ => {
-                    return Err(OssuaryError::InvalidPacket(
-                        "Received non-encrypted data on encrypted channel.".into()));
-                },
-            }
-        },
-        Err(OssuaryError::WouldBlock(b)) => {
-            return Err(OssuaryError::WouldBlock(b));
-        },
-        Err(_e) => {
-            conn.reset_state(None);
-            return Err(OssuaryError::InvalidPacket("Packet header did not parse.".into()));
-        },
-    }
-    Ok((bytes_read, bytes_written))
-}
-
-pub fn crypto_flush<R,V>(conn: &mut ConnectionContext,
-                         mut out_buf: R) -> Result<usize, OssuaryError>
-where R: std::ops::DerefMut<Target = V>,
-      V: std::io::Write {
-    return write_stored_packet(conn, &mut out_buf);
-}
-
 #[cfg(test)]
 mod tests {
     use crate::*;

diff --git a/tests/basic.rs b/tests/basic.rs
line changes: +5/-7
index 2bc5a12..b422754
--- a/tests/basic.rs
+++ b/tests/basic.rs
@@ -1,6 +1,4 @@
 use ossuary::{ConnectionContext, ConnectionType};
-use ossuary::{crypto_send_handshake,crypto_recv_handshake, crypto_handshake_done};
-use ossuary::{crypto_send_data,crypto_recv_data};
 use ossuary::OssuaryError;
 
 use std::thread;
@@ -11,10 +9,10 @@ fn event_loop<T>(mut conn: ConnectionContext,
                  is_server: bool) -> Result<(), std::io::Error>
 where T: std::io::Read + std::io::Write {
     // Run the opaque handshake until the connection is established
-    while crypto_handshake_done(&conn).unwrap() == false {
-        if crypto_send_handshake(&mut conn, &mut stream).is_ok() {
+    while conn.crypto_handshake_done().unwrap() == false {
+        if conn.crypto_send_handshake(&mut stream).is_ok() {
             loop {
-                match crypto_recv_handshake(&mut conn, &mut stream) {
+                match conn.crypto_recv_handshake(&mut stream) {
                     Ok(_) => break,
                     Err(OssuaryError::WouldBlock(_)) => {},
                     _ => panic!("Handshake failed."),
@@ -29,12 +27,12 @@ where T: std::io::Read + std::io::Write {
         true => (strings.0.as_bytes(), strings.1.as_bytes()),
         false => (strings.1.as_bytes(), strings.0.as_bytes()),
     };
-    let _ = crypto_send_data(&mut conn, &mut plaintext, &mut stream);
+    let _ = conn.crypto_send_data(&mut plaintext, &mut stream);
 
     // Read a message from the other party
     let mut recv_plaintext = vec!();
     loop {
-        match crypto_recv_data(&mut conn, &mut stream, &mut recv_plaintext) {
+        match conn.crypto_recv_data(&mut stream, &mut recv_plaintext) {
             Ok(_) => {
                 println!("(basic) received: {:?}",
                          String::from_utf8(recv_plaintext.clone()).unwrap());

diff --git a/tests/basic_auth.rs b/tests/basic_auth.rs
line changes: +5/-7
index fad8807..e1aab10
--- a/tests/basic_auth.rs
+++ b/tests/basic_auth.rs
@@ -1,6 +1,4 @@
 use ossuary::{ConnectionContext, ConnectionType};
-use ossuary::{crypto_send_handshake,crypto_recv_handshake, crypto_handshake_done};
-use ossuary::{crypto_send_data,crypto_recv_data};
 use ossuary::OssuaryError;
 
 use std::thread;
@@ -11,10 +9,10 @@ fn event_loop<T>(mut conn: ConnectionContext,
                  is_server: bool) -> Result<(), std::io::Error>
 where T: std::io::Read + std::io::Write {
     // Run the opaque handshake until the connection is established
-    while crypto_handshake_done(&conn).unwrap() == false {
-        if crypto_send_handshake(&mut conn, &mut stream).is_ok() {
+    while conn.crypto_handshake_done().unwrap() == false {
+        if conn.crypto_send_handshake(&mut stream).is_ok() {
             loop {
-                match crypto_recv_handshake(&mut conn, &mut stream) {
+                match conn.crypto_recv_handshake(&mut stream) {
                     Ok(_) => break,
                     Err(OssuaryError::WouldBlock(_)) => {},
                     _ => panic!("Handshake failed."),
@@ -29,12 +27,12 @@ where T: std::io::Read + std::io::Write {
         true => (strings.0.as_bytes(), strings.1.as_bytes()),
         false => (strings.1.as_bytes(), strings.0.as_bytes()),
     };
-    let _ = crypto_send_data(&mut conn, &mut plaintext, &mut stream);
+    let _ = conn.crypto_send_data(&mut plaintext, &mut stream);
 
     // Read a message from the other party
     let mut recv_plaintext = vec!();
     loop {
-        match crypto_recv_data(&mut conn, &mut stream, &mut recv_plaintext) {
+        match conn.crypto_recv_data(&mut stream, &mut recv_plaintext) {
             Ok(_) => {
                 println!("(basic_auth) received: {:?}",
                          String::from_utf8(recv_plaintext.clone()).unwrap());