void_core/shard/

writer.rs

//! Shard writing and serialization

use rand::RngCore;

use crate::{Result, VoidError};

/// Padding strategy for shard sizes.
///
/// Padding shards to fixed size buckets hides true size from the server,
/// improving privacy by preventing size-based inference attacks.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum PaddingStrategy {
    /// No padding (original behavior)
    None,
    /// Pad to power-of-2 boundaries (1KB, 2KB, 4KB, 8KB, ...)
    #[default]
    PowerOfTwo,
    /// Pad to fixed bucket sizes (1KB, 4KB, 16KB, 64KB, 256KB, 1MB, 4MB)
    Buckets,
    /// Pad all shards to a fixed maximum size
    Fixed(usize),
}

/// Default bucket sizes for Buckets padding strategy.
pub const DEFAULT_BUCKETS: &[usize] = &[
    1024,    // 1 KB
    4096,    // 4 KB
    16384,   // 16 KB
    65536,   // 64 KB
    262144,  // 256 KB
    1048576, // 1 MB
    4194304, // 4 MB
];

/// Size of the padding length suffix (u64 little-endian).
const PADDING_SUFFIX_SIZE: usize = 8;

/// Magic bytes to identify padded shards: "VOIDPAD\0"
/// This distinguishes padded shards from old format shards.
const PADDING_MAGIC: [u8; 8] = *b"VOIDPAD\0";

/// Total size of padding footer: magic (8) + padding size (8)
const PADDING_FOOTER_SIZE: usize = 16;

/// Calculates the target padded size for a given data size.
pub fn calculate_padded_size(size: usize, strategy: PaddingStrategy) -> usize {
    match strategy {
        PaddingStrategy::None => size,
        PaddingStrategy::PowerOfTwo => {
            // Account for padding footer (magic + size)
            let total = size.saturating_add(PADDING_FOOTER_SIZE);
            if total == 0 {
                PADDING_FOOTER_SIZE
            } else {
                total.next_power_of_two()
            }
        }
        PaddingStrategy::Buckets => {
            let total = size.saturating_add(PADDING_FOOTER_SIZE);
            DEFAULT_BUCKETS
                .iter()
                .find(|&&b| b >= total)
                .copied()
                .unwrap_or(total) // If larger than all buckets, no padding
        }
        PaddingStrategy::Fixed(max) => max.max(size.saturating_add(PADDING_FOOTER_SIZE)),
    }
}

/// Reads the padding size from the end of shard data.
///
/// Returns `None` if this is an old format shard (no padding footer).
/// Returns `Some(padding_size)` if a valid padding footer is found.
/// Padded shards have a footer with: magic (8 bytes) + padding_size (8 bytes)
pub fn read_padding_info(data: &[u8]) -> Option<usize> {
    if data.len() < PADDING_FOOTER_SIZE {
        return None;
    }

    // Check for magic bytes before the padding size
    let magic_start = data.len() - PADDING_FOOTER_SIZE;
    let magic_end = magic_start + 8;
    if &data[magic_start..magic_end] != &PADDING_MAGIC {
        return None; // No padding (old format)
    }

    let padding = u64::from_le_bytes(
        data[data.len() - PADDING_SUFFIX_SIZE..]
            .try_into()
            .unwrap_or([0u8; 8]),
    ) as usize;

    // Sanity check: padding can't be larger than data minus the footer
    if padding.saturating_add(PADDING_FOOTER_SIZE) > data.len() {
        return None; // Invalid padding
    }
    Some(padding)
}

/// Builder for creating shards.
///
/// A shard is a zstd-compressed blob of concatenated file contents.
/// File indexing is handled by the TreeManifest, not the shard itself.
/// Shards are opaque blocks — like filesystem blocks.
pub struct ShardWriter {
    file_count: u32,
    body: Vec<u8>,
}

impl ShardWriter {
    /// Creates a new shard writer.
    pub fn new() -> Self {
        Self {
            file_count: 0,
            body: Vec::new(),
        }
    }

    /// Adds a file to the shard.
    ///
    /// # Arguments
    /// * `path` - File path (validated for safety)
    /// * `content` - File content bytes
    ///
    /// # Errors
    /// Returns `VoidError::Shard` if the path is invalid.
    pub fn add_file(&mut self, path: &str, content: &[u8]) -> Result<()> {
        if path.is_empty() {
            return Err(VoidError::Shard("empty path".into()));
        }

        // Normalize path (forward slashes, no leading slash)
        let normalized = path.replace('\\', "/").trim_start_matches('/').to_string();

        // Reject path traversal attempts (defense-in-depth)
        let path_check = std::path::Path::new(&normalized);
        for component in path_check.components() {
            match component {
                std::path::Component::ParentDir => {
                    return Err(VoidError::Shard("path contains '..'".into()));
                }
                std::path::Component::RootDir | std::path::Component::Prefix(_) => {
                    return Err(VoidError::Shard("absolute path not allowed".into()));
                }
                _ => {}
            }
        }

        self.file_count += 1;
        self.body.extend_from_slice(content);

        Ok(())
    }

    /// Finishes building and returns the zstd-compressed shard bytes.
    ///
    /// # Errors
    /// Returns `VoidError::Compression` if compression fails.
    pub fn finish(self, compression_level: i32) -> Result<Vec<u8>> {
        zstd::encode_all(self.body.as_slice(), compression_level)
            .map_err(|e| VoidError::Compression(e.to_string()))
    }

    /// Returns the number of files added.
    pub fn file_count(&self) -> u32 {
        self.file_count
    }

    /// Returns the uncompressed body size.
    pub fn body_size(&self) -> usize {
        self.body.len()
    }

    /// Returns true if this shard has no content.
    pub fn is_empty(&self) -> bool {
        self.file_count == 0
    }

    /// Finishes building and returns serialized shard bytes with optional size padding.
    ///
    /// Padding is applied to hide the true shard size from the server.
    /// The padding size is stored as a little-endian u64 at the end of the shard.
    ///
    /// # Errors
    /// Returns `VoidError::Compression` if compression fails.
    pub fn finish_padded(self, strategy: PaddingStrategy, compression_level: i32) -> Result<Vec<u8>> {
        let mut data = self.finish(compression_level)?;

        if matches!(strategy, PaddingStrategy::None) {
            return Ok(data);
        }

        let original_size = data.len();
        let target_size = calculate_padded_size(original_size, strategy);
        // Account for the full footer size (magic + padding size)
        let padding_size = target_size.saturating_sub(original_size + PADDING_FOOTER_SIZE);

        if padding_size > 0 {
            // Generate random padding
            let mut padding = vec![0u8; padding_size];
            rand::thread_rng().fill_bytes(&mut padding);
            data.extend(padding);
        }

        // Append magic bytes to identify this as a padded shard
        data.extend(&PADDING_MAGIC);
        // Append padding size as little-endian u64
        data.extend(&(padding_size as u64).to_le_bytes());

        Ok(data)
    }
}

impl Default for ShardWriter {
    fn default() -> Self {
        Self::new()
    }
}