src/command.rs

   1 use std::io;
   2 use std::rand::{ random, task_rng, distributions };
   3 use std::rand::distributions::IndependentSample;
   4 use crypto;
   5
   6 pub enum ReadingError {
   7    IOReadError(io::IoError),
   8    UnknownPacketTypeReadError, // If the first byte is unknown.
   9    UnconsistentDataSizeError, // The payload is too small or too big.
  10 }
  11
  12 macro_rules! try_read_io(
  13    ($e:expr) => (match $e { Ok(e) => e, Err(e) => return Err(IOReadError(e)) })
  14 )
  15
  16 pub enum WritingError {
  17    WriteIOError(io::IoError)
  18    // TODO...
  19 }
  20
  21 macro_rules! try_write_io(
  22    ($e:expr) => (match $e { Ok(e) => e, Err(e) => return Err(WriteIOError(e)) })
  23 )
  24
  25 pub type ReadingResult = Result<Packet, ReadingError>;
  26 pub type WritingResult = Result<(), WritingError>;
  27
  28 static MIN_PAYLOAD_SIZE: uint = 6;
  29 static MAX_PAYLOAD_SIZE: uint = 38;
  30
  31 pub struct CommandPacket {
  32    timestamp: u64,
  33    id: u8,
  34    payload: Vec<u8> // May vary from 'MIN_PAYLOAD_SIZE' to 'MAX_PAYLOAD_SIZE' bytes.
  35 }
  36
  37 pub enum ErrorType {
  38    DecryptError,
  39    AuthError
  40 }
  41
  42 pub struct ErrorPacket {
  43    t: ErrorType,
  44    timestamp: u64
  45 }
  46
  47 pub enum Packet {
  48    Command(CommandPacket),
  49    Error(ErrorPacket)
  50 }
  51
  52 impl Packet {
  53    pub fn new_random_command(timestamp: u64) -> Packet {
  54       let mut rng = task_rng();
  55       Command(CommandPacket {
  56          timestamp: timestamp,
  57          id: random::<u8>(),
  58          payload: Vec::from_fn(distributions::Range::new(MIN_PAYLOAD_SIZE, MAX_PAYLOAD_SIZE + 1).ind_sample(&mut rng), |idx| random::<u8>())
  59       })
  60    }
  61
  62    pub fn write(&self, output: &mut io::Writer) ->  WritingResult {
  63       match self {
  64          &Command(ref command) => {
  65             // Data.
  66             let mut data = Vec::with_capacity(command.payload.len() + 1);
  67             data.push(command.id);
  68             data.push(command.payload.len() as u8);
  69             data.push_all(command.payload.as_slice());
  70
  71             // Padding.
  72             let padding_size = if data.len() % 16 == 0 { 16 } else { data.len() % 16 } ;
  73             let padding = Vec::from_elem(padding_size, padding_size as u8);
  74             data.push_all(padding.as_slice());
  75
  76             // Encrypt.
  77             let encrypted_data = crypto::encrypt(data.as_slice(), iv_from_timestamp(command.timestamp).as_slice());
  78
  79             // Data size.
  80             try_write_io!(output.write_be_u16(encrypted_data.len() as u16));
  81
  82             // Packet type.
  83             try_write_io!(output.write_u8(0));
  84
  85             // Timestamp.
  86             try_write_io!(output.write_be_u64(command.timestamp));
  87
  88             // Write encrypted data.
  89             try_write_io!(output.write(encrypted_data.as_slice()));
  90
  91             // MAC.
  92             try_write_io!(output.write(crypto::compute_mac(data.as_slice())));
  93          },
  94          &Error(error) => {
  95             // Padding as data: 16 * '0'.
  96             let padding = Vec::from_elem(16, '0' as u8);
  97             let encrypted_data = crypto::encrypt(padding.as_slice(), iv_from_timestamp(error.timestamp).as_slice());
  98
  99             // Data size.
 100             try_write_io!(output.write_be_u16(encrypted_data.len() as u16));
 101
 102             // Error type.
 103             try_write_io!(match error.t {
 104                DecryptError => output.write_u8(0x0A),
 105                AuthError => output.write_u8(0x0B)
 106             });
 107
 108             // Timestamp.
 109             try_write_io!(output.write_be_u64(error.timestamp));
 110
 111             // Encrypt the padding.
 112             try_write_io!(output.write(encrypted_data.as_slice()));
 113
 114             // MAC.
 115             try_write_io!(output.write(crypto::compute_mac(padding.as_slice())));
 116          }
 117       }
 118
 119       Ok(())
 120    }
 121
 122    pub fn read(input: &mut io::Reader) -> ReadingResult {
 123       let data_size = try_read_io!(input.read_be_u16());
 124
 125       match try_read_io!(input.read_byte()) {
 126          0 => {
 127             let mut command = CommandPacket { timestamp: try_read_io!(input.read_be_u64()), id: 0, payload: vec!() };
 128             let mut encrypted_data = Vec::from_elem(data_size as uint, 0u8);
 129             if try_read_io!(input.read(encrypted_data.as_mut_slice_())) != data_size as uint {
 130                return Err(IOReadError(io::IoError { kind: io::EndOfFile, desc: "Unable to read all encrypted data", detail: None }));
 131             }
 132
 133             let data = crypto::decrypt(encrypted_data.as_slice(), iv_from_timestamp(command.timestamp).as_slice());
 134             if data.len() < MIN_PAYLOAD_SIZE + 2 || data.len() > MAX_PAYLOAD_SIZE + 2 {
 135                return Err(UnconsistentDataSizeError)
 136             }
 137
 138             Ok(Command(command))
 139          },
 140          t if t == 0x0A || t == 0x0B  => {
 141             let mut error = ErrorPacket { t: if t == 0x0A { DecryptError } else { AuthError }, timestamp: try_read_io!(input.read_be_u64()) };
 142             // TODO
 143             Ok(Error(error))
 144          },
 145          _ => Err(UnknownPacketTypeReadError),
 146       }
 147    }
 148 }
 149
 150 fn iv_from_timestamp(t: u64) -> Vec<u8> {
 151    // Initialization vector = 64 * 0u8 + timestamp.
 152    let mut iv = io::MemWriter::with_capacity(16);
 153    iv.write_be_u64(0u64);
 154    iv.write_be_u64(t);
 155    iv.unwrap()
 156 }