From 4450899ad4743813ffcf5e2acdcdb6a316245824 Mon Sep 17 00:00:00 2001
From: morganamilo <morganamilo@archlinux.org>
Date: Tue, 19 Mar 2024 22:23:13 +0000
Subject: [PATCH] Add next_header() to ArchiveIterator

This allows for skipping files without allocating the contents.
---
 src/iterator.rs           | 34 +++++++++++++++++++++++++--
 tests/integration_test.rs | 49 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/src/iterator.rs b/src/iterator.rs
index 5b58bee..16e1d11 100644
--- a/src/iterator.rs
+++ b/src/iterator.rs
@@ -74,7 +74,7 @@ impl<R: Read + Seek> Iterator for ArchiveIterator<R> {
             let next = if self.in_file {
                 unsafe { self.next_data_chunk() }
             } else {
-                unsafe { self.next_header() }
+                unsafe { self.unsafe_next_header() }
             };
 
             match &next {
@@ -108,6 +108,36 @@ impl<R: Read + Seek> Iterator for ArchiveIterator<R> {
     }
 }
 
+impl<R: Read + Seek> ArchiveIterator<R> {
+    pub fn next_header(&mut self) -> Option<ArchiveContents> {
+        debug_assert!(!self.closed);
+
+        if self.error {
+            return None;
+        }
+
+        let next = unsafe { self.unsafe_next_header() };
+
+        match &next {
+            ArchiveContents::StartOfEntry(name, stat) => {
+                if let Some(filter) = &self.filter {
+                    if !filter(name, stat) {
+                        return None;
+                    }
+                }
+
+                self.in_file = true;
+                Some(next)
+            }
+            ArchiveContents::Err(_) => {
+                self.error = true;
+                Some(next)
+            }
+            _ => None,
+        }
+    }
+}
+
 impl<R: Read + Seek> Drop for ArchiveIterator<R> {
     fn drop(&mut self) {
         drop(self.free());
@@ -299,7 +329,7 @@ impl<R: Read + Seek> ArchiveIterator<R> {
         Ok(())
     }
 
-    unsafe fn next_header(&mut self) -> ArchiveContents {
+    unsafe fn unsafe_next_header(&mut self) -> ArchiveContents {
         match ffi::archive_read_next_header(self.archive_reader, &mut self.archive_entry) {
             ffi::ARCHIVE_EOF => ArchiveContents::EndOfEntry,
             ffi::ARCHIVE_OK | ffi::ARCHIVE_WARN => {
diff --git a/tests/integration_test.rs b/tests/integration_test.rs
index 9c12d54..af22d98 100644
--- a/tests/integration_test.rs
+++ b/tests/integration_test.rs
@@ -591,6 +591,55 @@ fn iterate_tar() {
     assert_eq!(contents, expected);
 }
 
+fn collect_iterate_names_with_encoding(
+    source: std::fs::File,
+    decode: DecodeCallback,
+) -> Vec<String> {
+    let mut results = Vec::new();
+
+    let mut iter =
+        ArchiveIterator::from_read_with_encoding(source, decode).expect("Failed to get the file");
+
+    while let Some(content) = iter.next_header() {
+        match content {
+            ArchiveContents::StartOfEntry(file_name, _) => {
+                results.push(file_name);
+            }
+            ArchiveContents::DataChunk(_) => {
+                panic!("expected StartOfntry got DataChunk")
+            }
+            ArchiveContents::EndOfEntry => {
+                panic!("expected StartOfntry got EndOfEntry")
+            }
+            ArchiveContents::Err(e) => panic!("{:?}", e),
+        }
+    }
+
+    iter.close().unwrap();
+
+    results
+}
+
+#[test]
+fn iterate_tar_names() {
+    let source = std::fs::File::open("tests/fixtures/tree.tar").unwrap();
+
+    let contents = collect_iterate_names_with_encoding(source, decode_utf8);
+
+    let expected: Vec<String> = vec![
+        "tree/",
+        "tree/branch1/",
+        "tree/branch1/leaf",
+        "tree/branch2/",
+        "tree/branch2/leaf",
+    ]
+    .into_iter()
+    .map(|a| a.into())
+    .collect();
+
+    assert_eq!(contents, expected);
+}
+
 #[test]
 fn iterate_7z() {
     let source = std::fs::File::open("tests/fixtures/tree.7z").unwrap();