From 282f80b06859330d0aaa95f22bcc646bf61008ea Mon Sep 17 00:00:00 2001
From: Qinkun Bao <qinkunbao@gmail.com>
Date: Wed, 10 Jun 2020 14:02:51 -0400
Subject: [PATCH] read/write files as small chunks

---
 cli/src/main.rs   |  9 ++++-----
 crypto/src/lib.rs | 30 ++++++++++++++++++++++++------
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/cli/src/main.rs b/cli/src/main.rs
index b38c3b292..7b4b60e47 100644
--- a/cli/src/main.rs
+++ b/cli/src/main.rs
@@ -118,10 +118,9 @@ fn decrypt(opt: EncryptDecryptOpt) -> Result<CMac> {
         }
         TeaclaveFile128Key::SCHEMA => {
             let key = TeaclaveFile128Key::new(&key)?;
-            let mut content = vec![];
-            let res = key.decrypt(opt.input_file, &mut content)?;
+            let mut output_file = fs::File::create(opt.output_file)?;
+            let res = key.decrypt(opt.input_file, &mut output_file)?;
             cmac.copy_from_slice(&res);
-            fs::write(opt.output_file, content)?;
         }
         _ => bail!("Invalid crypto algorithm"),
     }
@@ -151,8 +150,8 @@ fn encrypt(opt: EncryptDecryptOpt) -> Result<CMac> {
         }
         TeaclaveFile128Key::SCHEMA => {
             let key = TeaclaveFile128Key::new(&key)?;
-            let content = fs::read(opt.input_file)?;
-            let res = key.encrypt(opt.output_file, &content)?;
+            let content = fs::File::open(opt.input_file)?;
+            let res = key.encrypt(opt.output_file, content)?;
             cmac.copy_from_slice(&res);
         }
         _ => bail!("Invalid crypto algorithm"),
diff --git a/crypto/src/lib.rs b/crypto/src/lib.rs
index 95e320f61..1834bdfa8 100644
--- a/crypto/src/lib.rs
+++ b/crypto/src/lib.rs
@@ -28,6 +28,7 @@ use rand::prelude::RngCore;
 use ring::aead;
 use serde::{Deserialize, Serialize};
 use std::format;
+use std::io::{Read, Write};
 use std::path::Path;
 
 const AES_GCM_128_KEY_LENGTH: usize = 16;
@@ -37,6 +38,8 @@ const AES_GCM_256_KEY_LENGTH: usize = 32;
 const AES_GCM_256_IV_LENGTH: usize = 12;
 const TEACLAVE_FILE_128_ROOT_KEY_LENGTH: usize = 16;
 const CMAC_LENGTH: usize = 16;
+const FILE_CHUNK_SIZE: usize = 1024 * 1024;
+
 type CMac = [u8; CMAC_LENGTH];
 
 #[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]
@@ -200,18 +203,33 @@ impl TeaclaveFile128Key {
         Ok(TeaclaveFile128Key { key })
     }
 
-    pub fn decrypt<P: AsRef<Path>>(&self, path: P, out: &mut Vec<u8>) -> Result<CMac> {
-        use std::io::Read;
+    pub fn decrypt<P: AsRef<Path>>(&self, path: P, output: &mut impl Write) -> Result<CMac> {
         let mut file = ProtectedFile::open_ex(path.as_ref(), &self.key)?;
-        file.read_to_end(out)?;
+        let mut buffer = std::vec![0; FILE_CHUNK_SIZE];
+        loop {
+            let n = file.read(&mut buffer)?;
+            if n > 0 {
+                output.write(&buffer[..n])?;
+            } else {
+                break;
+            }
+        }
+        output.flush()?;
         let cmac = file.current_meta_gmac()?;
         Ok(cmac)
     }
 
-    pub fn encrypt<P: AsRef<Path>>(&self, path: P, content: &[u8]) -> Result<CMac> {
-        use std::io::Write;
+    pub fn encrypt<P: AsRef<Path>>(&self, path: P, mut content: impl Read) -> Result<CMac> {
         let mut file = ProtectedFile::create_ex(path.as_ref(), &self.key)?;
-        file.write_all(content)?;
+        let mut buffer = std::vec![0; FILE_CHUNK_SIZE];
+        loop {
+            let n = content.read(&mut buffer[..])?;
+            if n > 0 {
+                file.write(&buffer[..n])?;
+            } else {
+                break;
+            }
+        }
         file.flush()?;
         let cmac = file.current_meta_gmac()?;
         Ok(cmac)