kreuzberg 4.7.2

//! Security utilities for document extractors.
//!
//! This module provides validation and protection mechanisms against common attacks:
//! - ZIP bomb detection (decompression bombs)
//! - XML entity expansion limits
//! - Nesting depth limits
//! - Input size limits
//! - Entity length validation

#[cfg(feature = "archives")]
use std::io::{Read, Seek};

/// Configuration for security limits across extractors.
///
/// All limits are intentionally conservative to prevent DoS attacks
/// while still supporting legitimate documents.
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
#[serde(default)]
pub struct SecurityLimits {
    /// Maximum uncompressed size for archives (500 MB)
    pub max_archive_size: usize,

    /// Maximum compression ratio before flagging as potential bomb (100:1)
    pub max_compression_ratio: usize,

    /// Maximum number of files in archive (10,000)
    pub max_files_in_archive: usize,

    /// Maximum nesting depth for structures (100)
    pub max_nesting_depth: usize,

    /// Maximum entity/string length (32)
    pub max_entity_length: usize,

    /// Maximum string growth per document (100 MB)
    pub max_content_size: usize,

    /// Maximum iterations per operation
    pub max_iterations: usize,

    /// Maximum XML depth (100 levels)
    pub max_xml_depth: usize,

    /// Maximum cells per table (100,000)
    pub max_table_cells: usize,
}

impl Default for SecurityLimits {
    fn default() -> Self {
        Self {
            max_archive_size: 500 * 1024 * 1024,
            max_compression_ratio: 100,
            max_files_in_archive: 10_000,
            max_nesting_depth: 100,
            max_entity_length: 32,
            max_content_size: 100 * 1024 * 1024,
            max_iterations: 10_000_000,
            max_xml_depth: 100,
            max_table_cells: 100_000,
        }
    }
}

/// Security validation errors.
#[derive(Debug, Clone)]
pub enum SecurityError {
    /// Potential ZIP bomb detected
    ZipBombDetected {
        compressed_size: u64,
        uncompressed_size: u64,
        ratio: f64,
    },

    /// Archive exceeds maximum size
    ArchiveTooLarge { size: u64, max: usize },

    /// Archive contains too many files
    TooManyFiles { count: usize, max: usize },

    /// Nesting too deep
    NestingTooDeep { depth: usize, max: usize },

    /// Content exceeds maximum size
    ContentTooLarge { size: usize, max: usize },

    /// Entity/string too long
    EntityTooLong { length: usize, max: usize },

    /// Too many iterations
    TooManyIterations { count: usize, max: usize },

    /// XML depth exceeded
    XmlDepthExceeded { depth: usize, max: usize },

    /// Too many table cells
    TooManyCells { cells: usize, max: usize },
}

impl std::fmt::Display for SecurityError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            SecurityError::ZipBombDetected {
                compressed_size,
                uncompressed_size,
                ratio,
            } => {
                write!(
                    f,
                    "Potential ZIP bomb detected: compressed {}B -> uncompressed {}B (ratio: {:.1}:1)",
                    compressed_size, uncompressed_size, ratio
                )
            }
            SecurityError::ArchiveTooLarge { size, max } => {
                write!(f, "Archive too large: {} bytes (max: {} bytes)", size, max)
            }
            SecurityError::TooManyFiles { count, max } => {
                write!(f, "Archive has too many files: {} (max: {})", count, max)
            }
            SecurityError::NestingTooDeep { depth, max } => {
                write!(f, "Nesting too deep: {} levels (max: {})", depth, max)
            }
            SecurityError::ContentTooLarge { size, max } => {
                write!(f, "Content too large: {} bytes (max: {} bytes)", size, max)
            }
            SecurityError::EntityTooLong { length, max } => {
                write!(f, "Entity too long: {} chars (max: {})", length, max)
            }
            SecurityError::TooManyIterations { count, max } => {
                write!(f, "Too many iterations: {} (max: {})", count, max)
            }
            SecurityError::XmlDepthExceeded { depth, max } => {
                write!(f, "XML depth exceeded: {} (max: {})", depth, max)
            }
            SecurityError::TooManyCells { cells, max } => {
                write!(f, "Too many table cells: {} (max: {})", cells, max)
            }
        }
    }
}

impl std::error::Error for SecurityError {}

/// Helper struct for validating ZIP archives for security issues.
#[cfg(feature = "archives")]
pub struct ZipBombValidator {
    limits: SecurityLimits,
}

