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