use std::io::{Read, Seek};
use crate::error::ConversionError;
use crate::format::InputFormat;
use crate::result::ConversionResult;
use crate::source::SourceDocument;
use crate::DocumentConverter;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ArchiveLimits {
pub max_entries: usize,
pub max_entry_size: u64,
pub max_total_size: u64,
pub max_compression_ratio: u64,
}
impl Default for ArchiveLimits {
fn default() -> Self {
Self {
max_entries: 10_000,
max_entry_size: 256 << 20,
max_total_size: 1 << 30,
max_compression_ratio: 200,
}
}
}
impl ArchiveLimits {
pub fn from_env() -> Self {
let d = Self::default();
let mb = |name: &str, default: u64| {
docling_core::env::parse::<u64>(name).map_or(default, |v| v.saturating_mul(1 << 20))
};
Self {
max_entries: docling_core::env::parse("DOCLING_RS_ZIP_MAX_ENTRIES")
.unwrap_or(d.max_entries),
max_entry_size: mb("DOCLING_RS_ZIP_MAX_ENTRY_MB", d.max_entry_size),
max_total_size: mb("DOCLING_RS_ZIP_MAX_TOTAL_MB", d.max_total_size),
max_compression_ratio: docling_core::env::parse("DOCLING_RS_ZIP_MAX_RATIO")
.unwrap_or(d.max_compression_ratio),
}
}
}
const RATIO_FLOOR: u64 = 1 << 20;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ArchiveEntryInfo {
pub index: usize,
pub path: String,
pub format: Option<InputFormat>,
pub skipped: Option<String>,
pub size: u64,
}
pub struct Archive<R> {
zip: zip::ZipArchive<R>,
entries: Vec<ArchiveEntryInfo>,
}
impl<R: Read + Seek> Archive<R> {
pub fn open(reader: R, limits: &ArchiveLimits) -> Result<Self, ConversionError> {
let mut zip = zip::ZipArchive::new(reader)
.map_err(|e| ConversionError::Parse(format!("zip: {e}")))?;
let mut entries = Vec::new();
let mut total: u64 = 0;
for index in 0..zip.len() {
let Ok(file) = zip.by_index_raw(index) else {
entries.push(ArchiveEntryInfo {
index,
path: format!("#{index}"),
format: None,
skipped: Some("unreadable entry header".into()),
size: 0,
});
continue;
};
if file.is_dir() {
continue;
}
let size = file.size();
let compressed = file.compressed_size();
let safe = file
.enclosed_name()
.map(|p| p.to_string_lossy().replace('\\', "/"));
let path = safe.clone().unwrap_or_else(|| file.name().to_string());
let skip = |why: &str| Some(why.to_string());
let mut format = None;
let skipped = if entries.len() >= limits.max_entries {
skip("over the archive's entry limit")
} else if safe.is_none() {
skip("unsafe path")
} else if file.encrypted() {
skip("encrypted")
} else if is_macos_metadata(&path) {
skip("macOS metadata")
} else if is_archive(&path) {
skip("nested archive")
} else if let Some(f) = extension(&path).and_then(InputFormat::from_extension) {
if size > limits.max_entry_size {
skip("larger than the per-entry size limit")
} else if size > RATIO_FLOOR
&& size / compressed.max(1) > limits.max_compression_ratio
{
skip("compression ratio over the limit")
} else if total.saturating_add(size) > limits.max_total_size {
skip("over the archive's total size limit")
} else {
total += size;
format = Some(f);
None
}
} else {
skip("unsupported file type")
};
entries.push(ArchiveEntryInfo {
index,
path,
format,
skipped,
size,
});
}
Ok(Self { zip, entries })
}
pub fn entries(&self) -> &[ArchiveEntryInfo] {
&self.entries
}
pub fn read(&mut self, index: usize) -> Result<SourceDocument, ConversionError> {
let info = self
.entries
.iter()
.find(|e| e.index == index)
.ok_or_else(|| ConversionError::Parse(format!("zip: no entry #{index}")))?;
let format = match (&info.format, &info.skipped) {
(Some(f), None) => *f,
(_, reason) => {
return Err(ConversionError::Parse(format!(
"zip: {} is skipped ({})",
info.path,
reason.as_deref().unwrap_or("not convertible")
)))
}
};
let (path, declared) = (info.path.clone(), info.size);
let file = self
.zip
.by_index(index)
.map_err(|e| ConversionError::Parse(format!("zip: {path}: {e}")))?;
let mut bytes = Vec::with_capacity(declared.min(64 << 20) as usize);
file.take(declared.saturating_add(1))
.read_to_end(&mut bytes)
.map_err(|e| ConversionError::Parse(format!("zip: {path}: {e}")))?;
if bytes.len() as u64 > declared {
return Err(ConversionError::Parse(format!(
"zip: {path}: larger than its header declares"
)));
}
let name = path
.rsplit('/')
.next()
.map(|f| f.rsplit_once('.').map_or(f, |(stem, _)| stem))
.filter(|s| !s.is_empty())
.unwrap_or("document")
.to_string();
Ok(SourceDocument::from_bytes(name, format, bytes))
}
}
#[derive(Debug)]
pub enum ArchiveOutcome {
Converted(Box<ConversionResult>),
Skipped(String),
Failed(ConversionError),
}
#[derive(Debug)]
pub struct ArchiveItem {
pub path: String,
pub outcome: ArchiveOutcome,
}
pub struct ArchiveConversion<'c, R> {
converter: &'c DocumentConverter,
archive: Archive<R>,
next: usize,
}
impl<R: Read + Seek> Iterator for ArchiveConversion<'_, R> {
type Item = ArchiveItem;
fn next(&mut self) -> Option<ArchiveItem> {
let info = self.archive.entries.get(self.next)?.clone();
self.next += 1;
let outcome = match info.skipped {
Some(reason) => ArchiveOutcome::Skipped(reason),
None => match self.archive.read(info.index) {
Ok(source) => match self.converter.convert(source) {
Ok(result) => ArchiveOutcome::Converted(Box::new(result)),
Err(e) => ArchiveOutcome::Failed(e),
},
Err(e) => ArchiveOutcome::Failed(e),
},
};
Some(ArchiveItem {
path: info.path,
outcome,
})
}
}
impl DocumentConverter {
pub fn convert_archive<R: Read + Seek>(
&self,
reader: R,
) -> Result<ArchiveConversion<'_, R>, ConversionError> {
Ok(ArchiveConversion {
converter: self,
archive: Archive::open(reader, &self.archive_limits_ref())?,
next: 0,
})
}
}
pub fn is_archive(path: &str) -> bool {
matches!(
extension(path).map(|e| e.to_ascii_lowercase()).as_deref(),
Some("zip" | "7z" | "rar" | "tar" | "gz" | "tgz" | "bz2" | "xz" | "zst")
)
}
fn is_macos_metadata(path: &str) -> bool {
path.starts_with("__MACOSX/") || path.rsplit('/').next().is_some_and(|f| f.starts_with("._"))
}
fn extension(path: &str) -> Option<&str> {
let file = path.rsplit('/').next()?;
file.rsplit_once('.')
.map(|(_, ext)| ext)
.filter(|e| !e.is_empty())
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{Cursor, Write};
fn zip_of(entries: &[(&str, &[u8])]) -> Vec<u8> {
let mut out = Cursor::new(Vec::new());
let mut w = zip::ZipWriter::new(&mut out);
let opts = zip::write::SimpleFileOptions::default();
for (name, data) in entries {
if name.ends_with('/') {
w.add_directory(name.trim_end_matches('/'), opts).unwrap();
} else {
w.start_file(*name, opts).unwrap();
w.write_all(data).unwrap();
}
}
w.finish().unwrap();
out.into_inner()
}
#[test]
fn entries_are_classified_before_inflating() {
let data = zip_of(&[
("docs/", b""),
("docs/a.md", b"# A\n\nalpha"),
("docs/b.html", b"<html><body><p>beta</p></body></html>"),
("tool.exe", b"MZ"),
("inner.zip", b"PK\x03\x04"),
("__MACOSX/docs/._a.md", b"junk"),
("noext", b"x"),
]);
let archive = Archive::open(Cursor::new(data), &ArchiveLimits::default()).unwrap();
let got: Vec<(&str, Option<&str>)> = archive
.entries()
.iter()
.map(|e| (e.path.as_str(), e.skipped.as_deref()))
.collect();
assert_eq!(
got,
[
("docs/a.md", None),
("docs/b.html", None),
("tool.exe", Some("unsupported file type")),
("inner.zip", Some("nested archive")),
("__MACOSX/docs/._a.md", Some("macOS metadata")),
("noext", Some("unsupported file type")),
]
);
}
#[test]
fn convert_archive_reports_every_entry() {
let data = zip_of(&[
("a.md", b"# A\n\nalpha"),
("broken.docx", b"not a zip at all"),
("c.exe", b"MZ"),
]);
let conv = DocumentConverter::new();
let items: Vec<ArchiveItem> = conv.convert_archive(Cursor::new(data)).unwrap().collect();
assert_eq!(items.len(), 3);
match &items[0].outcome {
ArchiveOutcome::Converted(r) => {
assert_eq!(r.input_name, "a");
assert!(r.document.export_to_markdown().contains("alpha"));
}
other => panic!("a.md: {other:?}"),
}
assert!(
matches!(items[1].outcome, ArchiveOutcome::Failed(_)),
"{:?}",
items[1]
);
assert!(matches!(items[2].outcome, ArchiveOutcome::Skipped(_)));
}
#[test]
fn limits_stop_entries_before_decompression() {
let big = vec![b'a'; 3 << 20]; let data = zip_of(&[("a.md", b"alpha"), ("bomb.md", &big), ("c.md", b"gamma")]);
let limits = ArchiveLimits::default();
let archive = Archive::open(Cursor::new(data.clone()), &limits).unwrap();
assert_eq!(
archive.entries()[1].skipped.as_deref(),
Some("compression ratio over the limit")
);
let tight = ArchiveLimits {
max_entries: 2,
max_entry_size: 4,
..limits
};
let archive = Archive::open(Cursor::new(data), &tight).unwrap();
let skipped: Vec<_> = archive
.entries()
.iter()
.map(|e| e.skipped.as_deref())
.collect();
assert_eq!(
skipped,
[
Some("larger than the per-entry size limit"),
Some("larger than the per-entry size limit"),
Some("over the archive's entry limit"),
]
);
}
#[test]
fn unsafe_paths_are_never_extracted() {
let data = zip_of(&[("../escape.md", b"x"), ("/abs.md", b"y"), ("ok.md", b"z")]);
let mut archive = Archive::open(Cursor::new(data), &ArchiveLimits::default()).unwrap();
let listed: Vec<(&str, Option<&str>)> = archive
.entries()
.iter()
.map(|e| (e.path.as_str(), e.skipped.as_deref()))
.collect();
assert_eq!(
listed,
[
("../escape.md", Some("unsafe path")),
("abs.md", None),
("ok.md", None)
]
);
assert!(archive.read(0).is_err());
assert_eq!(archive.read(1).unwrap().bytes, b"y");
}
#[test]
fn not_a_zip_is_an_error() {
assert!(Archive::open(Cursor::new(b"plain".to_vec()), &ArchiveLimits::default()).is_err());
assert_eq!(InputFormat::from_extension("zip"), None);
}
}