From 5f3f906804cd2e6871f0f2e982977244afa5e3ec Mon Sep 17 00:00:00 2001
From: qima <qi.ma@maidsafe.net>
Date: Wed, 6 Sep 2023 23:55:22 +0800
Subject: [PATCH] feat: stream encryptor write encrypted chunks to disk

---
 src/lib.rs   | 31 ++++++++++++++++++++++---------
 src/tests.rs | 32 ++++++++++++++++++++++++++------
 2 files changed, 48 insertions(+), 15 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 0cefdf727..bdcf2b2b8 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -157,11 +157,17 @@ pub struct StreamSelfEncryptor {
     data_map: Vec<ChunkInfo>,
     // Progressing collection of source chunks' names
     src_hashes: BTreeMap<usize, XorName>,
+    // File path to flush encrypted_chunks into.
+    chunk_dir: Option<Box<PathBuf>>,
 }
 
 impl StreamSelfEncryptor {
-    /// For encryption, return with an intialized streaming encryptor
-    pub fn encrypt_from_file(file_path: Box<PathBuf>) -> Result<Self> {
+    /// For encryption, return with an intialized streaming encryptor.
+    /// If a `chunk_dir` is provided, the encrypted_chunks will be written into the specified dir as well.
+    pub fn encrypt_from_file(
+        file_path: Box<PathBuf>,
+        chunk_dir: Option<Box<PathBuf>>,
+    ) -> Result<Self> {
         let file = File::open(&*file_path)?;
         let metadata = file.metadata()?;
         let file_size = metadata.len();
@@ -174,6 +180,7 @@ impl StreamSelfEncryptor {
             chunk_index: 0,
             data_map: Vec::new(),
             src_hashes: BTreeMap::new(),
+            chunk_dir,
         })
     }
 
@@ -201,13 +208,19 @@ impl StreamSelfEncryptor {
             src_size: end_pos - start_pos,
         });
 
-        Ok((
-            Some(EncryptedChunk {
-                index,
-                content: encrypted_content,
-            }),
-            None,
-        ))
+        let encrypted_chunk = EncryptedChunk {
+            index,
+            content: encrypted_content,
+        };
+
+        if let Some(chunk_dir) = self.chunk_dir.clone() {
+            let file_path = chunk_dir.join(hex::encode(dst_hash));
+            let result = File::create(file_path);
+            let mut output_file = result?;
+            output_file.write_all(&encrypted_chunk.content)?;
+        }
+
+        Ok((Some(encrypted_chunk), None))
     }
 
     fn read_chunk(&mut self, chunk_index: usize) -> Result<(XorName, Bytes)> {
diff --git a/src/tests.rs b/src/tests.rs
index de03d4624..d8a0844fb 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -15,7 +15,7 @@ use bytes::Bytes;
 use itertools::Itertools;
 use rand::prelude::SliceRandom;
 use std::{
-    fs::File,
+    fs::{create_dir_all, File},
     io::{Read, Write},
 };
 use tempfile::tempdir;
@@ -30,8 +30,14 @@ fn test_stream_self_encryptor() -> Result<(), Error> {
     let data = random_bytes(file_size);
     file.write_all(&data)?;
 
+    let chunk_path = dir.path().join("chunk_path");
+    create_dir_all(chunk_path.clone())?;
+
     // Encrypt the file using StreamSelfEncryptor
-    let mut encryptor = StreamSelfEncryptor::encrypt_from_file(Box::new(file_path))?;
+    let mut encryptor = StreamSelfEncryptor::encrypt_from_file(
+        Box::new(file_path),
+        Some(Box::new(chunk_path.clone())),
+    )?;
     let mut encrypted_chunks = Vec::new();
     let mut data_map = None;
     while let Ok((chunk, map)) = encryptor.next_encryption() {
@@ -44,6 +50,7 @@ fn test_stream_self_encryptor() -> Result<(), Error> {
             break;
         }
     }
+    let data_map = data_map.unwrap();
 
     // Shuffle the encrypted chunks
     let mut rng = rand::thread_rng();
@@ -51,10 +58,8 @@ fn test_stream_self_encryptor() -> Result<(), Error> {
 
     // Decrypt the shuffled chunks using StreamSelfDecryptor
     let decrypted_file_path = dir.path().join("decrypted");
-    let mut decryptor = StreamSelfDecryptor::decrypt_to_file(
-        Box::new(decrypted_file_path.clone()),
-        &data_map.unwrap(),
-    )?;
+    let mut decryptor =
+        StreamSelfDecryptor::decrypt_to_file(Box::new(decrypted_file_path.clone()), &data_map)?;
     for chunk in encrypted_chunks {
         let _ = decryptor.next_encrypted(chunk)?;
     }
@@ -65,6 +70,21 @@ fn test_stream_self_encryptor() -> Result<(), Error> {
     let _ = decrypted_file.read_to_end(&mut decrypted_data)?;
     assert_eq!(data, decrypted_data);
 
+    // Use the flushed encrypted chunks to recover the file and verify with the original data
+    let mut flushed_encrypted_chunks = Vec::new();
+    for chunk_info in data_map.infos() {
+        let file_path = chunk_path.join(&hex::encode(chunk_info.dst_hash));
+        let mut chunk_file = File::open(file_path)?;
+        let mut chunk_data = Vec::new();
+        let _ = chunk_file.read_to_end(&mut chunk_data)?;
+        flushed_encrypted_chunks.push(EncryptedChunk {
+            index: chunk_info.index,
+            content: chunk_data.into(),
+        });
+    }
+    let decrypted_flushed_data = decrypt_full_set(&data_map, &flushed_encrypted_chunks)?;
+    assert_eq!(data, decrypted_flushed_data);
+
     Ok(())
 }