containerarchive/rust/src/restore.rs

/// Restore pipeline: reads a snapshot manifest, looks up chunks in the global
/// index, reads from pack files, decrypts, decompresses, and writes to a Unix socket.

use tokio::io::AsyncWriteExt;
use tokio::net::UnixStream;

use crate::compression;
use crate::encryption;
use crate::error::ArchiveError;
use crate::hasher;
use crate::pack_reader;
use crate::repository::Repository;
use crate::snapshot;

/// Restore a snapshot (or a specific item) to a Unix socket.
pub async fn restore(
    repo: &Repository,
    snapshot_id: &str,
    socket_path: &str,
    item_name: Option<&str>,
) -> Result<(), ArchiveError> {
    // Load snapshot manifest
    let snap = snapshot::load_snapshot(&repo.path, snapshot_id).await?;

    // Determine which items to restore
    let items_to_restore: Vec<&snapshot::SnapshotItem> = if let Some(name) = item_name {
        snap.items.iter()
            .filter(|i| i.name == name)
            .collect()
    } else {
        snap.items.iter().collect()
    };

    if items_to_restore.is_empty() {
        return Err(ArchiveError::NotFound(format!(
            "No items found in snapshot {}{}",
            snapshot_id,
            item_name.map(|n| format!(" with name '{}'", n)).unwrap_or_default()
        )));
    }

    // Connect to the Unix socket where TypeScript will read the restored data
    let mut stream = UnixStream::connect(socket_path).await
        .map_err(|e| ArchiveError::Io(e))?;

    tracing::info!("Connected to restore socket: {}", socket_path);

    let mut restored_bytes: u64 = 0;
    let mut chunks_read: u64 = 0;

    for item in items_to_restore {
        for hash_hex in &item.chunks {
            // Look up chunk in global index
            let index_entry = repo.index.get(hash_hex)
                .ok_or_else(|| ArchiveError::NotFound(format!(
                    "Chunk {} not found in index", hash_hex
                )))?;

            // Determine pack file path
            let shard = &index_entry.pack_id[..2];
            let pack_path = std::path::Path::new(&repo.path)
                .join("packs")
                .join("data")
                .join(shard)
                .join(format!("{}.pack", index_entry.pack_id));

            // Read chunk data from pack
            let stored_data = pack_reader::read_chunk(
                &pack_path,
                index_entry.offset,
                index_entry.compressed_size,
            ).await?;

            // Decrypt if encrypted
            let compressed = if let Some(ref key) = repo.master_key {
                // We need the nonce. Read it from the IDX file.
                let idx_path = std::path::Path::new(&repo.path)
                    .join("packs")
                    .join("data")
                    .join(shard)
                    .join(format!("{}.idx", index_entry.pack_id));

                let entries = pack_reader::load_idx(&idx_path).await?;
                let hash_bytes = hasher::hex_to_hash(hash_hex)
                    .map_err(|_| ArchiveError::Corruption(format!("Invalid hash: {}", hash_hex)))?;

                let idx_entry = pack_reader::find_in_idx(&entries, &hash_bytes)
                    .ok_or_else(|| ArchiveError::NotFound(format!(
                        "Chunk {} not found in pack index {}", hash_hex, index_entry.pack_id
                    )))?;

                encryption::decrypt_chunk(&stored_data, key, &idx_entry.nonce)?
            } else {
                stored_data
            };

            // Decompress
            let plaintext = compression::decompress(&compressed)?;

            // Verify hash
            let actual_hash = hasher::hash_chunk(&plaintext);
            let expected_hash = hasher::hex_to_hash(hash_hex)
                .map_err(|_| ArchiveError::Corruption(format!("Invalid hash: {}", hash_hex)))?;

            if actual_hash != expected_hash {
                return Err(ArchiveError::Corruption(format!(
                    "Hash mismatch for chunk {}: expected {}, got {}",
                    hash_hex,
                    hash_hex,
                    hasher::hash_to_hex(&actual_hash)
                )));
            }

            // Write to output socket
            stream.write_all(&plaintext).await?;

            restored_bytes += plaintext.len() as u64;
            chunks_read += 1;
        }
    }

    // Close the write side
    stream.shutdown().await?;

    tracing::info!(
        "Restore complete: {} bytes, {} chunks from snapshot {}",
        restored_bytes, chunks_read, snapshot_id
    );

    Ok(())
}