self_encryption 0.34.1

// Copyright 2025 MaidSafe.net limited.
//
// This SAFE Network Software is licensed to you under The General Public License (GPL), version 3.
// Unless required by applicable law or agreed to in writing, the SAFE Network Software distributed
// under the GPL Licence is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. Please review the Licences for the specific language governing
// permissions and limitations relating to use of the SAFE Network Software.

//! Tests for streaming encryption functionality to ensure consistency between different methods.

use bytes::Bytes;
use self_encryption::{
    stream_encrypt, streaming_encrypt_from_file, test_helpers::random_bytes, DataMap, Error, Result,
};
use std::{
    collections::HashMap,
    io::Write,
    sync::{Arc, Mutex},
};
use tempfile::NamedTempFile;
use xor_name::XorName;

/// Helper function to create a temporary file with test data
fn create_temp_file_with_data(data: &[u8]) -> Result<NamedTempFile> {
    let mut temp_file = NamedTempFile::new()?;
    temp_file.write_all(data)?;
    Ok(temp_file)
}

/// Helper function to collect chunks from stream_encrypt
fn collect_stream_encrypt_chunks(
    data_size: usize,
    data_iter: impl Iterator<Item = Bytes>,
) -> Result<(DataMap, HashMap<XorName, Vec<u8>>)> {
    let mut stream = stream_encrypt(data_size, data_iter)?;
    let mut chunks = HashMap::new();

    // Collect all chunks
    for chunk_result in stream.chunks() {
        let (hash, content) = chunk_result?;
        chunks.insert(hash, content.to_vec());
    }

    // Get the datamap
    let datamap = stream
        .datamap()
        .expect("Should have DataMap after iteration")
        .clone();

    Ok((datamap, chunks))
}

/// Helper function to collect chunks from streaming_encrypt_from_file
fn collect_file_encrypt_chunks(
    file_path: &std::path::Path,
) -> Result<(DataMap, HashMap<XorName, Vec<u8>>)> {
    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 datamap = streaming_encrypt_from_file(file_path, store)?;

    // Extract all stored chunks
    let chunks = storage.lock().unwrap().clone();

    Ok((datamap, chunks))
}

/// Test that both streaming encryption methods produce identical results for the same input data
fn test_encryption_consistency(file_size: usize) -> Result<()> {
    println!("\n=== Testing encryption consistency for {file_size} bytes ===");

    // Generate test data
    let test_data = random_bytes(file_size);

    // Create temporary file for file-based encryption
    let temp_file = create_temp_file_with_data(&test_data)?;

    // Method 1: stream_encrypt (iterator-based)
    let data_iter = test_data
        .chunks(8192)
        .map(|chunk| Bytes::from(chunk.to_vec()));
    let (stream_datamap, stream_chunks) = collect_stream_encrypt_chunks(file_size, data_iter)?;

    println!(
        "Stream encrypt: {} chunks, datamap child level: {:?}",
        stream_chunks.len(),
        stream_datamap.child()
    );

    // Method 2: streaming_encrypt_from_file (file-based)
    let (file_datamap, file_chunks) = collect_file_encrypt_chunks(temp_file.path())?;

    println!(
        "File encrypt: {} chunks, datamap child level: {:?}",
        file_chunks.len(),
        file_datamap.child()
    );

    // STRICT VERIFICATION: Both methods must produce exactly identical results

    // 1. Identical DataMaps
    assert_eq!(
        stream_datamap.len(),
        file_datamap.len(),
        "DataMap lengths must be identical (stream: {}, file: {})",
        stream_datamap.len(),
        file_datamap.len()
    );

    assert_eq!(
        stream_datamap.child(),
        file_datamap.child(),
        "DataMap child levels must be identical (stream: {:?}, file: {:?})",
        stream_datamap.child(),
        file_datamap.child()
    );

    // 2. Identical chunk counts
    assert_eq!(
        stream_chunks.len(),
        file_chunks.len(),
        "Chunk counts must be identical (stream: {}, file: {})",
        stream_chunks.len(),
        file_chunks.len()
    );

    // 3. Identical chunk sets - every chunk from stream must exist in file with same content
    for (stream_hash, stream_content) in &stream_chunks {
        match file_chunks.get(stream_hash) {
            Some(file_content) => {
                assert_eq!(
                    stream_content,
                    file_content,
                    "Chunk content must be identical for hash {}",
                    hex::encode(stream_hash)
                );
            }
            None => {
                panic!(
                    "Chunk {} found in stream_encrypt but not in streaming_encrypt_from_file",
                    hex::encode(stream_hash)
                );
            }
        }
    }

    // 4. Verify no extra chunks in file method
    for file_hash in file_chunks.keys() {
        assert!(
            stream_chunks.contains_key(file_hash),
            "Chunk {} found in streaming_encrypt_from_file but not in stream_encrypt",
            hex::encode(file_hash)
        );
    }

    // 5. Verify DataMap chunk infos are identical
    let stream_infos = stream_datamap.infos();
    let file_infos = file_datamap.infos();

    for (stream_info, file_info) in stream_infos.iter().zip(file_infos.iter()) {
        assert_eq!(stream_info.index, file_info.index, "Chunk index must match");
        assert_eq!(
            stream_info.src_size, file_info.src_size,
            "Chunk src_size must match"
        );
        assert_eq!(
            stream_info.src_hash, file_info.src_hash,
            "Chunk src_hash must match"
        );
        assert_eq!(
            stream_info.dst_hash, file_info.dst_hash,
            "Chunk dst_hash must match"
        );
    }

    println!("✅ RIGOROUS VERIFICATION PASSED: Both methods produce identical chunks, counts, and DataMaps for {file_size} bytes");
    Ok(())
}

