lab1_rust/src/packet.rs

   1 use std::io;
   2 use std::fmt;
   3 use std::rand::{ Rng, StdRng, SeedableRng, distributions };
   4 use std::rand::distributions::IndependentSample;
   5 use serialize::hex::{ ToHex };
   6 use self::PacketType::{ Command, Answer, Error };
   7 use crypto;
   8 use utils::from_elem;
   9
  10 #[derive(Show, Copy)]
  11 pub enum Variant {
  12    Weak, // The MAC is computed on data without padding.
  13    Fixed // The MAC is computed on data and padding.
  14 }
  15
  16 // There are all the errors that may occur when reading an encrypted and authenticated packet.
  17 #[derive(Show)]
  18 pub enum ReadingError {
  19    IO(io::IoError),
  20    UnknownPacketType, // If the first byte is unknown.
  21    UnconsistentEncryptedSize,
  22    UnconsistentDataSize, // The data size is not valid.
  23    UnconsistentMACSize, // The MAC hasn't the correct size.
  24    MACMismatch, // The uncrypted received data doesn't match to the received MAC.
  25    Padding, // Padding format error.
  26    Data, // The data are invalid.
  27    InvalidTimestamp
  28 }
  29
  30 // A macro to return a 'Err(ReadingError::IO(..))' in case of error.
  31 macro_rules! try_read_io(
  32    ($e:expr) => (match $e { Ok(e) => e, Err(e) => return Err(ReadingError::IO(e)) })
  33 );
  34
  35 // There are all the errors that may occur when encrypting, authenticating and writing a packet.
  36 #[derive(Show)]
  37 pub enum WritingError {
  38    IO(io::IoError),
  39    Encrypt,
  40 }
  41
  42 // A macro to return a 'Err(WritingError::IO(..))' in case of error.
  43 macro_rules! try_write_io(
  44    ($e:expr) => (match $e { Ok(e) => e, Err(e) => return Err(WritingError::IO(e)) })
  45 );
  46
  47 pub type ReadingResult = Result<Packet, ReadingError>;
  48 pub type WritingResult = Result<(), WritingError>;
  49
  50 static MIN_PAYLOAD_SIZE: uint = 7;
  51 static MAX_PAYLOAD_SIZE: uint = 39;
  52 static FIXED_PACKET_SIZE: uint = 1 + 8 + 10; // Packet type + timestamp + MAC.
  53
  54 #[derive(Show, Clone)]
  55 pub struct PacketData {
  56    id: u8,
  57    payload: Vec<u8> // The size can vary from 'MIN_PAYLOAD_SIZE' to 'MAX_PAYLOAD_SIZE' bytes.
  58 }
  59
  60 #[derive(Show, Clone)]
  61 pub enum ErrorType {
  62    Crypt,
  63    Auth
  64 }
  65
  66 #[derive(Clone)]
  67 pub enum PacketType {
  68    Command(PacketData),
  69    Answer(PacketData),
  70    Error(ErrorType),
  71 }
  72
  73 /// Serialized packet format : |LL|P|TTTTTTTT|D...D|MMMMMMMMMM|
  74 /// Where:
  75 ///   LL: Size of the following data
  76 ///   P: Packet type:
  77 ///      0x00: Command
  78 ///      OxFF: Answer
  79 ///      0x0A: Decrypt error
  80 ///      0x0B: Authentication error
  81 ///   TTTTTTTT: Timestamp (64 bits)
  82 ///   D...D: Encrypted data (AES-256 CBC mode) of:
  83 ///      |I|C...C|P...P| for command and answer packet:
  84 ///         I: Command ID
  85 ///         C: Command payload (from 7 to 39 bytes)
  86 ///         P: Padding from 1 to 16, |I|C...C|P...P| size must be a multiple of 16
  87 ///      |0000000000000000| for error packet (16 bytes length)
  88 ///   MMMMMMMMMM: first 10 bytes (most significant) of the HMAC-SHA256 of:
  89 ///      for command and answer packet:
  90 ///         |I|C...C| for weak variant
  91 ///         |I|C...C|P...P|for fixed variant
  92 ///      |0000000000000000| for error packet
  93 #[derive(Show)]
  94 pub struct Packet {
  95    pub t: PacketType,
  96    pub timestamp: u64
  97 }
  98
  99 impl fmt::Show for PacketType {
 100    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
 101       fn data_to_str(data: &PacketData) -> String {
 102          format!("id: {}, payload({}): \"{}\"", data.id, data.payload.len(), data.payload.as_slice().to_hex())
 103       }
 104       match self {
 105          &Command(ref data) => write!(formatter, "Command {{ {} }}", data_to_str(data)),
 106          &Answer(ref data) => write!(formatter, "Answer {{ {} }}", data_to_str(data)),
 107          &Error(ref error_type) => write!(formatter, "Error {{ errorType: {} }}", error_type)
 108       }
 109    }
 110 }
 111
 112 impl Packet {
 113    pub fn random_packet_data(seed: &[uint]) -> PacketData {
 114       let mut rng = if seed.is_empty() { StdRng::new().unwrap() } else { SeedableRng::from_seed(seed) };
 115       let mut payload = from_elem(distributions::Range::new(MIN_PAYLOAD_SIZE, MAX_PAYLOAD_SIZE + 1).ind_sample(&mut rng), 0u8);
 116       rng.fill_bytes(payload.as_mut_slice());
 117       PacketData {
 118          id: rng.gen::<u8>(),
 119          payload: payload
 120       }
 121    }
 122
 123    pub fn new_packet_data(id: u8, payload: Vec<u8>) -> PacketData {
 124       PacketData { id: id, payload: payload }
 125    }
 126
 127    pub fn write(&self, output: &mut io::Writer, variant: Variant) -> WritingResult {
 128       self.write_with_padding_fun(output, variant, |_, padding_length: uint| -> u8 {
 129          padding_length as u8
 130       })
 131    }
 132
 133    /// 'padd_fun' is function defining the padding. The first argument is the index of the current byte, starting at 0.
 134    /// The second argument is the padding length.
 135    pub fn write_with_padding_fun(&self, output: &mut io::Writer, variant: Variant, padd_fun: |uint, uint| -> u8) -> WritingResult {
 136       fn packet_data(p: &PacketData) -> Vec<u8> {
 137          let mut d = Vec::new();
 138          d.push(p.id);
 139          d.push_all(p.payload.as_slice());
 140          d
 141       }
 142
 143       // Data to be encrypted.
 144       let mut data =
 145          match self.t {
 146             Command(ref p) | Answer(ref p) => packet_data(p),
 147             Error(_) => from_elem(16, 0) // Padding as data: 16 * 0.
 148          };
 149
 150       let data_size = data.len();
 151
 152       // Padding.
 153       match self.t {
 154          Command(_) | Answer(_) => {
 155             let padding_size = if data.len() % 16 == 0 { 16 } else { 16 - data.len() % 16 } ;
 156             data.reserve(padding_size);
 157             for i in range(0, padding_size) {
 158                data.push(padd_fun(i, padding_size));
 159             }
 160          },
 161          _ => ()
 162       }
 163
 164       // Computes the MAC. It depends of the choosen variant.
 165       let mac = crypto::compute_mac(data.slice_to(match variant { Variant::Weak => data_size, _ => data.len() }));
 166
 167       // Encrypts.
 168       let encrypted_data = match crypto::encrypt(data.as_slice(), iv_from_timestamp(self.timestamp).as_slice()) {
 169          Some(d) => d,
 170          _ => return Err(WritingError::Encrypt)
 171       };
 172
 173       // Writes packet length.
 174       try_write_io!(output.write_be_u16((encrypted_data.len() + FIXED_PACKET_SIZE) as u16));
 175
 176       // Writes packet type.
 177       try_write_io!(output.write_u8(
 178          match self.t {
 179             Command(_) => 0x00,
 180             Answer(_) => 0xFF,
 181             Error(ErrorType::Crypt) => 0x0A,
 182             Error(ErrorType::Auth) => 0x0B
 183          }
 184       ));
 185
 186       // Writes timestamp.
 187       try_write_io!(output.write_be_u64(self.timestamp));
 188
 189       // Writes encrypted data.
 190       try_write_io!(output.write(encrypted_data.as_slice()));
 191
 192       // Writes the MAC.
 193       try_write_io!(output.write(&mac));
 194
 195       Ok(())
 196    }
 197
 198    pub fn read(input: &mut io::Reader, variant: Variant) -> ReadingResult {
 199       fn consume(input: &mut io::Reader, nb_byte: uint) {
 200          let _ = input.read_exact(nb_byte);
 201       }
 202
 203       let data_size = try_read_io!(input.read_be_u16());
 204
 205       // Reads and checks the packet type.
 206       let packet_type = try_read_io!(input.read_u8());
 207       if ![0x00, 0xFF, 0x0A, 0x0B].iter().any(|p| *p == packet_type) {
 208          consume(input, data_size as uint - 1);
 209          return Err(ReadingError::UnknownPacketType)
 210       }
 211
 212       let timestamp = try_read_io!(input.read_be_u64());
 213
 214       let mut encrypted_data = from_elem(data_size as uint - FIXED_PACKET_SIZE, 0u8);
 215       if try_read_io!(input.read(encrypted_data.as_mut_slice())) != encrypted_data.len() {
 216          return Err(ReadingError::UnconsistentEncryptedSize)
 217       }
 218       let mut data = match crypto::decrypt(encrypted_data.as_slice(), iv_from_timestamp(timestamp).as_slice()) {
 219          Some(d) => d,
 220          _ => return Err(ReadingError::UnconsistentEncryptedSize)
 221       };
 222
 223       // Reads the MAC.
 224       let mut mac_read = [0u8; 10];
 225       if try_read_io!(input.read(&mut mac_read)) != mac_read.len() {
 226          return Err(ReadingError::UnconsistentMACSize)
 227       }
 228
 229       match variant { Variant::Fixed if mac_read != crypto::compute_mac(data.as_slice()) => return Err(ReadingError::MACMismatch), _ => () };
 230
 231       // Controls the size and the content of the padding then removes it.
 232       if packet_type == 0x00 || packet_type == 0xFF {
 233          match data.last() {
 234             Some(&padding_size) => {
 235                if padding_size as uint > data.len() || padding_size == 0 || data.slice_from(data.len() - padding_size as uint).iter().any(|b| *b != padding_size) {
 236                   return Err(ReadingError::Padding)
 237                }
 238                let data_length = data.len() - padding_size as uint;
 239                data.truncate(data_length);
 240             },
 241             None =>
 242                return Err(ReadingError::Padding)
 243          }
 244       }
 245
 246       match variant { Variant::Weak if mac_read != crypto::compute_mac(data.as_slice()) => return Err(ReadingError::MACMismatch), _ => () };
 247
 248       Ok(Packet {
 249          t: match packet_type {
 250             // Command or answer.
 251             0x00 | 0xFF => {
 252                if data.len() < MIN_PAYLOAD_SIZE + 1 || data.len() > MAX_PAYLOAD_SIZE + 1 {
 253                   return Err(ReadingError::UnconsistentDataSize)
 254                }
 255                let pd = PacketData { id: data[0], payload: data.tail().to_vec() }; // match data.as_slice() { [id, payload..] => PacketData { id: id, payload: payload.to_vec() } };
 256                match packet_type { 0x00 => Command(pd), _ => Answer(pd) }
 257             },
 258             // Error.
 259             _ => {
 260                if data.len() != 16 {
 261                   return Err(ReadingError::UnconsistentDataSize)
 262                } else if data != from_elem(16, 0) {
 263                   return Err(ReadingError::Data)
 264                }
 265                match packet_type { 0x0A => Error(ErrorType::Crypt), _ => Error(ErrorType::Auth) }
 266             }
 267          },
 268          timestamp: timestamp
 269       })
 270    }
 271 }
 272
 273 // Builds an initialization vector: 64 * 0u8 + timestamp (128 bits).
 274 fn iv_from_timestamp(timestamp: u64) -> Vec<u8> {
 275    let mut iv = io::MemWriter::with_capacity(16);
 276    let _ = iv.write_be_u64(0u64);
 277    let _ = iv.write_be_u64(timestamp);
 278    iv.into_inner()
 279 }