chacha20 stream wrapper

mod.rs

use crypto::chacha20::ChaCha20;
use crypto::symmetriccipher::SynchronousStreamCipher;

use std::io::prelude::*;

// Typedefs
type ReadResult = std::io::Result<usize>;
type WriteResult = std::io::Result<usize>;
type FlushResult = std::io::Result<()>;

/**
 * To be able to wrap the reader, the reader must implement
 * these traits. This is just a convienence to keep from having
 * to update the trait list everywhere else (eg, if I wanted to
 * add 'Seek' as a required trait as well).
 */
pub trait ChaChaRead: Read + Sized {}
impl<T> ChaChaRead for T where T: Read + Sized {}

/**
 * Same as ChaChaRead, but for Writers.
 */
pub trait ChaChaWrite: Write + Sized {}
impl<T> ChaChaWrite for T where T: Write + Sized {}

/**
 * The ChaChaReadStream struct contains the state needed to read
 * from a ChaChaRead compliant type.
 */
pub struct ChaChaReadStream<'a, R: ChaChaRead> {
  chacha: ChaCha20,
  reader: &'a mut R,
  buffer: [u8; 4096]
}

impl <'a, R: ChaChaRead> ChaChaReadStream<'a, R> {
  pub fn new(reader: &'a mut R, key: &[u8], nonce: &[u8]) -> Self {
    Self {
      chacha: ChaCha20::new(key, nonce),
      reader: reader,
      buffer: [0; 4096]
    }
  }
}

/**
 * The ChaChaWriteStream struct contains the state needed to
 * writeo to a ChaChaWrite compliant type.
 */
pub struct ChaChaWriteStream<'a, W: ChaChaWrite> {
  chacha: ChaCha20,
  writer: &'a mut W,
  buffer: [u8; 4096]
}

impl <'a, W: ChaChaWrite> ChaChaWriteStream<'a, W> {
  pub fn new(writer: &'a mut W, key: &[u8], nonce: &[u8]) -> Self {
    Self {
      chacha: ChaCha20::new(key, nonce),
      writer: writer,
      buffer: [0; 4096]
    }
  }
}

/**
 * The Read implementation for ChaChaStream will decrypt the file
 * contents on the fly. We process at most 4096 bytes at a time
 * because reasons.
 */
impl <'a, R: ChaChaRead> std::io::Read for ChaChaReadStream<'a, R> {
  fn read(&mut self, buf: &mut [u8]) -> ReadResult {
    if buf.len() == 0 {
      return Ok(0);
    }

    // If the slice we got is larger than our buffer, we read in
    // at most sizeof(our_buffer). Otherwise we read in at most
    // sizeof(buf).
    // We could loop over this buf.len() / self.buffer.len() times.
    let to_read = std::cmp::min(buf.len(), self.buffer.len());
    let mut read = 0;
    
    // XXX: This whole thing can almost certainly be optimized - 
    // I'm just not sure how to do it "the rust way".
    for n in 0..to_read {
      let mut input: [u8; 1] = [0];

      // Get the next byte from the file.
      let read_result = self.reader.read_exact(&mut input);

      // If we hit the end, we're done. Otherwise raise the error.
      if let Err(e) = read_result {
        if e.kind() != std::io::ErrorKind::UnexpectedEof {
          return Err(e)
        }
      }

      // Add it to our buffer and update the read count.
      self.buffer[n] = input[0];
      read = read + 1;
    }

    // Decrypt!
    self.chacha.process(&self.buffer[0..read], &mut buf[0..read]);

    // Report num bytes consumed.
    Ok(read)
  }
}

/**
 * On the fly encryption, writes in chunks of at most 4096 bytes.
 */
impl <'a, W: ChaChaWrite> std::io::Write for ChaChaWriteStream<'a, W> {
  fn write(&mut self, buf: &[u8]) -> WriteResult {
    // We could loop over this buf.len() / self.buffer.len() times.
    let to_write = std::cmp::min(buf.len(), self.buffer.len());

    self.chacha.process(&buf[0..to_write], &mut self.buffer[0..to_write]);
    let status = self.writer.write_all(&self.buffer[0..to_write]);

    // Failed to write all... Disk full/missing?
    if let Err(error) = status {
      return Err(error);
    }

    Ok(to_write)
  }

  // Flush the underlaying file.
  fn flush(&mut self) -> FlushResult {
    self.writer.flush()
  }
}

#[cfg(test)]
mod tests {
  use std::io::Cursor;
  use std::vec::Vec;

  use super::*;

  #[test]
  fn test_symetric() {
    // Pick something larger than our internal buffer size of 4096.
    const BUF_SIZE: usize = 5000;

    // Start with a zeroed out buffer.
    let buffer: Vec<u8> = vec![0; BUF_SIZE];
    let mut cursor = Cursor::new(buffer);

    // Use some static key and nonce.
    let key: [u8; 32] = [0xab; 32];
    let nonce: [u8; 12] = [0x12; 12];

    {
      // Encrypt our data, which is just counting from zero to 255
      // over and over until we reach the end of the buffer.
      let mut writer: ChaChaWriteStream<Cursor<Vec<u8>>> = ChaChaWriteStream::new(&mut cursor, &key, &nonce);

      for n in 0..BUF_SIZE {
        writer.write(&[(n % 255) as u8]).expect("failed to write");
      }
    }
    cursor.seek(std::io::SeekFrom::Start(0)).expect("failed to seek to start of cursor");

    {
      // Get the raw buffer back.
      let buf = cursor.into_inner();

      // Make sure that the data isn't what we wrote out (eg, it should
      // be encrypted. There's such a low chance of this being equal I
      // don't care that it's technically possible for these to collide).
      let mut ok = false;
      for n in 0..BUF_SIZE {
        if buf[n] != (n%255) as u8 {
          ok = true;
          break;
        }
      }
      assert!(ok);

      // Reset the cursor
      cursor = Cursor::new(buf);
    }

    {
      // Now read it back, making sure we have valid data.
      let mut reader: ChaChaReadStream<Cursor<Vec<u8>>> = ChaChaReadStream::new(&mut cursor, &key, &nonce);

      for n in 0..BUF_SIZE {
        let mut buf: [u8; 1] = [0; 1];
        reader.read_exact(&mut buf).expect("failed to read");
        println!("buf[0]: {}, n%255: {}", buf[0], (n%255) as u8);
        assert!(buf[0] == (n%255) as u8);
      }
    }
  }
}