/// Test that both methods can decrypt to the same original data
fn test_decryption_consistency(file_size: usize) -> Result<()> {
    println!("\n=== Testing decryption consistency for {file_size} bytes ===");

    // Generate test data (use a fixed seed for reproducibility in debugging)
    let test_data = random_bytes(file_size);

    // Create temporary file for file-based encryption
    let temp_file = create_temp_file_with_data(&test_data)?;

    // Method 1: stream_encrypt (iterator-based, with shrinking)
    let data_iter = test_data
        .chunks(8192)
        .map(|chunk| Bytes::from(chunk.to_vec()));
    let (stream_datamap, stream_chunks) = collect_stream_encrypt_chunks(file_size, data_iter)?;

    // Method 2: streaming_encrypt_from_file (file-based, with shrinking)
    let (file_datamap, file_chunks) = collect_file_encrypt_chunks(temp_file.path())?;

    println!(
        "Stream encrypt: {} chunks collected, DataMap references {} chunks",
        stream_chunks.len(),
        stream_datamap.len()
    );
    println!(
        "File encrypt: {} chunks collected, DataMap references {} chunks",
        file_chunks.len(),
        file_datamap.len()
    );

    // Verify that all DataMap-referenced chunks are available
    println!("Verifying stream chunks...");
    for info in stream_datamap.infos() {
        let found = stream_chunks.contains_key(&info.dst_hash);
        println!(
            "  dst_hash {} -> found: {}",
            hex::encode(info.dst_hash),
            found
        );
        if !found {
            return Err(Error::Generic(format!(
                "Stream missing chunk: {}",
                hex::encode(info.dst_hash)
            )));
        }
    }

    println!("Verifying file chunks...");
    for info in file_datamap.infos() {
        let found = file_chunks.contains_key(&info.dst_hash);
        println!(
            "  dst_hash {} -> found: {}",
            hex::encode(info.dst_hash),
            found
        );
        if !found {
            return Err(Error::Generic(format!(
                "File missing chunk: {}",
                hex::encode(info.dst_hash)
            )));
        }
    }

    // Convert chunks to EncryptedChunk format for decryption
    let mut stream_encrypted_chunks = Vec::new();
    for content in stream_chunks.values() {
        use self_encryption::EncryptedChunk;
        stream_encrypted_chunks.push(EncryptedChunk {
            content: Bytes::from(content.clone()),
        });
    }

    let mut file_encrypted_chunks = Vec::new();
    for content in file_chunks.values() {
        use self_encryption::EncryptedChunk;
        file_encrypted_chunks.push(EncryptedChunk {
            content: Bytes::from(content.clone()),
        });
    }

    // Decrypt both results using their respective chunks and DataMaps
    println!("Attempting decryption...");
    let stream_decrypted = self_encryption::decrypt(&stream_datamap, &stream_encrypted_chunks)?;
    let file_decrypted = self_encryption::decrypt(&file_datamap, &file_encrypted_chunks)?;

    // Both should decrypt to the original data
    assert_eq!(
        stream_decrypted, test_data,
        "Stream encrypt decryption should match original data"
    );

    assert_eq!(
        file_decrypted, test_data,
        "File encrypt decryption should match original data"
    );

    // Both decrypted results should be identical (they should both match the original)
    assert_eq!(
        stream_decrypted, file_decrypted,
        "Both methods should decrypt to identical data"
    );

    println!(
        "✓ Decryption consistency verified for {} bytes (stream: {} chunks, file: {} chunks)",
        file_size,
        stream_chunks.len(),
        file_chunks.len()
    );
    Ok(())
}

