1use std::io::{Read, Seek};
25
26use crate::error::ConversionError;
27use crate::format::InputFormat;
28use crate::result::ConversionResult;
29use crate::source::SourceDocument;
30use crate::DocumentConverter;
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub struct ArchiveLimits {
39 pub max_entries: usize,
42 pub max_entry_size: u64,
44 pub max_total_size: u64,
46 pub max_compression_ratio: u64,
49}
50
51impl Default for ArchiveLimits {
52 fn default() -> Self {
53 Self {
54 max_entries: 10_000,
55 max_entry_size: 256 << 20,
56 max_total_size: 1 << 30,
57 max_compression_ratio: 200,
58 }
59 }
60}
61
62impl ArchiveLimits {
63 pub fn from_env() -> Self {
67 let d = Self::default();
68 let mb = |name: &str, default: u64| {
69 docling_core::env::parse::<u64>(name).map_or(default, |v| v.saturating_mul(1 << 20))
70 };
71 Self {
72 max_entries: docling_core::env::parse("DOCLING_RS_ZIP_MAX_ENTRIES")
73 .unwrap_or(d.max_entries),
74 max_entry_size: mb("DOCLING_RS_ZIP_MAX_ENTRY_MB", d.max_entry_size),
75 max_total_size: mb("DOCLING_RS_ZIP_MAX_TOTAL_MB", d.max_total_size),
76 max_compression_ratio: docling_core::env::parse("DOCLING_RS_ZIP_MAX_RATIO")
77 .unwrap_or(d.max_compression_ratio),
78 }
79 }
80}
81
82const RATIO_FLOOR: u64 = 1 << 20;
85
86#[derive(Debug, Clone, PartialEq, Eq)]
88pub struct ArchiveEntryInfo {
89 pub index: usize,
91 pub path: String,
94 pub format: Option<InputFormat>,
97 pub skipped: Option<String>,
99 pub size: u64,
101}
102
103pub struct Archive<R> {
105 zip: zip::ZipArchive<R>,
106 entries: Vec<ArchiveEntryInfo>,
107}
108
109impl<R: Read + Seek> Archive<R> {
110 pub fn open(reader: R, limits: &ArchiveLimits) -> Result<Self, ConversionError> {
113 let mut zip = zip::ZipArchive::new(reader)
114 .map_err(|e| ConversionError::Parse(format!("zip: {e}")))?;
115 let mut entries = Vec::new();
116 let mut total: u64 = 0;
117 for index in 0..zip.len() {
118 let Ok(file) = zip.by_index_raw(index) else {
120 entries.push(ArchiveEntryInfo {
121 index,
122 path: format!("#{index}"),
123 format: None,
124 skipped: Some("unreadable entry header".into()),
125 size: 0,
126 });
127 continue;
128 };
129 if file.is_dir() {
130 continue;
131 }
132 let size = file.size();
133 let compressed = file.compressed_size();
134 let safe = file
135 .enclosed_name()
136 .map(|p| p.to_string_lossy().replace('\\', "/"));
137 let path = safe.clone().unwrap_or_else(|| file.name().to_string());
138 let skip = |why: &str| Some(why.to_string());
139 let mut format = None;
140 let skipped = if entries.len() >= limits.max_entries {
141 skip("over the archive's entry limit")
142 } else if safe.is_none() {
143 skip("unsafe path")
146 } else if file.encrypted() {
147 skip("encrypted")
148 } else if is_macos_metadata(&path) {
149 skip("macOS metadata")
150 } else if is_archive(&path) {
151 skip("nested archive")
152 } else if let Some(f) = extension(&path).and_then(InputFormat::from_extension) {
153 if size > limits.max_entry_size {
154 skip("larger than the per-entry size limit")
155 } else if size > RATIO_FLOOR
156 && size / compressed.max(1) > limits.max_compression_ratio
157 {
158 skip("compression ratio over the limit")
159 } else if total.saturating_add(size) > limits.max_total_size {
160 skip("over the archive's total size limit")
161 } else {
162 total += size;
163 format = Some(f);
164 None
165 }
166 } else {
167 skip("unsupported file type")
168 };
169 entries.push(ArchiveEntryInfo {
170 index,
171 path,
172 format,
173 skipped,
174 size,
175 });
176 }
177 Ok(Self { zip, entries })
178 }
179
180 pub fn entries(&self) -> &[ArchiveEntryInfo] {
182 &self.entries
183 }
184
185 pub fn read(&mut self, index: usize) -> Result<SourceDocument, ConversionError> {
190 let info = self
191 .entries
192 .iter()
193 .find(|e| e.index == index)
194 .ok_or_else(|| ConversionError::Parse(format!("zip: no entry #{index}")))?;
195 let format = match (&info.format, &info.skipped) {
196 (Some(f), None) => *f,
197 (_, reason) => {
198 return Err(ConversionError::Parse(format!(
199 "zip: {} is skipped ({})",
200 info.path,
201 reason.as_deref().unwrap_or("not convertible")
202 )))
203 }
204 };
205 let (path, declared) = (info.path.clone(), info.size);
206 let file = self
207 .zip
208 .by_index(index)
209 .map_err(|e| ConversionError::Parse(format!("zip: {path}: {e}")))?;
210 let mut bytes = Vec::with_capacity(declared.min(64 << 20) as usize);
212 file.take(declared.saturating_add(1))
213 .read_to_end(&mut bytes)
214 .map_err(|e| ConversionError::Parse(format!("zip: {path}: {e}")))?;
215 if bytes.len() as u64 > declared {
216 return Err(ConversionError::Parse(format!(
217 "zip: {path}: larger than its header declares"
218 )));
219 }
220 let name = path
221 .rsplit('/')
222 .next()
223 .map(|f| f.rsplit_once('.').map_or(f, |(stem, _)| stem))
224 .filter(|s| !s.is_empty())
225 .unwrap_or("document")
226 .to_string();
227 Ok(SourceDocument::from_bytes(name, format, bytes))
228 }
229}
230
231#[derive(Debug)]
233pub enum ArchiveOutcome {
234 Converted(Box<ConversionResult>),
236 Skipped(String),
239 Failed(ConversionError),
241}
242
243#[derive(Debug)]
245pub struct ArchiveItem {
246 pub path: String,
247 pub outcome: ArchiveOutcome,
248}
249
250pub struct ArchiveConversion<'c, R> {
256 converter: &'c DocumentConverter,
257 archive: Archive<R>,
258 next: usize,
259}
260
261impl<R: Read + Seek> Iterator for ArchiveConversion<'_, R> {
262 type Item = ArchiveItem;
263
264 fn next(&mut self) -> Option<ArchiveItem> {
265 let info = self.archive.entries.get(self.next)?.clone();
266 self.next += 1;
267 let outcome = match info.skipped {
268 Some(reason) => ArchiveOutcome::Skipped(reason),
269 None => match self.archive.read(info.index) {
270 Ok(source) => match self.converter.convert(source) {
271 Ok(result) => ArchiveOutcome::Converted(Box::new(result)),
272 Err(e) => ArchiveOutcome::Failed(e),
273 },
274 Err(e) => ArchiveOutcome::Failed(e),
275 },
276 };
277 Some(ArchiveItem {
278 path: info.path,
279 outcome,
280 })
281 }
282}
283
284impl DocumentConverter {
285 pub fn convert_archive<R: Read + Seek>(
291 &self,
292 reader: R,
293 ) -> Result<ArchiveConversion<'_, R>, ConversionError> {
294 Ok(ArchiveConversion {
295 converter: self,
296 archive: Archive::open(reader, &self.archive_limits_ref())?,
297 next: 0,
298 })
299 }
300}
301
302pub fn is_archive(path: &str) -> bool {
305 matches!(
306 extension(path).map(|e| e.to_ascii_lowercase()).as_deref(),
307 Some("zip" | "7z" | "rar" | "tar" | "gz" | "tgz" | "bz2" | "xz" | "zst")
308 )
309}
310
311fn is_macos_metadata(path: &str) -> bool {
314 path.starts_with("__MACOSX/") || path.rsplit('/').next().is_some_and(|f| f.starts_with("._"))
315}
316
317fn extension(path: &str) -> Option<&str> {
318 let file = path.rsplit('/').next()?;
319 file.rsplit_once('.')
320 .map(|(_, ext)| ext)
321 .filter(|e| !e.is_empty())
322}
323
324#[cfg(test)]
325mod tests {
326 use super::*;
327 use std::io::{Cursor, Write};
328
329 fn zip_of(entries: &[(&str, &[u8])]) -> Vec<u8> {
330 let mut out = Cursor::new(Vec::new());
331 let mut w = zip::ZipWriter::new(&mut out);
332 let opts = zip::write::SimpleFileOptions::default();
333 for (name, data) in entries {
334 if name.ends_with('/') {
335 w.add_directory(name.trim_end_matches('/'), opts).unwrap();
336 } else {
337 w.start_file(*name, opts).unwrap();
338 w.write_all(data).unwrap();
339 }
340 }
341 w.finish().unwrap();
342 out.into_inner()
343 }
344
345 #[test]
346 fn entries_are_classified_before_inflating() {
347 let data = zip_of(&[
348 ("docs/", b""),
349 ("docs/a.md", b"# A\n\nalpha"),
350 ("docs/b.html", b"<html><body><p>beta</p></body></html>"),
351 ("tool.exe", b"MZ"),
352 ("inner.zip", b"PK\x03\x04"),
353 ("__MACOSX/docs/._a.md", b"junk"),
354 ("noext", b"x"),
355 ]);
356 let archive = Archive::open(Cursor::new(data), &ArchiveLimits::default()).unwrap();
357 let got: Vec<(&str, Option<&str>)> = archive
358 .entries()
359 .iter()
360 .map(|e| (e.path.as_str(), e.skipped.as_deref()))
361 .collect();
362 assert_eq!(
363 got,
364 [
365 ("docs/a.md", None),
366 ("docs/b.html", None),
367 ("tool.exe", Some("unsupported file type")),
368 ("inner.zip", Some("nested archive")),
369 ("__MACOSX/docs/._a.md", Some("macOS metadata")),
370 ("noext", Some("unsupported file type")),
371 ]
372 );
373 }
374
375 #[test]
376 fn convert_archive_reports_every_entry() {
377 let data = zip_of(&[
378 ("a.md", b"# A\n\nalpha"),
379 ("broken.docx", b"not a zip at all"),
380 ("c.exe", b"MZ"),
381 ]);
382 let conv = DocumentConverter::new();
383 let items: Vec<ArchiveItem> = conv.convert_archive(Cursor::new(data)).unwrap().collect();
384 assert_eq!(items.len(), 3);
385 match &items[0].outcome {
386 ArchiveOutcome::Converted(r) => {
387 assert_eq!(r.input_name, "a");
388 assert!(r.document.export_to_markdown().contains("alpha"));
389 }
390 other => panic!("a.md: {other:?}"),
391 }
392 assert!(
393 matches!(items[1].outcome, ArchiveOutcome::Failed(_)),
394 "{:?}",
395 items[1]
396 );
397 assert!(matches!(items[2].outcome, ArchiveOutcome::Skipped(_)));
398 }
399
400 #[test]
401 fn limits_stop_entries_before_decompression() {
402 let big = vec![b'a'; 3 << 20]; let data = zip_of(&[("a.md", b"alpha"), ("bomb.md", &big), ("c.md", b"gamma")]);
404 let limits = ArchiveLimits::default();
405 let archive = Archive::open(Cursor::new(data.clone()), &limits).unwrap();
406 assert_eq!(
407 archive.entries()[1].skipped.as_deref(),
408 Some("compression ratio over the limit")
409 );
410 let tight = ArchiveLimits {
411 max_entries: 2,
412 max_entry_size: 4,
413 ..limits
414 };
415 let archive = Archive::open(Cursor::new(data), &tight).unwrap();
416 let skipped: Vec<_> = archive
417 .entries()
418 .iter()
419 .map(|e| e.skipped.as_deref())
420 .collect();
421 assert_eq!(
422 skipped,
423 [
424 Some("larger than the per-entry size limit"),
425 Some("larger than the per-entry size limit"),
426 Some("over the archive's entry limit"),
427 ]
428 );
429 }
430
431 #[test]
432 fn unsafe_paths_are_never_extracted() {
433 let data = zip_of(&[("../escape.md", b"x"), ("/abs.md", b"y"), ("ok.md", b"z")]);
436 let mut archive = Archive::open(Cursor::new(data), &ArchiveLimits::default()).unwrap();
437 let listed: Vec<(&str, Option<&str>)> = archive
438 .entries()
439 .iter()
440 .map(|e| (e.path.as_str(), e.skipped.as_deref()))
441 .collect();
442 assert_eq!(
443 listed,
444 [
445 ("../escape.md", Some("unsafe path")),
446 ("abs.md", None),
447 ("ok.md", None)
448 ]
449 );
450 assert!(archive.read(0).is_err());
451 assert_eq!(archive.read(1).unwrap().bytes, b"y");
452 }
453
454 #[test]
455 fn not_a_zip_is_an_error() {
456 assert!(Archive::open(Cursor::new(b"plain".to_vec()), &ArchiveLimits::default()).is_err());
457 assert_eq!(InputFormat::from_extension("zip"), None);
459 }
460}