Initial AES-CTR implementation

.rustfmt.toml

@@ -1,5 +1,6 @@
 chain_split_single_child = true
 closure_block_indent_threshold = 1
+fn_args_density = "Compressed"
 wrap_comments = true
 control_brace_style = "ClosingNextLine"
 max_width = 80

src/algorithm.rs

@@ -28,10 +28,8 @@ pub static COMPRESSION: &[CompressionAlgorithm] =
     &[CompressionAlgorithm::None, CompressionAlgorithm::Zlib];
 
 /// Find the best matching algorithm
-pub fn negotiate<A: PartialEq + Copy>(
-    server: &[A],
-    client: &[A],
-) -> ConnectionResult<A> {
+pub fn negotiate<A: PartialEq + Copy>(server: &[A], client: &[A])
+    -> ConnectionResult<A> {
     for algorithm in client.iter() {
         if server.iter().any(|a| a == algorithm) {
             return Ok(*algorithm);

src/connection.rs

@@ -1,6 +1,8 @@
+use std::borrow::BorrowMut;
 use std::io::{self, BufRead, BufReader, Read, Write};
 use std::sync::Arc;
 
+use encryption::{AesCtr, Decryptor, Encryption};
 use error::{ConnectionError, ConnectionResult};
 use key_exchange::{self, KexResult, KeyExchange};
 use message::MessageType;
@@ -35,6 +37,7 @@ pub struct Connection {
     key_exchange: Option<Box<KeyExchange>>,
     stream: Box<Write>,
     session_id: Option<Vec<u8>>,
+    encryption: Option<(Box<Encryption>, Box<Encryption>)>,
 }
 
 impl<'a> Connection {
@@ -46,6 +49,7 @@ impl<'a> Connection {
             key_exchange: None,
             stream: Box::new(stream),
             session_id: None,
+            encryption: None,
         }
     }
 
@@ -56,7 +60,14 @@ impl<'a> Connection {
         self.read_id(&mut reader)?;
 
         loop {
-            let packet = Packet::read_from(&mut reader)?;
+            let packet = if let Some((ref mut c2s, _)) = self.encryption {
+                println!("decrypting!!!");
+                let mut decryptor = Decryptor::new(&mut **c2s, &mut reader);
+                Packet::read_from(&mut decryptor)?
+            }
+            else {
+                Packet::read_from(&mut reader)?
+            };
             trace!("Packet received: {:?}", packet);
             self.process(packet)?;
         }
@@ -96,11 +107,8 @@ impl<'a> Connection {
         Ok(())
     }
 
-    fn generate_key(
-        &mut self,
-        id: &[u8],
-        len: usize,
-    ) -> ConnectionResult<Vec<u8>> {
+    fn generate_key(&mut self, id: &[u8], len: usize)
+        -> ConnectionResult<Vec<u8>> {
         use self::ConnectionError::KeyGenerationError;
 
         let kex = self.key_exchange.take().ok_or(KeyGenerationError)?;
@@ -139,7 +147,15 @@ impl<'a> Connection {
                 let int_c2s = self.generate_key(b"E", 256)?;
                 let int_s2c = self.generate_key(b"F", 256)?;
 
-                println!("c2s enc key: {:?}", enc_c2s);
+                self.encryption =
+                    Some((
+                        Box::new(
+                            AesCtr::new(enc_c2s.as_slice(), iv_c2s.as_slice()),
+                        ),
+                        Box::new(
+                            AesCtr::new(enc_s2c.as_slice(), iv_s2c.as_slice()),
+                        ),
+                    ));
 
                 Ok(())
             }

src/encryption/aes_ctr.rs

@@ -0,0 +1,24 @@
+use crypto::aes::{KeySize, ctr};
+use crypto::symmetriccipher::SynchronousStreamCipher;
+
+use encryption::Encryption;
+
+pub struct AesCtr {
+    cipher: Box<SynchronousStreamCipher + 'static>,
+}
+
+impl AesCtr {
+    pub fn new(key: &[u8], iv: &[u8]) -> AesCtr {
+        AesCtr { cipher: ctr(KeySize::KeySize256, key, iv) }
+    }
+}
+
+impl Encryption for AesCtr {
+    fn encrypt(&mut self, data: &[u8], buf: &mut [u8]) {
+        self.cipher.process(data, buf);
+    }
+
+    fn decrypt(&mut self, data: &[u8], buf: &mut [u8]) {
+        self.encrypt(data, buf);
+    }
+}

src/encryption/mod.rs

@@ -0,0 +1,34 @@
+use std::io::{self, Read};
+
+mod aes_ctr;
+
+pub use self::aes_ctr::AesCtr;
+
+pub trait Encryption {
+    fn encrypt(&mut self, data: &[u8], buf: &mut [u8]);
+    fn decrypt(&mut self, data: &[u8], buf: &mut [u8]);
+}
+
+pub struct Decryptor<'a> {
+    encryption: &'a mut Encryption,
+    stream: &'a mut Read,
+}
+
+impl<'a> Decryptor<'a> {
+    pub fn new(encryption: &'a mut Encryption, stream: &'a mut Read)
+        -> Decryptor<'a> {
+        Decryptor {
+            encryption: encryption,
+            stream: stream,
+        }
+    }
+}
+
+impl<'a> Read for Decryptor<'a> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let mut tmp = vec![0; buf.len()];
+        self.stream.read(tmp.as_mut_slice())?;
+        self.encryption.decrypt(tmp.as_slice(), buf);
+        Ok(buf.len())
+    }
+}

src/lib.rs

@@ -11,6 +11,7 @@ mod packet;
 mod message;
 mod connection;
 mod key_exchange;
+mod encryption;
 
 pub mod public_key;
 pub mod server;

src/packet.rs

@@ -71,10 +71,8 @@ impl Packet {
         &mut self.payload
     }
 
-    pub fn with_writer(
-        &mut self,
-        f: &Fn(&mut Write) -> Result<()>,
-    ) -> Result<()> {
+    pub fn with_writer(&mut self, f: &Fn(&mut Write) -> Result<()>)
+        -> Result<()> {
         f(&mut self.payload)
     }