/// Test that both methods produce the same chunk content (when properly aligned)
fn test_chunk_content_consistency(file_size: usize) -> Result<()> {
    println!("\n=== Testing chunk content consistency for {file_size} bytes ===");

    // Generate test data
    let test_data = random_bytes(file_size);

    // Create temporary file for file-based encryption
    let temp_file = create_temp_file_with_data(&test_data)?;

    // Method 1: stream_encrypt
    let data_iter = test_data
        .chunks(8192)
        .map(|chunk| Bytes::from(chunk.to_vec()));
    let (_stream_datamap, stream_chunks) = collect_stream_encrypt_chunks(file_size, data_iter)?;

    // Method 2: streaming_encrypt_from_file
    let (_file_datamap, file_chunks) = collect_file_encrypt_chunks(temp_file.path())?;

    // For chunk content consistency, we need to compare the unshrunk results
    // since shrinking may produce different chunk hashes

    // Both methods should produce similar numbers of initial chunks
    // (may differ by 1-2 due to implementation differences in shrinking)
    let chunk_count_diff = (stream_chunks.len() as i32 - file_chunks.len() as i32).abs();
    assert!(
        chunk_count_diff <= 2,
        "Chunk count difference should be small (got {} vs {}, diff: {})",
        stream_chunks.len(),
        file_chunks.len(),
        chunk_count_diff
    );

    println!(
        "✓ Chunk content consistency verified for {} bytes (stream: {} chunks, file: {} chunks)",
        file_size,
        stream_chunks.len(),
        file_chunks.len()
    );
    Ok(())
}

