use crate::{VfError, VfResult, api::ListDirOptions};
use std::path::Path;
pub struct TraversalBudget {
entries: usize,
bytes: usize,
}
impl TraversalBudget {
pub fn new(entries: usize, bytes: usize) -> Self {
Self { entries, bytes }
}
pub fn remaining_entries(&self) -> usize {
self.entries
}
pub fn remaining_path_bytes(&self) -> usize {
self.bytes
}
pub fn remaining_options(&self, options: ListDirOptions) -> ListDirOptions {
options
.with_max_entries(self.entries)
.with_max_path_bytes(self.bytes)
}
pub fn charge(&mut self, path: &Path) -> VfResult<()> {
self.charge_paths([path])
}
pub fn charge_paths<'a>(&mut self, paths: impl IntoIterator<Item = &'a Path>) -> VfResult<()> {
if self.entries == 0 {
return Err(limit_error());
}
let bytes = paths.into_iter().try_fold(0usize, |total, path| {
total
.checked_add(path.as_os_str().len())
.ok_or_else(limit_error)
})?;
self.charge_bytes(bytes)?;
self.entries -= 1;
Ok(())
}
pub fn charge_path(&mut self, path: &Path) -> VfResult<()> {
self.charge_bytes(path.as_os_str().len())
}
pub fn charge_bytes(&mut self, bytes: usize) -> VfResult<()> {
self.bytes = self.bytes.checked_sub(bytes).ok_or_else(limit_error)?;
Ok(())
}
}
fn limit_error() -> VfError {
VfError::client(0, libc::EFBIG as u32)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn root_path_charges_and_failed_entries_preserve_remaining_budget() {
let mut budget = TraversalBudget::new(1, 5);
budget.charge_path(Path::new("/a")).unwrap();
assert_eq!(budget.remaining_entries(), 1);
assert!(budget.charge(Path::new("/long")).is_err());
assert_eq!(budget.remaining_entries(), 1);
assert_eq!(
budget
.remaining_options(ListDirOptions::new())
.max_path_bytes_or(crate::api::DEFAULT_DIRECTORY_MAX_PATH_BYTES),
3
);
budget.charge(Path::new("/é")).unwrap();
assert_eq!(budget.remaining_entries(), 0);
assert_eq!(
budget
.remaining_options(ListDirOptions::new())
.max_path_bytes_or(crate::api::DEFAULT_DIRECTORY_MAX_PATH_BYTES),
0
);
assert!(budget.charge(Path::new("")).is_err());
}
#[test]
fn maximum_limits_do_not_wrap_and_remaining_options_retain_policy() {
let mut budget = TraversalBudget::new(usize::MAX, usize::MAX);
budget.charge(Path::new("/file")).unwrap();
let options = budget.remaining_options(
ListDirOptions::new()
.with_recursive(true)
.with_allow_symlink_components(true)
.with_max_depth(crate::api::DepthLimit::new(7).expect("finite depth")),
);
assert_eq!(
options.max_entries_or(crate::api::DEFAULT_DIRECTORY_MAX_ENTRIES),
usize::MAX - 1
);
assert_eq!(
options.max_path_bytes_or(crate::api::DEFAULT_DIRECTORY_MAX_PATH_BYTES),
usize::MAX - 5
);
assert!(options.recursive());
assert!(options.allow_symlink_components());
assert_eq!(options.max_depth_or(Default::default()).get(), 7);
}
}