Add some tests.
[crypto_lab1.git] / src / packet.rs
index 74bf102..ece32c4 100644 (file)
@@ -10,11 +10,11 @@ use crypto;
 pub enum ReadingError {
    IOReadError(io::IoError),
    UnknownPacketTypeError, // If the first byte is unknown.
-   UnconsistentEncryptedSizeError, 
+   UnconsistentEncryptedSizeError,
    UnconsistentDataSizeError, // The data size is not valid.
    UnconsistentMACSizeError, // The MAC hasn't the correct size.
-   MACMismatchError, // The uncrypted received data doesn't match to the received MAC. 
-   PaddingError, // Padding format error.   
+   MACMismatchError, // The uncrypted received data doesn't match to the received MAC.
+   PaddingError, // Padding format error.
    DataError, // The data are invalid.
    InvalidTimestampError
 }
@@ -104,94 +104,110 @@ impl Packet {
       let mut rng = if seed.is_empty() { StdRng::new().unwrap() } else { SeedableRng::from_seed(seed) };
       let mut payload = Vec::from_elem(distributions::Range::new(MIN_PAYLOAD_SIZE, MAX_PAYLOAD_SIZE + 1).ind_sample(&mut rng), 0u8);
       rng.fill_bytes(payload.as_mut_slice_());
-      PacketData { 
+      PacketData {
          id: rng.gen::<u8>(),
          payload: payload
       }
    }
 
+   pub fn new_packet_data(id: u8, payload: Vec<u8>) -> PacketData {
+      PacketData { id: id, payload: payload }
+   }
+
    pub fn write(&self, output: &mut io::Writer) -> WritingResult {
+      self.write_with_padding_fun(output, |_: uint, padding_length: uint| -> u8 {
+         padding_length as u8
+      })
+   }
+
+   /// 'padd_fun' is function to fill the padding. The first argument is the index of the current byte, starting at 0.
+   /// The second argument is the padding length.
+   pub fn write_with_padding_fun(&self, output: &mut io::Writer, padd_fun: |uint, uint| -> u8) -> WritingResult {
       fn packet_data(p: &PacketData) -> Vec<u8> {
          let mut d = Vec::new();
          d.push(p.id);
          d.push_all(p.payload.as_slice());
          d
       }
-      
-      // Data to be encrypted.    
-      let mut data = 
+
+      // Data to be encrypted.
+      let mut data =
          match self.t {
             Command(ref p) | Answer(ref p) => packet_data(p),
             Error(_) => Vec::from_elem(16, '0' as u8) // Padding as data: 16 * '0'.
          };
-         
+
       // Compute the MAC
-      let mac = crypto::compute_mac(data.as_slice());        
-   
+      let mac = crypto::compute_mac(data.as_slice());
+
       // Padding.
       match self.t {
          Command(_) | Answer(_) => {
             let padding_size = if data.len() % 16 == 0 { 16 } else { data.len() % 16 } ;
-            data.push_all(Vec::from_elem(padding_size, padding_size as u8).as_slice());
+            data.reserve_additional(padding_size);
+            for i in range(0, padding_size) {
+               data.push(padd_fun(i, padding_size));
+            }
          },
          _ => ()
       }
-         
+
       // Encrypt.
-      let encrypted_data = crypto::encrypt(data.as_slice(), iv_from_timestamp(self.timestamp).as_slice());         
-         
+      let encrypted_data = crypto::encrypt(data.as_slice(), iv_from_timestamp(self.timestamp).as_slice());
+
       // Write packet length.
       try_write_io!(output.write_be_u16((encrypted_data.len() + FIXED_PACKET_SIZE) as u16));
-      
+
       // Write packet type.
       try_write_io!(output.write_u8(
          match self.t {
             Command(_) => 0x00,
             Answer(_) => 0xFF,
             Error(CryptError) => 0x0A,
-            Error(AuthError) => 0x0B 
+            Error(AuthError) => 0x0B
          }
       ));
-      
+
       // Write timestamp.
       try_write_io!(output.write_be_u64(self.timestamp));
-      
+
       // Write encrypted data.
       try_write_io!(output.write(encrypted_data.as_slice()));
-      
+
       // Write the MAC.
       try_write_io!(output.write(mac));
-         
+
       Ok(())
    }
-   
+
    pub fn read(input: &mut io::Reader) -> ReadingResult {
       fn consume(input: &mut io::Reader, nb_byte: uint) {
-         let _ = input.read_exact(nb_byte);        
+         let _ = input.read_exact(nb_byte);
       }
-   
+
       let data_size = try_read_io!(input.read_be_u16());
-     
+
       // Read and check the packet type.
-      let packet_type = try_read_io!(input.read_u8());            
+      let packet_type = try_read_io!(input.read_u8());
       if ![0x00, 0xFF, 0x0A, 0x0B].iter().any(|p| *p == packet_type) {
          consume(input, data_size as uint - 1);
          return Err(UnknownPacketTypeError)
-      }      
-      
+      }
+
       let timestamp = try_read_io!(input.read_be_u64());
-      
-      let mut encrypted_data = Vec::from_elem(data_size as uint - FIXED_PACKET_SIZE, 0u8);      
+
+      let mut encrypted_data = Vec::from_elem(data_size as uint - FIXED_PACKET_SIZE, 0u8);
       if try_read_io!(input.read(encrypted_data.as_mut_slice_())) != encrypted_data.len() {
          return Err(UnconsistentEncryptedSizeError)
       }
       let mut data = crypto::decrypt(encrypted_data.as_slice(), iv_from_timestamp(timestamp).as_slice());
-                     
-      // Control the size and the content of the padding then remove it.      
+
+      // Control the size and the content of the padding then remove it.
       if packet_type == 0x00 || packet_type == 0xFF {
          match data.last() {
-            Some(&padding_size) => {            
+            Some(&padding_size) => {
                if padding_size as uint > data.len() || !data.slice_from(data.len() - padding_size as uint).iter().any(|b| *b == padding_size) {
+                  consume(input, 10);
                   return Err(PaddingError)
                }
                let data_length = data.len() - padding_size as uint;
@@ -201,7 +217,7 @@ impl Packet {
                return Err(PaddingError)
          }
       }
-            
+
       // Read an verify the MAC.
       let mut mac_read = [0u8, ..10];
       if try_read_io!(input.read(mac_read)) != mac_read.len() {
@@ -211,7 +227,7 @@ impl Packet {
       if mac_read != mac_data {
          return Err(MACMismatchError)
       }
-      
+
       Ok(Packet {
          t: match packet_type {
             // Command or answer.
@@ -223,13 +239,13 @@ impl Packet {
                match packet_type { 0x00 => Command(pd), _ => Answer(pd) }
             },
             // Error.
-            _ => { 
+            _ => {
                if data.len() != 16 {
                   return Err(UnconsistentDataSizeError)
-               } else if data != Vec::from_elem(16, '0' as u8) {               
+               } else if data != Vec::from_elem(16, '0' as u8) {
                   return Err(DataError)
                }
-               match packet_type { 0x0A => Error(CryptError), _ => Error(AuthError) }               
+               match packet_type { 0x0A => Error(CryptError), _ => Error(AuthError) }
             }
          },
          timestamp: timestamp