From db8678377c7ea4fa7183a54c1bfe095e051882a8 Mon Sep 17 00:00:00 2001 From: Ummon Date: Fri, 31 Oct 2014 10:11:16 +0100 Subject: [PATCH] Change the structure of a packet. --- src/client.rs | 24 +++--- src/command.rs | 156 ------------------------------------ src/main.rs | 11 +-- src/packet.rs | 212 +++++++++++++++++++++++++++++++++++++++++++++++++ src/server.rs | 30 +++++-- 5 files changed, 254 insertions(+), 179 deletions(-) delete mode 100644 src/command.rs create mode 100644 src/packet.rs diff --git a/src/client.rs b/src/client.rs index 575df05..14df2d4 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,5 +1,5 @@ use std::io::{ TcpStream, IoResult }; -use command::{ Command, Packet, Error }; +use packet::{ Command, Packet }; pub struct Client { socket: TcpStream, @@ -14,19 +14,23 @@ impl Client { }) } - pub fn close(&mut self) { - self.socket.close_read(); - self.socket.close_write(); + pub fn close(&mut self) -> IoResult<()> { + try!(self.socket.close_read()); + try!(self.socket.close_write()); + Ok(()) } pub fn start_tests(&mut self) { - - let command = Packet::new_random_command(self.current_timestamp); - self.send(command); - + self.current_timestamp += 1; + let timestamp = self.current_timestamp; + self.send(Packet { timestamp: timestamp, t: Command(Packet::random_packet_data()) }); } - fn send(&mut self, p: Packet) { - p.write(&mut self.socket); + fn send(&mut self, p: Packet) { + println!("Client sends packet: {}", p); + match p.write(&mut self.socket) { + Err(e) => println!("Error on sending: {}", e), + _ => () + } } } \ No newline at end of file diff --git a/src/command.rs b/src/command.rs deleted file mode 100644 index 31b7235..0000000 --- a/src/command.rs +++ /dev/null @@ -1,156 +0,0 @@ -use std::io; -use std::rand::{ random, task_rng, distributions }; -use std::rand::distributions::IndependentSample; -use crypto; - -pub enum ReadingError { - IOReadError(io::IoError), - UnknownPacketTypeReadError, // If the first byte is unknown. - UnconsistentDataSizeError, // The payload is too small or too big. -} - -macro_rules! try_read_io( - ($e:expr) => (match $e { Ok(e) => e, Err(e) => return Err(IOReadError(e)) }) -) - -pub enum WritingError { - WriteIOError(io::IoError) - // TODO... -} - -macro_rules! try_write_io( - ($e:expr) => (match $e { Ok(e) => e, Err(e) => return Err(WriteIOError(e)) }) -) - -pub type ReadingResult = Result; -pub type WritingResult = Result<(), WritingError>; - -static MIN_PAYLOAD_SIZE: uint = 6; -static MAX_PAYLOAD_SIZE: uint = 38; - -pub struct CommandPacket { - timestamp: u64, - id: u8, - payload: Vec // May vary from 'MIN_PAYLOAD_SIZE' to 'MAX_PAYLOAD_SIZE' bytes. -} - -pub enum ErrorType { - DecryptError, - AuthError -} - -pub struct ErrorPacket { - t: ErrorType, - timestamp: u64 -} - -pub enum Packet { - Command(CommandPacket), - Error(ErrorPacket) -} - -impl Packet { - pub fn new_random_command(timestamp: u64) -> Packet { - let mut rng = task_rng(); - Command(CommandPacket { - timestamp: timestamp, - id: random::(), - payload: Vec::from_fn(distributions::Range::new(MIN_PAYLOAD_SIZE, MAX_PAYLOAD_SIZE + 1).ind_sample(&mut rng), |idx| random::()) - }) - } - - pub fn write(&self, output: &mut io::Writer) -> WritingResult { - match self { - &Command(ref command) => { - // Data. - let mut data = Vec::with_capacity(command.payload.len() + 1); - data.push(command.id); - data.push(command.payload.len() as u8); - data.push_all(command.payload.as_slice()); - - // Padding. - let padding_size = if data.len() % 16 == 0 { 16 } else { data.len() % 16 } ; - let padding = Vec::from_elem(padding_size, padding_size as u8); - data.push_all(padding.as_slice()); - - // Encrypt. - let encrypted_data = crypto::encrypt(data.as_slice(), iv_from_timestamp(command.timestamp).as_slice()); - - // Data size. - try_write_io!(output.write_be_u16(encrypted_data.len() as u16)); - - // Packet type. - try_write_io!(output.write_u8(0)); - - // Timestamp. - try_write_io!(output.write_be_u64(command.timestamp)); - - // Write encrypted data. - try_write_io!(output.write(encrypted_data.as_slice())); - - // MAC. - try_write_io!(output.write(crypto::compute_mac(data.as_slice()))); - }, - &Error(error) => { - // Padding as data: 16 * '0'. - let padding = Vec::from_elem(16, '0' as u8); - let encrypted_data = crypto::encrypt(padding.as_slice(), iv_from_timestamp(error.timestamp).as_slice()); - - // Data size. - try_write_io!(output.write_be_u16(encrypted_data.len() as u16)); - - // Error type. - try_write_io!(match error.t { - DecryptError => output.write_u8(0x0A), - AuthError => output.write_u8(0x0B) - }); - - // Timestamp. - try_write_io!(output.write_be_u64(error.timestamp)); - - // Encrypt the padding. - try_write_io!(output.write(encrypted_data.as_slice())); - - // MAC. - try_write_io!(output.write(crypto::compute_mac(padding.as_slice()))); - } - } - - Ok(()) - } - - pub fn read(input: &mut io::Reader) -> ReadingResult { - let data_size = try_read_io!(input.read_be_u16()); - - match try_read_io!(input.read_byte()) { - 0 => { - let mut command = CommandPacket { timestamp: try_read_io!(input.read_be_u64()), id: 0, payload: vec!() }; - let mut encrypted_data = Vec::from_elem(data_size as uint, 0u8); - if try_read_io!(input.read(encrypted_data.as_mut_slice_())) != data_size as uint { - return Err(IOReadError(io::IoError { kind: io::EndOfFile, desc: "Unable to read all encrypted data", detail: None })); - } - - let data = crypto::decrypt(encrypted_data.as_slice(), iv_from_timestamp(command.timestamp).as_slice()); - if data.len() < MIN_PAYLOAD_SIZE + 2 || data.len() > MAX_PAYLOAD_SIZE + 2 { - return Err(UnconsistentDataSizeError) - } - - Ok(Command(command)) - }, - t if t == 0x0A || t == 0x0B => { - let mut error = ErrorPacket { t: if t == 0x0A { DecryptError } else { AuthError }, timestamp: try_read_io!(input.read_be_u64()) }; - // TODO - Ok(Error(error)) - }, - _ => Err(UnknownPacketTypeReadError), - } - } -} - -fn iv_from_timestamp(t: u64) -> Vec { - // Initialization vector = 64 * 0u8 + timestamp. - let mut iv = io::MemWriter::with_capacity(16); - iv.write_be_u64(0u64); - iv.write_be_u64(t); - iv.unwrap() -} diff --git a/src/main.rs b/src/main.rs index cec3957..fbffa21 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,7 +6,7 @@ use std::io; use std::os; mod crypto; -mod command; +mod packet; mod client; mod server; @@ -19,7 +19,7 @@ TODO const PORT: u16 = 4221; fn print_usage() { - println!("{} | ...", os::args()[0]); + println!("{} | | ...", os::args()[0]); } fn main() { @@ -32,7 +32,9 @@ fn main() { Ok(key) => println!("key: {}", key), Err(e) => println!("Unable to generate a key. Error: {}", e) } - } else { + } else { + println!("Starting server.., Press any key to quit"); + match server::Server::new(PORT) { Ok(mut server) => { println!("Server started"); @@ -40,7 +42,7 @@ fn main() { match client::Client::new("127.0.0.1", PORT) { Ok(mut client) => { client.start_tests(); - client.close(); + let _ = client.close(); }, Err(e) => { println!("Unable to create a client. Error: {}", e); @@ -48,7 +50,6 @@ fn main() { } } - println!("Press any key to quit"); io::stdin().read_line().ok().expect("Failed to read line"); server.close().ok().expect("Failed to close the server"); }, diff --git a/src/packet.rs b/src/packet.rs new file mode 100644 index 0000000..49ae2be --- /dev/null +++ b/src/packet.rs @@ -0,0 +1,212 @@ +use std::io; +use std::fmt; +use std::rand::{ random, task_rng, distributions }; +use std::rand::distributions::IndependentSample; +use crypto; + +// There are all the errors that may occur when reading an encrypted and authenticated packet. +#[deriving(Show)] +pub enum ReadingError { + IOReadError(io::IoError), + UnknownPacketTypeReadError, // If the first byte is unknown. + UnconsistentEncryptedSizeError, + UnconsistentDataSizeError, // The data size is not valid. + UnconsistentMACSizeError, // The MAC hasn't the correct size. + MACMismatchError, // The uncrypted received data doesn't match to the received MAC. + PaddingError, // Padding format error. + DataError, // The data are invalid. +} + +// A macro to return a 'IOReadError' in case of error. +macro_rules! try_read_io( + ($e:expr) => (match $e { Ok(e) => e, Err(e) => return Err(IOReadError(e)) }) +) + +// There are all the errors that may occur when encrypting, authenticating and writing a packet. +#[deriving(Show)] +pub enum WritingError { + WriteIOError(io::IoError) + // TODO... +} + +// A macro to return a 'IOWritingError' in case of error. +macro_rules! try_write_io( + ($e:expr) => (match $e { Ok(e) => e, Err(e) => return Err(WriteIOError(e)) }) +) + +pub type ReadingResult = Result; +pub type WritingResult = Result<(), WritingError>; + +static MIN_PAYLOAD_SIZE: uint = 7; +static MAX_PAYLOAD_SIZE: uint = 39; +static FIXED_PACKET_SIZE: uint = 1 + 8 + 10; // Packet type + timestamp + MAC. + +#[deriving(Show)] +pub struct PacketData { + id: u8, + payload: Vec // The size can vary from 'MIN_PAYLOAD_SIZE' to 'MAX_PAYLOAD_SIZE' bytes. +} + +#[deriving(Show)] +pub enum ErrorType { + CryptError, + AuthError +} + +#[deriving(Show)] +pub enum PacketType { + Command(PacketData), + Answer(PacketData), + Error(ErrorType), +} + +pub struct Packet { + pub t: PacketType, + pub timestamp: u64 +} + +impl fmt::Show for Packet { + fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + fn data_to_str(data: &PacketData) -> String { + format!("id: {}, payload({}): {}", data.id, data.payload.len(), data.payload) + } + match self.t { + Command(ref data) => write!(formatter, "Command packet {{ timestamp: {}, {} }}", self.timestamp, data_to_str(data)), + Answer(ref data) => write!(formatter, "Answer packet {{ timestamp: {}, {} }}", self.timestamp, data_to_str(data)), + Error(error_type) => write!(formatter, "Error packet {{ timestamp: {}, errorType: {} }}", self.timestamp, error_type) + } + } +} + +impl Packet { + pub fn random_packet_data() -> PacketData { + let mut rng = task_rng(); + PacketData { + id: random::(), + payload: Vec::from_fn(distributions::Range::new(MIN_PAYLOAD_SIZE, MAX_PAYLOAD_SIZE + 1).ind_sample(&mut rng), |_| random::()) + } + } + + pub fn write(&self, output: &mut io::Writer) -> WritingResult { + fn packet_data(p: &PacketData) -> Vec { + let mut d = Vec::new(); + d.push(p.id); + d.push_all(p.payload.as_slice()); + d + } + + // Data to be encrypted. + let mut data = + match self.t { + Command(ref p) | Answer(ref p) => packet_data(p), + Error(_) => Vec::from_elem(16, '0' as u8) // Padding as data: 16 * '0'. + }; + + // Compute the MAC + let mac = crypto::compute_mac(data.as_slice()); + + // Padding. + match self.t { + Command(_) | Answer(_) => { + let padding_size = if data.len() % 16 == 0 { 16 } else { data.len() % 16 } ; + data.push_all(Vec::from_elem(padding_size, padding_size as u8).as_slice()); + }, + _ => () + } + + // Encrypt. + let encrypted_data = crypto::encrypt(data.as_slice(), iv_from_timestamp(self.timestamp).as_slice()); + + // Write packet length. + try_write_io!(output.write_be_u16((encrypted_data.len() + FIXED_PACKET_SIZE) as u16)); + + // Write packet type. + try_write_io!(output.write_u8( + match self.t { + Command(_) => 0x00, + Answer(_) => 0xFF, + Error(CryptError) => 0x0A, + Error(AuthError) => 0x0B + } + )); + + // Write timestamp. + try_write_io!(output.write_be_u64(self.timestamp)); + + // Write encrypted data. + try_write_io!(output.write(encrypted_data.as_slice())); + + // Write the MAC. + try_write_io!(output.write(mac)); + + Ok(()) + } + + pub fn read(input: &mut io::Reader) -> ReadingResult { + let data_size = try_read_io!(input.read_be_u16()); + let packet_type = try_read_io!(input.read_u8()); + let timestamp = try_read_io!(input.read_be_u64()); + + let mut encrypted_data = Vec::from_elem(data_size as uint - FIXED_PACKET_SIZE, 0u8); + if try_read_io!(input.read(encrypted_data.as_mut_slice_())) != encrypted_data.len() { + return Err(UnconsistentEncryptedSizeError) + } + let mut data = crypto::decrypt(encrypted_data.as_slice(), iv_from_timestamp(timestamp).as_slice()); + + // Control the size and the content of the padding then remove it. + if packet_type == 0x00 || packet_type == 0xFF { + match data.last() { + Some(&padding_size) => { + if padding_size as uint > data.len() || !data.slice_from(data.len() - padding_size as uint).iter().any(|b| *b == padding_size) { + return Err(PaddingError) + } + let data_length = data.len() - padding_size as uint; + data.truncate(data_length); + }, + None => + return Err(PaddingError) + } + } + + // Read an verify the MAC. + let mut mac_read = [0u8, ..10]; + if try_read_io!(input.read(mac_read)) != mac_read.len() { + return Err(UnconsistentMACSizeError) + } + let mac_data = crypto::compute_mac(data.as_slice()); + if mac_read != mac_data { + return Err(MACMismatchError) + } + + Ok(Packet { + t: match packet_type { + 0x00 | 0xFF => { + if data.len() < MIN_PAYLOAD_SIZE + 1 || data.len() > MAX_PAYLOAD_SIZE + 1 { + return Err(UnconsistentDataSizeError) + } + let pd = PacketData { id: data[0], payload: data.tail().to_vec() }; // match data.as_slice() { [id, payload..] => PacketData { id: id, payload: payload.to_vec() } }; + match packet_type { 0x00 => Command(pd), _ => Answer(pd) } + }, + 0x0A | 0x0B => { + if data.len() != 16 { + return Err(UnconsistentDataSizeError) + } else if data != Vec::from_elem(16, '0' as u8) { + return Err(DataError) + } + match packet_type { 0x0A => Error(CryptError), _ => Error(AuthError) } + }, + _ => + return Err(UnknownPacketTypeReadError) + }, + timestamp: timestamp + }) + } +} + +// Build an initialization vector: 64 * 0u8 + timestamp (128 bits). +fn iv_from_timestamp(timestamp: u64) -> Vec { + let mut iv = io::MemWriter::with_capacity(16); + let _ = iv.write_be_u64(0u64); + let _ = iv.write_be_u64(timestamp); + iv.unwrap() +} diff --git a/src/server.rs b/src/server.rs index 1ce5f95..b655ae7 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,6 +1,6 @@ -use std::io::{ TcpListener, TcpStream, Acceptor, Listener, IoError, IoResult }; +use std::io::{ TcpListener, TcpStream, Acceptor, Listener, IoError, IoResult, EndOfFile }; use std::io::net::tcp::{ TcpAcceptor }; -use command::{ Command, Packet, Error }; +use packet::{ Packet, IOReadError }; pub struct Server { acceptor: TcpAcceptor @@ -34,14 +34,28 @@ impl Server { self.acceptor.close_accept() } - fn handle_client(mut stream: TcpStream) { - + fn handle_client(mut stream: TcpStream) { let mut current_timestamp = 0u64; - match Packet::read(&mut stream) { - Ok(Command(packet)) => println!("CommandPacket"), - Ok(Error(packet)) => println!("AnswerPacket"), - Err(io_error) => println!("IO Error") + loop { + // stream.set_timeout(Some(1000000000)); + match Packet::read(&mut stream) { + Ok(packet) => { + if packet.timestamp <= current_timestamp { + println!("Error, timestamp mismatch, current timestamp: {}, packet received: {}", current_timestamp, packet); + } else { + current_timestamp += 1; + println!("Receive packet: {}", packet); + } + }, + // Socket has been closed. + Err(IOReadError(IoError { kind: EndOfFile, .. })) => { + return; + }, + Err(e) => { + println!("Read error: {}", e) + } + } } } } \ No newline at end of file -- 2.45.2