-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<Packet, ReadingError>;
-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<u8> // 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::<u8>(),
- payload: Vec::from_fn(distributions::Range::new(MIN_PAYLOAD_SIZE, MAX_PAYLOAD_SIZE + 1).ind_sample(&mut rng), |idx| random::<u8>())
- })
- }
-
- 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<u8> {
- // 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()
-}