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