Skip to content

Commit

Permalink
fix: clippy and fmt
Browse files Browse the repository at this point in the history
  • Loading branch information
dirvine committed Dec 17, 2024
1 parent 53d5c8d commit 49f3bd9
Show file tree
Hide file tree
Showing 6 changed files with 127 additions and 116 deletions.
5 changes: 3 additions & 2 deletions examples/parallel_streaming_decryptor.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
use bytes::Bytes;
use clap::Parser;
use rayon::prelude::*;
use self_encryption::{deserialize, streaming_decrypt_from_storage, DataMap, Error, Result};
use std::{fs::File, io::Read, path::Path};
use xor_name::XorName;
use clap::{Parser, error::ErrorKind};

/// Parallel streaming decryptor for self-encrypted files
#[derive(Parser, Debug)]
Expand Down Expand Up @@ -56,7 +56,8 @@ fn validate_paths(args: &Args) -> Result<()> {
)));
}
// Try to verify write permissions
if !parent.metadata()
if !parent
.metadata()
.map(|m| m.permissions().readonly())
.unwrap_or(true)
{
Expand Down
22 changes: 10 additions & 12 deletions src/encrypt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,25 +111,23 @@ pub(crate) fn encrypt_chunk(content: Bytes, pki: (Pad, Key, Iv)) -> Result<Bytes
/// - For chunk 0: Uses hashes of the last two chunks
/// - For chunk 1: Uses hash of chunk 0 and the last chunk
/// - For chunks 2+: Uses hashes of the previous two chunks
pub(crate) fn encrypt_stream(
chunks: Vec<RawChunk>,
) -> Result<DataMap> {
pub(crate) fn encrypt_stream(chunks: Vec<RawChunk>) -> Result<DataMap> {
// Create a sorted vector of all hashes - we still need this for encryption
let src_hashes: Vec<_> = chunks.iter().map(|c| c.hash).collect();
let mut keys = Vec::with_capacity(chunks.len());

// First, process chunks 2 onwards in parallel since they only need their previous two hashes
let later_chunks: Vec<_> = chunks.iter().skip(2).collect();
let later_chunk_infos: Vec<ChunkInfo> = later_chunks
.into_par_iter()
.map(|chunk| {
let RawChunk { index, data, hash } = chunk;
let src_size = data.len();

let pki = get_pad_key_and_iv(*index, &src_hashes);
let encrypted_content = encrypt_chunk(data.clone(), pki)?;
let dst_hash = XorName::from_content(encrypted_content.as_ref());

Ok(ChunkInfo {
index: *index,
dst_hash,
Expand All @@ -138,15 +136,15 @@ pub(crate) fn encrypt_stream(
})
})
.collect::<Result<Vec<_>>>()?;

keys.extend(later_chunk_infos);

// Process chunk 1 (needs hash 0 and last hash)
let chunk = &chunks[1];
let pki = get_pad_key_and_iv(1, &src_hashes);
let encrypted_content = encrypt_chunk(chunk.data.clone(), pki)?;
let dst_hash = XorName::from_content(encrypted_content.as_ref());

// Insert at beginning since this is chunk 1
keys.insert(
0,
Expand All @@ -157,13 +155,13 @@ pub(crate) fn encrypt_stream(
src_size: chunk.data.len(),
},
);

// Process chunk 0 (needs last two hashes)
let chunk = &chunks[0];
let pki = get_pad_key_and_iv(0, &src_hashes);
let encrypted_content = encrypt_chunk(chunk.data.clone(), pki)?;
let dst_hash = XorName::from_content(encrypted_content.as_ref());

// Insert at beginning since this is chunk 0
keys.insert(
0,
Expand All @@ -174,6 +172,6 @@ pub(crate) fn encrypt_stream(
src_size: chunk.data.len(),
},
);

Ok(DataMap::new(keys))
}
91 changes: 51 additions & 40 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -532,36 +532,36 @@ where

// Create a buffered reader with a reasonable buffer size
let mut reader = BufReader::with_capacity(1024 * 1024, file);

// Read all chunks first to get their hashes
let mut chunks = Vec::with_capacity(num_chunks);
for chunk_index in 0..num_chunks {
let (start, end) = get_start_end_positions(file_size, chunk_index);
let chunk_size = end - start;
let mut chunk_data = vec![0u8; chunk_size];
reader.read_exact(&mut chunk_data)?;

let hash = XorName::from_content(&chunk_data);
chunks.push(crate::chunk::RawChunk {
index: chunk_index,
data: Bytes::from(chunk_data),
hash,
});
}

// Process chunks and store them immediately
let data_map = encrypt::encrypt_stream(chunks.clone())?;

// Now encrypt and store each chunk
let src_hashes: Vec<_> = chunks.iter().map(|c| c.hash).collect();

for chunk in chunks {
let pki = get_pad_key_and_iv(chunk.index, &src_hashes);
let encrypted_content = encrypt::encrypt_chunk(chunk.data, pki)?;
let hash = XorName::from_content(&encrypted_content);
chunk_store(hash, encrypted_content)?;
}

// Shrink the data map and store additional chunks if needed
let (shrunk_data_map, _) = shrink_data_map(data_map, |hash, content| {
chunk_store(hash, content)?;
Expand Down Expand Up @@ -624,10 +624,7 @@ where
.iter()
.map(|info| {
let content = chunk_cache.get(&info.dst_hash).ok_or_else(|| {
Error::Generic(format!(
"Chunk not found for hash: {:?}",
info.dst_hash
))
Error::Generic(format!("Chunk not found for hash: {:?}", info.dst_hash))
})?;
Ok(EncryptedChunk {
content: content.clone(),
Expand Down Expand Up @@ -756,10 +753,10 @@ mod tests {
std::fs::write(&file_path, small_data)?;

let store = |_: XorName, _: Bytes| -> Result<()> { Ok(()) };

let result = streaming_encrypt_from_file(&file_path, store);
assert!(result.is_err());

Ok(())
}

Expand All @@ -773,15 +770,15 @@ mod tests {

let storage = Arc::new(Mutex::new(HashMap::new()));
let storage_clone = storage.clone();

let store = move |hash: XorName, content: Bytes| -> Result<()> {
let _ = storage_clone.lock().unwrap().insert(hash, content.to_vec());
Ok(())
};

let data_map = streaming_encrypt_from_file(&file_path, store)?;
assert!(data_map.chunk_identifiers.len() >= 3);

Ok(())
}

Expand All @@ -797,7 +794,7 @@ mod tests {

let storage = Arc::new(Mutex::new(HashMap::new()));
let storage_clone = storage.clone();

let store = move |hash: XorName, content: Bytes| -> Result<()> {
let _ = storage_clone.lock().unwrap().insert(hash, content.to_vec());
Ok(())
Expand All @@ -806,16 +803,22 @@ mod tests {
// First get the number of chunks directly
let bytes = Bytes::from(large_data.clone());
let (num_chunks, _) = chunk::batch_chunks(bytes);
assert!(num_chunks > 3,
"Should have more than 3 chunks before shrinking. Got: {}", num_chunks);
assert!(
num_chunks > 3,
"Should have more than 3 chunks before shrinking. Got: {}",
num_chunks
);

// Now test the streaming encryption
let data_map = streaming_encrypt_from_file(&file_path, store)?;

// After shrinking, should be exactly 3 chunks
assert_eq!(data_map.chunk_identifiers.len(), 3,
"Final data map should have exactly 3 chunks after shrinking");

assert_eq!(
data_map.chunk_identifiers.len(),
3,
"Final data map should have exactly 3 chunks after shrinking"
);

// Verify chunk indices are sequential
let mut prev_index = None;
for chunk_info in &data_map.chunk_identifiers {
Expand All @@ -824,7 +827,7 @@ mod tests {
}
prev_index = Some(chunk_info.index);
}

Ok(())
}

Expand All @@ -851,7 +854,7 @@ mod tests {

let result = streaming_encrypt_from_file(&file_path, store);
assert!(result.is_err());

Ok(())
}

Expand All @@ -865,15 +868,15 @@ mod tests {

let storage = Arc::new(Mutex::new(HashMap::new()));
let storage_clone = storage.clone();

let store = move |hash: XorName, content: Bytes| -> Result<()> {
let _ = storage_clone.lock().unwrap().insert(hash, content.to_vec());
Ok(())
};

// Encrypt the file
let data_map = streaming_encrypt_from_file(&file_path, store)?;

// Convert stored chunks to EncryptedChunk format
let stored = storage.lock().unwrap();
let encrypted_chunks: Vec<_> = stored
Expand All @@ -886,7 +889,7 @@ mod tests {
// Decrypt and verify
let decrypted = decrypt(&data_map, &encrypted_chunks)?;
assert_eq!(&original_data[..], &decrypted[..]);

Ok(())
}

Expand Down Expand Up @@ -921,10 +924,11 @@ mod tests {

// Create a very large data map (12 chunks)
let original_map = create_dummy_data_map(100000);
let (shrunk_map, _shrink_chunks) = shrink_data_map(original_map.clone(), |hash, content| {
let _ = storage_clone.lock().unwrap().insert(hash, content);
Ok(())
})?;
let (shrunk_map, _shrink_chunks) =
shrink_data_map(original_map.clone(), |hash, content| {
let _ = storage_clone.lock().unwrap().insert(hash, content);
Ok(())
})?;

// Verify multiple levels of shrinking occurred
assert!(shrunk_map.child().unwrap() > 0);
Expand All @@ -942,30 +946,36 @@ mod tests {
fn test_encryption_algorithm_consistency() -> Result<()> {
// Create deterministic test data
let test_data = vec![42u8; MIN_ENCRYPTABLE_BYTES * 2]; // Repeating value for predictability

// First encryption
let storage1 = Arc::new(Mutex::new(HashMap::new()));
let storage1_clone = storage1.clone();

let store1 = move |hash: XorName, content: Bytes| -> Result<()> {
let _ = storage1_clone.lock().unwrap().insert(hash, content.to_vec());
let _ = storage1_clone
.lock()
.unwrap()
.insert(hash, content.to_vec());
Ok(())
};

// Second encryption
let storage2 = Arc::new(Mutex::new(HashMap::new()));
let storage2_clone = storage2.clone();

let store2 = move |hash: XorName, content: Bytes| -> Result<()> {
let _ = storage2_clone.lock().unwrap().insert(hash, content.to_vec());
let _ = storage2_clone
.lock()
.unwrap()
.insert(hash, content.to_vec());
Ok(())
};

// Create temporary files with same content
let temp_dir = tempfile::TempDir::new()?;
let file_path1 = temp_dir.path().join("test1.bin");
let file_path2 = temp_dir.path().join("test2.bin");

std::fs::write(&file_path1, &test_data)?;
std::fs::write(&file_path2, &test_data)?;

Expand All @@ -983,7 +993,7 @@ mod tests {
// Compare stored chunks
let stored1 = storage1.lock().unwrap();
let stored2 = storage2.lock().unwrap();

assert_eq!(
stored1.len(),
stored2.len(),
Expand All @@ -992,10 +1002,11 @@ mod tests {

// Compare each chunk's content
for (hash, content1) in stored1.iter() {
let content2 = stored2.get(hash).expect("Chunk should exist in both storages");
let content2 = stored2
.get(hash)
.expect("Chunk should exist in both storages");
assert_eq!(
content1,
content2,
content1, content2,
"Encrypted chunks should be identical for same input"
);
}
Expand Down
Loading

0 comments on commit 49f3bd9

Please sign in to comment.