use std::marker::PhantomData;
use crate::ArchiveError;
use crate::Result;
use crate::SecurityConfig;
use crate::config::Validated;
#[derive(Debug)]
#[must_use]
pub struct QuotaPermit(PhantomData<()>);
#[derive(Debug)]
#[must_use]
pub struct EntryCount(PhantomData<()>);
#[derive(Debug, Default)]
pub struct QuotaTracker {
files_extracted: usize,
bytes_written: u64,
}
fn check_file_size(size: u64, max: u64) -> Result<()> {
if size > max {
core::hint::cold_path();
return Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::FileSize { size, max },
});
}
Ok(())
}
impl QuotaTracker {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[inline]
pub fn reserve(
&mut self,
size: u64,
config: &SecurityConfig<Validated>,
) -> Result<QuotaPermit> {
self.reserve_within(size, config.max_file_size, config)
}
pub(crate) fn reserve_count(
&mut self,
config: &SecurityConfig<Validated>,
) -> Result<EntryCount> {
self.bump_files()?;
self.check_file_count(config)?;
Ok(EntryCount(PhantomData))
}
pub(crate) fn charge(
&mut self,
count: EntryCount,
size: u64,
config: &SecurityConfig<Validated>,
) -> Result<QuotaPermit> {
drop(count);
check_file_size(size, config.max_file_size)?;
self.add_bytes(size)?;
self.check_total_size(config)?;
Ok(QuotaPermit(PhantomData))
}
pub(crate) fn reserve_without_file_limit(
&mut self,
size: u64,
config: &SecurityConfig<Validated>,
) -> Result<QuotaPermit> {
self.reserve_within(size, u64::MAX, config)
}
#[inline]
fn reserve_within(
&mut self,
size: u64,
max_file_size: u64,
config: &SecurityConfig<Validated>,
) -> Result<QuotaPermit> {
if max_file_size == u64::MAX
&& config.max_file_count == usize::MAX
&& config.max_total_size == u64::MAX
{
self.files_extracted = self.files_extracted.checked_add(1).ok_or_else(|| {
core::hint::cold_path();
ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::IntegerOverflow,
}
})?;
self.bytes_written = self.bytes_written.checked_add(size).ok_or_else(|| {
core::hint::cold_path();
ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::IntegerOverflow,
}
})?;
return Ok(QuotaPermit(PhantomData));
}
self.record_file_checked(size, max_file_size, config)?;
Ok(QuotaPermit(PhantomData))
}
#[inline(never)]
fn record_file_checked(
&mut self,
size: u64,
max_file_size: u64,
config: &SecurityConfig<Validated>,
) -> Result<()> {
check_file_size(size, max_file_size)?;
self.bump_files()?;
self.add_bytes(size)?;
self.check_file_count(config)?;
self.check_total_size(config)
}
fn bump_files(&mut self) -> Result<()> {
self.files_extracted = self.files_extracted.checked_add(1).ok_or_else(|| {
core::hint::cold_path();
ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::IntegerOverflow,
}
})?;
Ok(())
}
fn add_bytes(&mut self, size: u64) -> Result<()> {
self.bytes_written = self.bytes_written.checked_add(size).ok_or_else(|| {
core::hint::cold_path();
ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::IntegerOverflow,
}
})?;
Ok(())
}
fn check_file_count(&self, config: &SecurityConfig<Validated>) -> Result<()> {
if self.files_extracted > config.max_file_count {
core::hint::cold_path();
return Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::FileCount {
current: self.files_extracted,
max: config.max_file_count,
},
});
}
Ok(())
}
fn check_total_size(&self, config: &SecurityConfig<Validated>) -> Result<()> {
if self.bytes_written > config.max_total_size {
core::hint::cold_path();
return Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::TotalSize {
current: self.bytes_written,
max: config.max_total_size,
},
});
}
Ok(())
}
#[must_use]
pub fn files_extracted(&self) -> usize {
self.files_extracted
}
#[must_use]
pub fn bytes_written(&self) -> u64 {
self.bytes_written
}
}
#[cfg(test)]
#[allow(clippy::field_reassign_with_default, clippy::expect_used)]
mod tests {
use super::*;
use std::assert_matches;
#[test]
fn test_quota_tracker_new() {
let tracker = QuotaTracker::new();
assert_eq!(tracker.files_extracted(), 0);
assert_eq!(tracker.bytes_written(), 0);
}
#[test]
fn test_quota_tracker_reserve() {
let mut tracker = QuotaTracker::new();
let config = SecurityConfig::default().validate().expect("valid config");
assert!(tracker.reserve(1000, &config).is_ok());
assert_eq!(tracker.files_extracted(), 1);
assert_eq!(tracker.bytes_written(), 1000);
}
#[test]
fn test_quota_tracker_exceed_file_count() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_count = 2;
let config = config.validate().expect("valid config");
assert!(tracker.reserve(100, &config).is_ok());
assert!(tracker.reserve(100, &config).is_ok());
let result = tracker.reserve(100, &config);
assert_matches!(result, Err(ArchiveError::QuotaExceeded { .. }));
}
#[test]
fn test_quota_tracker_exceed_total_size() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_total_size = 1000;
let config = config.validate().expect("valid config");
assert!(tracker.reserve(600, &config).is_ok());
let result = tracker.reserve(500, &config);
assert_matches!(result, Err(ArchiveError::QuotaExceeded { .. }));
}
#[test]
fn test_quota_tracker_exceed_file_size() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_size = 1000;
let config = config.validate().expect("valid config");
let result = tracker.reserve(2000, &config);
assert_matches!(result, Err(ArchiveError::QuotaExceeded { .. }));
}
#[test]
fn test_quota_exactly_at_file_count_limit() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_count = 3;
config.max_total_size = u64::MAX;
config.max_file_size = u64::MAX;
let config = config.validate().expect("valid config");
assert!(
tracker.reserve(100, &config).is_ok(),
"file 1 should succeed"
);
assert!(
tracker.reserve(100, &config).is_ok(),
"file 2 should succeed"
);
assert!(
tracker.reserve(100, &config).is_ok(),
"file 3 should succeed"
);
assert_eq!(tracker.files_extracted(), 3, "should have 3 files");
let result = tracker.reserve(100, &config);
assert_matches!(
result,
Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::FileCount { current: 4, max: 3 }
}),
"file 4 should exceed quota"
);
}
#[test]
fn test_quota_exactly_at_total_size_limit() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_count = 100;
config.max_total_size = 1000;
config.max_file_size = u64::MAX;
let config = config.validate().expect("valid config");
assert!(tracker.reserve(600, &config).is_ok());
assert_eq!(tracker.bytes_written(), 600);
assert!(tracker.reserve(400, &config).is_ok());
assert_eq!(tracker.bytes_written(), 1000, "should be exactly at limit");
let result = tracker.reserve(1, &config);
assert_matches!(
result,
Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::TotalSize {
current: 1001,
max: 1000
}
}),
"exceeding total size should fail"
);
}
#[test]
fn test_quota_exactly_at_file_size_limit() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_count = 100;
config.max_total_size = u64::MAX;
config.max_file_size = 5000;
let config = config.validate().expect("valid config");
assert!(
tracker.reserve(5000, &config).is_ok(),
"file exactly at limit should succeed"
);
let result = tracker.reserve(5001, &config);
assert_matches!(
result,
Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::FileSize {
size: 5001,
max: 5000
}
}),
"file exceeding limit should fail"
);
}
#[test]
fn test_quota_off_by_one_file_count() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_count = 1;
config.max_total_size = u64::MAX;
config.max_file_size = u64::MAX;
let config = config.validate().expect("valid config");
assert!(tracker.reserve(100, &config).is_ok());
let result = tracker.reserve(100, &config);
assert_matches!(result, Err(ArchiveError::QuotaExceeded { .. }));
}
#[test]
fn test_quota_fast_path_unlimited() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_size = u64::MAX;
config.max_file_count = usize::MAX;
config.max_total_size = u64::MAX;
let config = config.validate().expect("valid config");
for i in 1..=1000 {
assert!(
tracker.reserve(1000, &config).is_ok(),
"file {i} should succeed with unlimited quotas"
);
}
assert_eq!(tracker.files_extracted(), 1000);
assert_eq!(tracker.bytes_written(), 1_000_000);
}
#[test]
fn test_quota_fast_path_overflow_detection() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_size = u64::MAX;
config.max_file_count = usize::MAX;
config.max_total_size = u64::MAX;
let config = config.validate().expect("valid config");
tracker.bytes_written = u64::MAX - 100;
let result = tracker.reserve(200, &config);
assert_matches!(
result,
Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::IntegerOverflow
}),
"fast path should still detect overflow"
);
}
#[test]
fn test_reserve_without_file_limit_skips_only_file_size() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_size = 10;
config.max_file_count = 2;
config.max_total_size = 1000;
let config = config.validate().expect("valid config");
assert_matches!(
tracker.reserve(500, &config),
Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::FileSize { size: 500, max: 10 }
})
);
assert!(tracker.reserve_without_file_limit(500, &config).is_ok());
assert_matches!(
tracker.reserve_without_file_limit(501, &config),
Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::TotalSize {
current: 1001,
max: 1000
}
})
);
}
#[test]
fn test_reserve_without_file_limit_enforces_file_count() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_size = 10;
config.max_file_count = 1;
let config = config.validate().expect("valid config");
assert!(tracker.reserve_without_file_limit(500, &config).is_ok());
assert_matches!(
tracker.reserve_without_file_limit(500, &config),
Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::FileCount { current: 2, max: 1 }
})
);
}
#[test]
fn test_reserve_without_file_limit_fast_path_accepts_oversized_file() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_size = 10;
config.max_file_count = usize::MAX;
config.max_total_size = u64::MAX;
let config = config.validate().expect("valid config");
assert!(tracker.reserve(500, &config).is_err());
assert!(tracker.reserve_without_file_limit(500, &config).is_ok());
assert_eq!(tracker.files_extracted(), 1);
assert_eq!(tracker.bytes_written(), 500);
}
#[test]
fn test_reserve_without_file_limit_fast_path_detects_overflow() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_size = 10;
config.max_file_count = usize::MAX;
config.max_total_size = u64::MAX;
let config = config.validate().expect("valid config");
tracker.bytes_written = u64::MAX - 100;
assert_matches!(
tracker.reserve_without_file_limit(200, &config),
Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::IntegerOverflow
})
);
}
#[test]
fn test_charge_adds_bytes_without_counting_again() {
let mut tracker = QuotaTracker::new();
let config = SecurityConfig::default().validate().expect("valid config");
let count = tracker.reserve_count(&config).expect("count");
assert_eq!(tracker.files_extracted(), 1);
assert_eq!(tracker.bytes_written(), 0);
drop(tracker.charge(count, 40, &config).expect("charge"));
assert_eq!(tracker.files_extracted(), 1);
assert_eq!(tracker.bytes_written(), 40);
}
#[test]
fn test_reserve_count_exceeds_file_count() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_count = 1;
let config = config.validate().expect("valid config");
let _first = tracker.reserve_count(&config).expect("first");
assert_matches!(
tracker.reserve_count(&config),
Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::FileCount { current: 2, max: 1 }
})
);
}
#[test]
fn test_charge_enforces_file_and_total_size() {
let mut tracker = QuotaTracker::new();
let mut config = SecurityConfig::default();
config.max_file_size = 100;
config.max_total_size = 150;
let config = config.validate().expect("valid config");
let a = tracker.reserve_count(&config).expect("a");
let b = tracker.reserve_count(&config).expect("b");
let c = tracker.reserve_count(&config).expect("c");
assert_matches!(
tracker.charge(a, 101, &config),
Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::FileSize { .. }
})
);
drop(tracker.charge(b, 100, &config).expect("within budget"));
assert_matches!(
tracker.charge(c, 100, &config),
Err(ArchiveError::QuotaExceeded {
resource: crate::QuotaResource::TotalSize { .. }
})
);
}
}