From 64e44de13a138382dd02057422cb33dac0988278 Mon Sep 17 00:00:00 2001 From: Thomas Gatzweiler Date: Tue, 18 Jul 2017 15:41:33 +0200 Subject: [PATCH] Initial AES-CTR implementation --- .rustfmt.toml | 1 + src/algorithm.rs | 6 ++---- src/connection.rs | 30 +++++++++++++++++++++++------- src/encryption/aes_ctr.rs | 24 ++++++++++++++++++++++++ src/encryption/mod.rs | 34 ++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/packet.rs | 6 ++---- 7 files changed, 87 insertions(+), 15 deletions(-) create mode 100644 src/encryption/aes_ctr.rs create mode 100644 src/encryption/mod.rs diff --git a/.rustfmt.toml b/.rustfmt.toml index 67befda..e775cae 100644 --- a/.rustfmt.toml +++ b/.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 diff --git a/src/algorithm.rs b/src/algorithm.rs index 10d2164..870434d 100644 --- a/src/algorithm.rs +++ b/src/algorithm.rs @@ -28,10 +28,8 @@ pub static COMPRESSION: &[CompressionAlgorithm] = &[CompressionAlgorithm::None, CompressionAlgorithm::Zlib]; /// Find the best matching algorithm -pub fn negotiate( - server: &[A], - client: &[A], -) -> ConnectionResult { +pub fn negotiate(server: &[A], client: &[A]) + -> ConnectionResult { for algorithm in client.iter() { if server.iter().any(|a| a == algorithm) { return Ok(*algorithm); diff --git a/src/connection.rs b/src/connection.rs index dbd27d6..355899e 100644 --- a/src/connection.rs +++ b/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>, stream: Box, session_id: Option>, + encryption: Option<(Box, Box)>, } 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> { + fn generate_key(&mut self, id: &[u8], len: usize) + -> ConnectionResult> { 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(()) } diff --git a/src/encryption/aes_ctr.rs b/src/encryption/aes_ctr.rs new file mode 100644 index 0000000..17f6bb2 --- /dev/null +++ b/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, +} + +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); + } +} diff --git a/src/encryption/mod.rs b/src/encryption/mod.rs new file mode 100644 index 0000000..66f921e --- /dev/null +++ b/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 { + let mut tmp = vec![0; buf.len()]; + self.stream.read(tmp.as_mut_slice())?; + self.encryption.decrypt(tmp.as_slice(), buf); + Ok(buf.len()) + } +} diff --git a/src/lib.rs b/src/lib.rs index 937b89e..ad00053 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/packet.rs b/src/packet.rs index 3ee9076..8ebf036 100644 --- a/src/packet.rs +++ b/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) }