/// Test detailed chunk comparison between both methods
fn test_detailed_chunk_comparison(file_size: usize) -> Result<()> {
    println!("\n=== Testing detailed chunk comparison for {file_size} bytes ===");

    // Generate test data
    let test_data = random_bytes(file_size);

    // Create temporary file for file-based encryption
    let temp_file = create_temp_file_with_data(&test_data)?;

    // Method 1: stream_encrypt (iterator-based, now with shrinking)
    let data_iter = test_data
        .chunks(8192)
        .map(|chunk| Bytes::from(chunk.to_vec()));
    let (stream_datamap, stream_chunks) = collect_stream_encrypt_chunks(file_size, data_iter)?;

    // Method 2: streaming_encrypt_from_file (file-based, with shrinking)
    let (file_datamap, file_chunks) = collect_file_encrypt_chunks(temp_file.path())?;

    println!(
        "Stream encrypt: {} chunks, datamap child level: {:?}, datamap chunks: {}",
        stream_chunks.len(),
        stream_datamap.child(),
        stream_datamap.len()
    );

    println!(
        "File encrypt: {} chunks, datamap child level: {:?}, datamap chunks: {}",
        file_chunks.len(),
        file_datamap.child(),
        file_datamap.len()
    );

    // STRICT VERIFICATION: DataMaps must be exactly identical
    assert_eq!(
        stream_datamap.len(),
        file_datamap.len(),
        "DataMap lengths must be identical"
    );

    assert_eq!(
        stream_datamap.child(),
        file_datamap.child(),
        "DataMap child levels must be identical"
    );

    // Verify chunk infos are identical
    let stream_infos = stream_datamap.infos();
    let file_infos = file_datamap.infos();
    assert_eq!(
        stream_infos.len(),
        file_infos.len(),
        "Number of chunk infos must match"
    );

    for (stream_info, file_info) in stream_infos.iter().zip(file_infos.iter()) {
        assert_eq!(stream_info.index, file_info.index, "Chunk index must match");
        assert_eq!(
            stream_info.src_size, file_info.src_size,
            "Chunk src_size must match"
        );
        assert_eq!(
            stream_info.src_hash, file_info.src_hash,
            "Chunk src_hash must match"
        );
        assert_eq!(
            stream_info.dst_hash, file_info.dst_hash,
            "Chunk dst_hash must match"
        );
    }

    // STRICT VERIFICATION: Both methods must produce the same set of chunks
    assert_eq!(
        stream_chunks.len(),
        file_chunks.len(),
        "Both methods must produce the same number of chunks"
    );

    // Verify each chunk from stream_encrypt exists in file_encrypt with identical content
    for (stream_hash, stream_content) in &stream_chunks {
        match file_chunks.get(stream_hash) {
            Some(file_content) => {
                assert_eq!(
                    stream_content,
                    file_content,
                    "Chunk content must be identical for hash {}",
                    hex::encode(stream_hash)
                );
            }
            None => {
                panic!(
                    "Chunk {} found in stream_encrypt but not in file_encrypt",
                    hex::encode(stream_hash)
                );
            }
        }
    }

    // Verify each chunk from file_encrypt exists in stream_encrypt
    for file_hash in file_chunks.keys() {
        assert!(
            stream_chunks.contains_key(file_hash),
            "Chunk {} found in file_encrypt but not in stream_encrypt",
            hex::encode(file_hash)
        );
    }

    println!(
        "✅ EXACT EQUIVALENCE VERIFIED: Both methods produce identical DataMaps and chunk sets"
    );
    Ok(())
}

/// Test that stream_encrypt works on its own (encrypt then decrypt)
fn test_stream_encrypt_roundtrip(file_size: usize) -> Result<()> {
    println!("\n=== Testing stream_encrypt roundtrip for {file_size} bytes ===");

    // Generate test data
    let test_data = random_bytes(file_size);

    // Method: stream_encrypt (iterator-based, with shrinking)
    let data_iter = test_data
        .chunks(8192)
        .map(|chunk| Bytes::from(chunk.to_vec()));
    let (stream_datamap, stream_chunks) = collect_stream_encrypt_chunks(file_size, data_iter)?;

    println!(
        "Stream encrypt: {} chunks collected, DataMap references {} chunks, child level: {:?}",
        stream_chunks.len(),
        stream_datamap.len(),
        stream_datamap.child()
    );

    // Verify that all DataMap-referenced chunks are available
    for info in stream_datamap.infos() {
        let found = stream_chunks.contains_key(&info.dst_hash);
        if !found {
            return Err(Error::Generic(format!(
                "Missing chunk: {}",
                hex::encode(info.dst_hash)
            )));
        }
    }

    // Convert chunks to EncryptedChunk format for decryption
    let mut encrypted_chunks = Vec::new();
    for content in stream_chunks.values() {
        use self_encryption::EncryptedChunk;
        encrypted_chunks.push(EncryptedChunk {
            content: Bytes::from(content.clone()),
        });
    }

    // Decrypt using the stream_encrypt results
    println!("Attempting decryption...");
    let decrypted = self_encryption::decrypt(&stream_datamap, &encrypted_chunks)?;

    // Should match original data
    assert_eq!(
        decrypted, test_data,
        "Stream encrypt roundtrip should match original data"
    );

    println!("✓ Stream encrypt roundtrip verified for {file_size} bytes");
    Ok(())
}

#[test]
fn test_5mb_stream_encrypt_roundtrip() -> Result<()> {
    test_stream_encrypt_roundtrip(5 * 1024 * 1024)
}

