Optimize (and fix) encrypted file reader for any data size

Break up reader in different functions.
Read correctly when the buffer is only one byte.
Read additional data produced when finalizing.
Always try to fill the given reading buffer.
Create an internal buffer to handle excess data.
Improve reading performance.

src/main.rs

@@ -12,6 +12,7 @@ extern crate serde_derive;
 extern crate serde_json;
 extern crate sha2;
 
+use std::cmp::min;
 use std::fmt;
 use std::fs::File;
 use std::io::{self, BufReader, Cursor, Read};
@@ -38,6 +39,7 @@ use reqwest::mime::APPLICATION_OCTET_STREAM;
 use reqwest::multipart::Part;
 use sha2::Sha256;
 
+/// The length in bytes of crytographic tags that are used.
 const TAG_LEN: usize = 16;
 
 fn main() {
@@ -167,7 +169,7 @@ fn hkdf<'a>(
     info: Option<&[u8]>,
 ) -> Vec<u8> {
     // Unwrap info or use empty info
-    let info = info.unwrap_or(b"");
+    let info = info.unwrap_or(&[]);
 
     // Derive a HKDF key with the given length
     Hkdf::<Sha256>::new(&ikm, &[])
@@ -270,19 +272,28 @@ impl Header for XFileMetadata {
 /// This greatly reduces memory usage for large files.
 ///
 /// This reader encrypts the file data with an appended GCM tag.
+///
+/// The reader uses a small internal buffer as data is encrypted in blocks,
+/// which may output more data than fits in the given buffer while reading.
+/// The excess data is then returned on the next read.
 struct EncryptedFileReaderTagged {
-    /// The file to read.
+    /// The raw file that is read from.
     file: File,
 
-    /// The cipher that is used for decryption.
+    /// The cipher type used for encrypting.
     cipher: Cipher,
 
-    /// The crypter used to encrypt the file data.
+    /// The crypter used for encrypting the read file.
     crypter: Crypter,
 
     /// A tag cursor that reads the tag to append,
     /// when the file is fully read and the tag is known.
     tag: Option<Cursor<Vec<u8>>>,
+
+    /// The internal buffer, containing encrypted data that has yet to be
+    /// outputted to the reader. This data is always outputted before any new
+    /// data is produced.
+    internal_buf: Vec<u8>,
 }
 
 impl EncryptedFileReaderTagged {
@@ -292,83 +303,162 @@ impl EncryptedFileReaderTagged {
     /// constructing, and constructs a reader that has a size similar to the
     /// file.
     pub fn new(file: File, cipher: Cipher, key: &[u8], iv: &[u8]) -> Self {
+        // Build the crypter
         // TODO: return proper errors from crypter
+        let crypter = Crypter::new(
+            cipher,
+            CrypterMode::Encrypt,
+            key,
+            Some(iv),
+        ).unwrap();
+
+        // Construct the encrypted reader
         EncryptedFileReaderTagged {
             file,
             cipher,
-            crypter: Crypter::new(
+            crypter,
-                cipher,
-                CrypterMode::Encrypt,
-                key,
-                Some(iv),
-            ).unwrap(),
             tag: None,
+            internal_buf: Vec::new(),
         }
     }
 
     /// Calculate the total length of the encrypted file with the appended
     /// tag.
+    /// Useful in combination with some progress monitor, to determine how much
+    /// of the file is read or for example; sent over the network.
     pub fn len(&self) -> Result<u64, io::Error> {
         Ok(self.file.metadata()?.len() + TAG_LEN as u64)
     }
+
+    /// Read data from the internal buffer if there is any data in it, into
+    /// the given `buf`.
+    ///
+    /// The number of bytes that were read into `buf` is returned.
+    ///
+    /// If there is no data to be read, or `buf` has a zero size, `0` is always
+    /// returned.
+    fn read_internal(&mut self, buf: &mut [u8]) -> usize {
+        // Return if there is no data to read
+        if self.internal_buf.is_empty() || buf.len() == 0 {
+            return 0;
+        }
+
+        // Determine how much data will be read
+        let len = min(buf.len(), self.internal_buf.len());
+
+        // Slice the section we will read from, copy to the reader
+        {
+            let (out, _) = self.internal_buf.split_at(len);
+            let (buf, _) = buf.split_at_mut(len);
+            buf.copy_from_slice(out);
+        }
+
+        // Drain the read data from the internal buffer
+        self.internal_buf.drain(..len);
+
+        len
+    }
+
+    /// Read data directly from the file, and encrypt it.
+    ///
+    /// Because data may be encrypted in blocks, it is possible more data
+    /// is produced than fits in the given `buf`. In that case the excess data
+    /// is stored in an internal buffer, and is ouputted the next time being
+    /// read from the reader.
+    ///
+    /// The number of bytes that is read into `buf` is returned.
+    fn read_file_encrypted(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
+        // Get the block size, determine the buffer size, create a data buffer
+        let block_size = self.cipher.block_size();
+        let mut data = vec![0u8; buf.len()];
+
+        // Read the file, return if nothing was read
+        let len = self.file.read(&mut data)?;
+        if len == 0 {
+            return Ok(0);
+        }
+
+        // Create an encrypted buffer, truncate the data buffer
+        let mut encrypted = vec![0u8; len + block_size];
+        data.truncate(len);
+
+        // Encrypt the data that was read
+        let len = self.crypter.update(&data, &mut encrypted).unwrap();
+
+        // Calculate how many bytes will be copied to the reader
+        let out_len = min(buf.len(), len);
+
+        // Fill the reader buffer
+        let (out, remaining) = encrypted.split_at(out_len);
+        let (buf, _) = buf.split_at_mut(out_len);
+        buf.copy_from_slice(out);
+
+        // Splice to the actual remaining bytes, store it for later
+        let (store, _) = remaining.split_at(len - out_len);
+        self.internal_buf.extend(store.iter());
+
+        // Return the number of bytes read to the reader
+        Ok(out_len)
+    }
+
+    /// Finalize the crypter once it is done encrypthing the whole file.
+    /// This finalization step produces a tag that is placed after the
+    /// encrypted file data.
+    ///
+    /// This step must be invoked to start reading the tag,
+    /// and after it has been invoked no data must be encrypted anymore.
+    ///
+    /// This method must only be invoked once.
+    fn finalize_file(&mut self) -> Result<(), io::Error> {
+        // Finalize the crypter, catch any remaining output
+        let mut output = vec![0u8; self.cipher.block_size()];
+        let len = self.crypter.finalize(&mut output)?;
+
+        // Move additional output in the internal buffer
+        if len > 0 {
+            self.internal_buf.extend(output.iter().take(len));
+        }
+
+        // Fetch the encryption tag, and create an internal reader for it
+        let mut tag = vec![0u8; TAG_LEN];
+        self.crypter.get_tag(&mut tag)?;
+        self.tag = Some(Cursor::new(tag));
+
+        Ok(())
+    }
 }
 
 impl Read for EncryptedFileReaderTagged {
+    /// Read from the encrypted file, and then the encryption tag.
     fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
-        // If the tag reader has been created, read from it
+        // Read from the internal buffer, return full or splice to empty
+        let len = self.read_internal(buf);
+        if len >= buf.len() {
+            return Ok(len);
+        }
+        let (_, buf) = buf.split_at_mut(len);
+
+        // Keep track of the total number of read bytes, to return
+        let mut total = len;
+
+        // If the tag reader has been created, only read from that one
         if let Some(ref mut tag) = self.tag {
-            return tag.read(buf);
+            return Ok(tag.read(buf)? + total);
         }
 
-        // Get the block size
+        // Read the encrypted file, return full or splice to empty
-        let block_size = self.cipher.block_size();
+        let len = self.read_file_encrypted(buf)?;
-
+        total += len;
-        // Create a raw file buffer
+        if len >= buf.len() {
-        let mut raw = vec![0u8; buf.len() - block_size];
+            return Ok(total);
-
-        // TODO: remove after debugging
-        println!("DEBUG: Reading raw: {} (buf size: {})", raw.len(), buf.len());
-
-        // Read from the file, and truncate the buffer
-        let len = self.file.read(&mut raw)?;
-        raw.truncate(len);
-
-        // Encrypt raw data if if something was read
-        if len > 0 {
-            // Encrypt the raw data
-            // TODO: store raw bytes that were not encrypted yet
-            let len_enc = self.crypter.update(&raw, buf).unwrap();
-
-            // TODO: remove after debugging
-            println!("DEBUG: Read: {}; Encrypted: {}", len, len_enc);
-
-            // Return the number of encrypted bytes
-            return Ok(len_enc);
         }
+        let (_, buf) = buf.split_at_mut(len);
 
-        // Create a buffer for data that might be returned when finalizing
+        // Finalize the file crypter, and build the tag
-        let mut output = vec![0u8; block_size];
+        self.finalize_file()?;
 
-        // Finalize the crypter, truncate the output
+        // Try to fill the remaining part of the buffer
-        let len = self.crypter.finalize(&mut output).unwrap();
+        Ok(self.read(buf)? + total)
-        //output.truncate(len);
-
-        // TODO: remove after debugging
-        if len > 0 {
-            println!("DEBUG: Read {} more bytes when finalized!", len);
-        }
-
-        // Create a buffer for the tag
-        let mut tag = vec![0u8; TAG_LEN];
-
-        // Get the tag
-        self.crypter.get_tag(&mut tag).unwrap();
-
-        // Set the tag
-        self.tag = Some(Cursor::new(tag));
-
-        // Read again, to start reading the tag
-        self.read(buf)
     }
 }