#[cfg(feature = "archives")]
impl ZipBombValidator {
    /// Create a new ZIP bomb validator.
    pub fn new(limits: SecurityLimits) -> Self {
        Self { limits }
    }

    /// Validate a ZIP archive for security issues.
    ///
    /// # Arguments
    /// * `archive` - Mutable ZIP archive to validate
    ///
    /// # Returns
    /// * `Ok(())` if archive is safe
    /// * `Err(SecurityError)` if security limit violated
    pub fn validate<R: Read + Seek>(&self, archive: &mut zip::ZipArchive<R>) -> Result<(), SecurityError> {
        let file_count = archive.len();

        if file_count > self.limits.max_files_in_archive {
            return Err(SecurityError::TooManyFiles {
                count: file_count,
                max: self.limits.max_files_in_archive,
            });
        }

        let mut total_uncompressed: u64 = 0;
        let mut total_compressed: u64 = 0;

        for i in 0..file_count {
            if let Ok(file) = archive.by_index(i) {
                let compressed_size = file.compressed_size();
                let uncompressed_size = file.size();

                total_uncompressed += uncompressed_size;
                total_compressed += compressed_size;

                if compressed_size > 0 && uncompressed_size > 0 {
                    let ratio = uncompressed_size as f64 / compressed_size as f64;
                    if ratio > self.limits.max_compression_ratio as f64 {
                        return Err(SecurityError::ZipBombDetected {
                            compressed_size,
                            uncompressed_size,
                            ratio,
                        });
                    }
                }
            }
        }

        if total_uncompressed > self.limits.max_archive_size as u64 {
            return Err(SecurityError::ArchiveTooLarge {
                size: total_uncompressed,
                max: self.limits.max_archive_size,
            });
        }

        if total_compressed > 0 {
            let ratio = total_uncompressed as f64 / total_compressed as f64;
            if ratio > self.limits.max_compression_ratio as f64 {
                return Err(SecurityError::ZipBombDetected {
                    compressed_size: total_compressed,
                    uncompressed_size: total_uncompressed,
                    ratio,
                });
            }
        }

        Ok(())
    }
}

/// Helper struct for tracking and validating string growth.
pub struct StringGrowthValidator {
    max_size: usize,
    current_size: usize,
}

impl StringGrowthValidator {
    /// Create a new string growth validator.
    pub fn new(max_size: usize) -> Self {
        Self {
            max_size,
            current_size: 0,
        }
    }

    /// Validate and update size after appending.
    ///
    /// # Returns
    /// * `Ok(())` if size is within limits
    /// * `Err(SecurityError)` if size exceeds limit
    pub fn check_append(&mut self, len: usize) -> Result<(), SecurityError> {
        self.current_size = self.current_size.saturating_add(len);
        if self.current_size > self.max_size {
            Err(SecurityError::ContentTooLarge {
                size: self.current_size,
                max: self.max_size,
            })
        } else {
            Ok(())
        }
    }

    /// Get current size.
    pub fn current_size(&self) -> usize {
        self.current_size
    }
}

/// Helper struct for validating iteration counts.
pub struct IterationValidator {
    max_iterations: usize,
    current_count: usize,
}

impl IterationValidator {
    /// Create a new iteration validator.
    pub fn new(max_iterations: usize) -> Self {
        Self {
            max_iterations,
            current_count: 0,
        }
    }

    /// Validate and increment iteration count.
    ///
    /// # Returns
    /// * `Ok(())` if count is within limits
    /// * `Err(SecurityError)` if count exceeds limit
    pub fn check_iteration(&mut self) -> Result<(), SecurityError> {
        self.current_count += 1;
        if self.current_count > self.max_iterations {
            Err(SecurityError::TooManyIterations {
                count: self.current_count,
                max: self.max_iterations,
            })
        } else {
            Ok(())
        }
    }

    /// Get current iteration count.
    pub fn current_count(&self) -> usize {
        self.current_count
    }
}

/// Helper struct for validating nesting depth.
pub struct DepthValidator {
    max_depth: usize,
    current_depth: usize,
}

impl DepthValidator {
    /// Create a new depth validator.
    pub fn new(max_depth: usize) -> Self {
        Self {
            max_depth,
            current_depth: 0,
        }
    }

