Read and writing of packet, work in progress.

[crypto_lab1.git] / src / command.rs

diff --git a/src/command.rs b/src/command.rs
index 7de2afc..31b7235 100644 (file)
--- a/src/command.rs
+++ b/src/command.rs
@@ -1,33 +1,46 @@
 use std::io;
 use std::rand::{ random, task_rng, distributions };
 use std::rand::distributions::IndependentSample;
+use crypto;
 
 pub enum ReadingError {
-   ReadIOError(io::IoError)
-   // TODO...
+   IOReadError(io::IoError),
+   UnknownPacketTypeReadError, // If the first byte is unknown.
+   UnconsistentDataSizeError, // The payload is too small or too big.
 }
 
+macro_rules! try_read_io(
+   ($e:expr) => (match $e { Ok(e) => e, Err(e) => return Err(IOReadError(e)) })
+)
+
 pub enum WritingError {
    WriteIOError(io::IoError)
    // TODO...
 }
 
+macro_rules! try_write_io(
+   ($e:expr) => (match $e { Ok(e) => e, Err(e) => return Err(WriteIOError(e)) })
+)
+
 pub type ReadingResult = Result<Packet, ReadingError>;
 pub type WritingResult = Result<(), WritingError>;
 
-pub enum PacketType {
-   CommandType,
-   AnswerType
-}
+static MIN_PAYLOAD_SIZE: uint = 6;
+static MAX_PAYLOAD_SIZE: uint = 38;
 
 pub struct CommandPacket {
-   t: PacketType,
    timestamp: u64,
-   commandID: u8,
-   commandPayload: Vec<u8>
+   id: u8,
+   payload: Vec<u8> // May vary from 'MIN_PAYLOAD_SIZE' to 'MAX_PAYLOAD_SIZE' bytes.
+}
+
+pub enum ErrorType {
+   DecryptError,
+   AuthError
 }
 
 pub struct ErrorPacket {
+   t: ErrorType,
    timestamp: u64
 }
 
@@ -37,33 +50,107 @@ pub enum Packet {
 }
 
 impl Packet {
-   pub fn newRandomCommand(timestamp: u64) -> Packet {     
-         
+   pub fn new_random_command(timestamp: u64) -> Packet {        
       let mut rng = task_rng();
       Command(CommandPacket {
-         t: CommandType,
-         timestamp: timestamp, 
-         commandID: random::<u8>(),
-         commandPayload: Vec::from_fn(distributions::Range::new(7u, 40u).ind_sample(&mut rng), |idx| 0u8)
+         timestamp: timestamp,
+         id: random::<u8>(),
+         payload: Vec::from_fn(distributions::Range::new(MIN_PAYLOAD_SIZE, MAX_PAYLOAD_SIZE + 1).ind_sample(&mut rng), |idx| random::<u8>())
       })
    }
 
-   pub fn write(&self, output: &mut io::Writer) ->  WritingResult {
-      
+   pub fn write(&self, output: &mut io::Writer) ->  WritingResult {            
       match self {
          &Command(ref command) => {
-            output.write_u8(0);
-            output.write_be_u64(command.timestamp);
+            // Data.
+            let mut data = Vec::with_capacity(command.payload.len() + 1);
+            data.push(command.id);
+            data.push(command.payload.len() as u8);
+            data.push_all(command.payload.as_slice());
+                        
+            // Padding.
+            let padding_size = if data.len() % 16 == 0 { 16 } else { data.len() % 16 } ;
+            let padding = Vec::from_elem(padding_size, padding_size as u8);
+            data.push_all(padding.as_slice());
+            
+            // Encrypt.
+            let encrypted_data = crypto::encrypt(data.as_slice(), iv_from_timestamp(command.timestamp).as_slice());
+            
+            // Data size.
+            try_write_io!(output.write_be_u16(encrypted_data.len() as u16));
+         
+            // Packet type.
+            try_write_io!(output.write_u8(0));
+            
+            // Timestamp.
+            try_write_io!(output.write_be_u64(command.timestamp));
             
+            // Write encrypted data.
+            try_write_io!(output.write(encrypted_data.as_slice()));
             
+            // MAC.
+            try_write_io!(output.write(crypto::compute_mac(data.as_slice())));
          },
-         &Error(error) => ()
+         &Error(error) => {            
+            // Padding as data: 16 * '0'.
+            let padding = Vec::from_elem(16, '0' as u8);
+            let encrypted_data = crypto::encrypt(padding.as_slice(), iv_from_timestamp(error.timestamp).as_slice());
+           
+            // Data size.
+            try_write_io!(output.write_be_u16(encrypted_data.len() as u16));
+            
+            // Error type.
+            try_write_io!(match error.t {
+               DecryptError => output.write_u8(0x0A),
+               AuthError => output.write_u8(0x0B)
+            });
+            
+            // Timestamp.
+            try_write_io!(output.write_be_u64(error.timestamp));
+            
+            // Encrypt the padding.
+            try_write_io!(output.write(encrypted_data.as_slice()));
+            
+            // MAC.
+            try_write_io!(output.write(crypto::compute_mac(padding.as_slice())));
+         }
       }
    
       Ok(())
    }
    
    pub fn read(input: &mut io::Reader) -> ReadingResult {
-      Ok(Packet::newRandomCommand(0))
+      let data_size = try_read_io!(input.read_be_u16());
+   
+      match try_read_io!(input.read_byte()) {
+         0 => {
+            let mut command = CommandPacket { timestamp: try_read_io!(input.read_be_u64()), id: 0, payload: vec!() };
+            let mut encrypted_data = Vec::from_elem(data_size as uint, 0u8);
+            if try_read_io!(input.read(encrypted_data.as_mut_slice_())) != data_size as uint {
+               return Err(IOReadError(io::IoError { kind: io::EndOfFile, desc: "Unable to read all encrypted data", detail: None }));
+            }
+            
+            let data = crypto::decrypt(encrypted_data.as_slice(), iv_from_timestamp(command.timestamp).as_slice());            
+            if data.len() < MIN_PAYLOAD_SIZE + 2 || data.len() > MAX_PAYLOAD_SIZE + 2 {
+               return Err(UnconsistentDataSizeError)
+            }           
+
+            Ok(Command(command))
+         },
+         t if t == 0x0A || t == 0x0B  => {
+            let mut error = ErrorPacket { t: if t == 0x0A { DecryptError } else { AuthError }, timestamp: try_read_io!(input.read_be_u64()) };
+            // TODO
+            Ok(Error(error))
+         },
+         _ => Err(UnknownPacketTypeReadError),
+      }
    }
-}
\ No newline at end of file
+}
+
+fn iv_from_timestamp(t: u64) -> Vec<u8> {
+   // Initialization vector = 64 * 0u8 + timestamp.
+   let mut iv = io::MemWriter::with_capacity(16);
+   iv.write_be_u64(0u64);
+   iv.write_be_u64(t);
+   iv.unwrap()
+}