/// Test that streaming_encrypt_from_file works on its own (encrypt then decrypt)
fn test_file_encrypt_roundtrip(file_size: usize) -> Result<()> {
    println!("\n=== Testing file encrypt roundtrip for {file_size} bytes ===");

    // Generate test data
    let test_data = random_bytes(file_size);

    // Create temporary file for file-based encryption
    let temp_file = create_temp_file_with_data(&test_data)?;

    // Method: streaming_encrypt_from_file (file-based, with shrinking)
    let (file_datamap, file_chunks) = collect_file_encrypt_chunks(temp_file.path())?;

    println!(
        "File encrypt: {} chunks collected, DataMap references {} chunks, child level: {:?}",
        file_chunks.len(),
        file_datamap.len(),
        file_datamap.child()
    );

    // Verify that all DataMap-referenced chunks are available
    for info in file_datamap.infos() {
        let found = file_chunks.contains_key(&info.dst_hash);
        if !found {
            return Err(Error::Generic(format!(
                "Missing chunk: {}",
                hex::encode(info.dst_hash)
            )));
        }
    }

    // Convert chunks to EncryptedChunk format for decryption
    let mut encrypted_chunks = Vec::new();
    for content in file_chunks.values() {
        use self_encryption::EncryptedChunk;
        encrypted_chunks.push(EncryptedChunk {
            content: Bytes::from(content.clone()),
        });
    }

    // Decrypt using the file encrypt results
    println!("Attempting decryption...");
    let decrypted = self_encryption::decrypt(&file_datamap, &encrypted_chunks)?;

    // Should match original data
    assert_eq!(
        decrypted, test_data,
        "File encrypt roundtrip should match original data"
    );

    println!("✓ File encrypt roundtrip verified for {file_size} bytes");
    Ok(())
}

#[test]
fn test_5mb_file_encrypt_roundtrip() -> Result<()> {
    test_file_encrypt_roundtrip(5 * 1024 * 1024)
}

#[test]
fn test_5mb_encryption_consistency() -> Result<()> {
    test_encryption_consistency(5 * 1024 * 1024)
}

#[test]
fn test_10mb_encryption_consistency() -> Result<()> {
    test_encryption_consistency(10 * 1024 * 1024)
}

#[test]
fn test_100mb_encryption_consistency() -> Result<()> {
    test_encryption_consistency(100 * 1024 * 1024)
}

#[test]
fn test_5mb_decryption_consistency() -> Result<()> {
    test_decryption_consistency(5 * 1024 * 1024)
}

#[test]
fn test_10mb_decryption_consistency() -> Result<()> {
    test_decryption_consistency(10 * 1024 * 1024)
}

#[test]
fn test_100mb_decryption_consistency() -> Result<()> {
    test_decryption_consistency(100 * 1024 * 1024)
}

#[test]
fn test_5mb_chunk_content_consistency() -> Result<()> {
    test_chunk_content_consistency(5 * 1024 * 1024)
}

#[test]
fn test_10mb_chunk_content_consistency() -> Result<()> {
    test_chunk_content_consistency(10 * 1024 * 1024)
}

#[test]
fn test_100mb_chunk_content_consistency() -> Result<()> {
    test_chunk_content_consistency(100 * 1024 * 1024)
}

#[test]
fn test_5mb_detailed_chunk_comparison() -> Result<()> {
    test_detailed_chunk_comparison(5 * 1024 * 1024)
}

/// Comprehensive test that runs all consistency checks for multiple file sizes
#[test]
fn test_all_encryption_methods_consistency() -> Result<()> {
    let test_sizes = vec![
        5 * 1024 * 1024,   // 5MB
        10 * 1024 * 1024,  // 10MB
        100 * 1024 * 1024, // 100MB
    ];

    for &size in &test_sizes {
        println!(
            "\n🔍 Testing file size: {} bytes ({:.1} MB)",
            size,
            size as f64 / (1024.0 * 1024.0)
        );

        // Test encryption consistency
        test_encryption_consistency(size)?;

        // Test decryption consistency
        test_decryption_consistency(size)?;

        // Test chunk content consistency
        test_chunk_content_consistency(size)?;

        println!("✅ All consistency checks passed for {size} bytes");
    }

    println!("\n🎉 All encryption method consistency tests passed!");
    Ok(())
}