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