1 files changed, 104 insertions, 5 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 58e488e..fe7a985 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -115,6 +115,9 @@ impl Aez {
 
     /// Encrypt the given data.
     ///
+    /// This is a convenience function that allocates a fresh buffer of the appropriate size and
+    /// copies the data.
+    ///
     /// Parameters:
     ///
     /// * `nonce` -- the nonce to use. Each nonce should only be used once, as re-using the nonce
@@ -137,21 +140,81 @@ impl Aez {
         tau: u32,
         data: &[u8],
     ) -> Vec<u8> {
-        let mut buffer = vec![0; data.len() + tau as usize];
-        buffer[..data.len()].copy_from_slice(data);
-        encrypt(&self.key, nonce, associated_data, tau, &mut buffer);
+        let mut buffer = Vec::from(data);
+        self.encrypt_vec(nonce, associated_data, tau, &mut buffer);
         buffer
     }
 
+    /// Encrypts the data in the given [`Vec`].
+    ///
+    /// This function extends the vector with enough space to hold `tau` bytes of authentication
+    /// data. Afterwards, the vector will hold the ciphertext.
+    ///
+    /// If `tau == 0`, the vector will not be expanded.
+    ///
+    /// The parameters are the same as for [`Aez::encrypt`].
+    pub fn encrypt_vec(
+        &self,
+        nonce: &[u8],
+        associated_data: &[&[u8]],
+        tau: u32,
+        data: &mut Vec<u8>,
+    ) {
+        data.resize(data.len() + tau as usize, 0);
+        encrypt(&self.key, nonce, associated_data, tau, data);
+    }
+
+    /// Encrypts the data inplace.
+    ///
+    /// This function will overwrite the last `tau` bytes of the given buffer with the
+    /// authentication block before encrypting the data.
+    ///
+    /// If the buffer is smaller than `tau`, this function panics.
+    pub fn encrypt_inplace(
+        &self,
+        nonce: &[u8],
+        associated_data: &[&[u8]],
+        tau: u32,
+        buffer: &mut [u8],
+    ) {
+        assert!(buffer.len() >= tau as usize);
+        let data_len = buffer.len() - tau as usize;
+        append_auth(data_len, buffer);
+        encrypt(&self.key, nonce, associated_data, tau as u32, buffer);
+    }
+
+    /// Encrypts the data in the given buffer, writing the output to the given output buffer.
+    ///
+    /// This function will infer `tau` from the size difference between input and output. If the
+    /// output is smaller than the input, this funcion will panic.
+    ///
+    /// The `nonce` and `associated_data` parameters are the same as for [`Aez::encrypt`].
+    pub fn encrypt_buffer(
+        &self,
+        nonce: &[u8],
+        associated_data: &[&[u8]],
+        input: &[u8],
+        output: &mut [u8],
+    ) {
+        assert!(output.len() >= input.len());
+        let tau = output.len() - input.len();
+        output[..input.len()].copy_from_slice(input);
+        append_auth(input.len(), output);
+        encrypt(&self.key, nonce, associated_data, tau as u32, output);
+    }
+
     /// Decrypts the given ciphertext.
     ///
+    /// This is a convenience function that returns an owned version of the plaintext. If the
+    /// original buffer may be modified, you can use [`Aez::decrypt_inplace`] to save an allocation.
+    ///
     /// Parameters:
     ///
     /// * `nonce`, `associated_data` and `tau` are as for [`Aez::encrypt`].
     /// * `data` -- the ciphertext to decrypt.
     ///
-    /// Returns the decrypted content. If the authentication check fails, returns `None` instead.
-    /// The returned vector has length `data.len() - tau`.
+    /// Returns the decrypted content. If the integrity check fails, returns `None` instead. The
+    /// returned vector has length `data.len() - tau`.
     pub fn decrypt(
         &self,
         nonce: &[u8],
@@ -167,6 +230,21 @@ impl Aez {
         buffer.truncate(len);
         Some(buffer)
     }
+
+    /// Decrypt the given buffer in-place.
+    ///
+    /// Returns a slice to the valid plaintext subslice, or `None`.
+    ///
+    /// The parameters are the same as for [`Aez::decrypt`].
+    pub fn decrypt_inplace<'a>(
+        &self,
+        nonce: &[u8],
+        associated_data: &[&[u8]],
+        tau: u32,
+        data: &'a mut [u8],
+    ) -> Option<&'a [u8]> {
+        decrypt(&self.key, nonce, associated_data, tau, data)
+    }
 }
 
 fn aesenc(mut block: Block, key: &Block) -> block::Block {
@@ -754,6 +832,9 @@ fn split_key(key: &Key) -> (Block, Block, Block) {
 #[cfg(test)]
 mod test {
     use super::*;
+
+    static PLAIN: &[u8] = include_bytes!("payload.txt");
+
     #[test]
     fn test_extract() {
         for (a, b) in testvectors::EXTRACT_VECTORS {
@@ -888,6 +969,24 @@ mod test {
     }
 
     #[test]
+    fn test_encrypt_decrypt_inplace() {
+        let mut buffer = Vec::from(PLAIN);
+        let aez = Aez::new(b"foobar");
+        aez.encrypt_inplace(&[0], &[], 16, &mut buffer);
+        let plain = aez.decrypt_inplace(&[0], &[], 16, &mut buffer).unwrap();
+        assert_eq!(plain, &PLAIN[..PLAIN.len() - 16]);
+    }
+
+    #[test]
+    fn test_encrypt_decrypt_buffer() {
+        let mut output = vec![0; PLAIN.len() + 16];
+        let aez = Aez::new(b"foobar");
+        aez.encrypt_buffer(&[0], &[], PLAIN, &mut output);
+        let plain = aez.decrypt_inplace(&[0], &[], 16, &mut output).unwrap();
+        assert_eq!(plain, PLAIN);
+    }
+
+    #[test]
     fn test_encrypt_decrypt_long() {
         let message = b"ene mene miste es rappelt in der kiste ene mene meck und du bist weg";
         let aez = Aez::new(b"foobar");