    /// Push a level (increase depth).
    ///
    /// # Returns
    /// * `Ok(())` if depth is within limits
    /// * `Err(SecurityError)` if depth exceeds limit
    pub fn push(&mut self) -> Result<(), SecurityError> {
        self.current_depth += 1;
        if self.current_depth > self.max_depth {
            Err(SecurityError::NestingTooDeep {
                depth: self.current_depth,
                max: self.max_depth,
            })
        } else {
            Ok(())
        }
    }

    /// Pop a level (decrease depth).
    pub fn pop(&mut self) {
        if self.current_depth > 0 {
            self.current_depth -= 1;
        }
    }

    /// Get current depth.
    pub fn current_depth(&self) -> usize {
        self.current_depth
    }
}

/// Helper struct for validating entity/string length.
pub struct EntityValidator {
    max_length: usize,
}

impl EntityValidator {
    /// Create a new entity validator.
    pub fn new(max_length: usize) -> Self {
        Self { max_length }
    }

    /// Validate entity length.
    ///
    /// # Returns
    /// * `Ok(())` if length is within limits
    /// * `Err(SecurityError)` if length exceeds limit
    pub fn validate(&self, content: &str) -> Result<(), SecurityError> {
        if content.len() > self.max_length {
            Err(SecurityError::EntityTooLong {
                length: content.len(),
                max: self.max_length,
            })
        } else {
            Ok(())
        }
    }
}

/// Helper struct for validating table cell counts.
pub struct TableValidator {
    max_cells: usize,
    current_cells: usize,
}

impl TableValidator {
    /// Create a new table validator.
    pub fn new(max_cells: usize) -> Self {
        Self {
            max_cells,
            current_cells: 0,
        }
    }

    /// Add cells to table and validate.
    ///
    /// # Returns
    /// * `Ok(())` if cell count is within limits
    /// * `Err(SecurityError)` if cell count exceeds limit
    pub fn add_cells(&mut self, count: usize) -> Result<(), SecurityError> {
        self.current_cells = self.current_cells.saturating_add(count);
        if self.current_cells > self.max_cells {
            Err(SecurityError::TooManyCells {
                cells: self.current_cells,
                max: self.max_cells,
            })
        } else {
            Ok(())
        }
    }

    /// Get current cell count.
    pub fn current_cells(&self) -> usize {
        self.current_cells
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_depth_validator() {
        let mut validator = DepthValidator::new(3);

        assert!(validator.push().is_ok());
        assert_eq!(validator.current_depth(), 1);

        assert!(validator.push().is_ok());
        assert_eq!(validator.current_depth(), 2);

        assert!(validator.push().is_ok());
        assert_eq!(validator.current_depth(), 3);

        assert!(validator.push().is_err());
        assert_eq!(validator.current_depth(), 4);

        validator.pop();
        assert_eq!(validator.current_depth(), 3);
    }

    #[test]
    fn test_entity_validator() {
        let validator = EntityValidator::new(10);

        assert!(validator.validate("short").is_ok());

        assert!(validator.validate("0123456789").is_ok());

        assert!(validator.validate("01234567890").is_err());
    }

    #[test]
    fn test_string_growth_validator() {
        let mut validator = StringGrowthValidator::new(100);

        assert!(validator.check_append(50).is_ok());
        assert_eq!(validator.current_size(), 50);

        assert!(validator.check_append(50).is_ok());
        assert_eq!(validator.current_size(), 100);

        assert!(validator.check_append(1).is_err());
    }

    #[test]
    fn test_iteration_validator() {
        let mut validator = IterationValidator::new(3);

        assert!(validator.check_iteration().is_ok());
        assert!(validator.check_iteration().is_ok());
        assert!(validator.check_iteration().is_ok());
        assert!(validator.check_iteration().is_err());
    }

    #[test]
    fn test_table_validator() {
        let mut validator = TableValidator::new(10);

        assert!(validator.add_cells(5).is_ok());
        assert_eq!(validator.current_cells(), 5);

        assert!(validator.add_cells(5).is_ok());
        assert_eq!(validator.current_cells(), 10);

        assert!(validator.add_cells(1).is_err());
    }

    #[test]
    fn test_default_limits() {
        let limits = SecurityLimits::default();
        assert_eq!(limits.max_archive_size, 500 * 1024 * 1024);
        assert_eq!(limits.max_nesting_depth, 100);
        assert_eq!(limits.max_entity_length, 32);
    }
}