create seekable_stream_reader context manager (#107)

Co-authored-by: Julio Martin-Hidalgo <julio.martin-hidalgo@airbus.com>
Co-authored-by: Charly C <changaco@changaco.oy.lc>

Author: 3 people

libarchive/__init__.py

@@ -2,7 +2,8 @@
 from .exception import ArchiveError
 from .extract import extract_fd, extract_file, extract_memory
 from .read import (
-    custom_reader, fd_reader, file_reader, memory_reader, stream_reader
+    custom_reader, fd_reader, file_reader, memory_reader, stream_reader,
+    seekable_stream_reader
 )
 from .write import custom_writer, fd_writer, file_writer, memory_writer
 
@@ -11,5 +12,6 @@
     ArchiveError,
     extract_fd, extract_file, extract_memory,
     custom_reader, fd_reader, file_reader, memory_reader, stream_reader,
+    seekable_stream_reader,
     custom_writer, fd_writer, file_writer, memory_writer
 )]

libarchive/ffi.py

@@ -48,6 +48,9 @@
 READ_CALLBACK = CFUNCTYPE(
     c_ssize_t, c_void_p, c_void_p, POINTER(c_void_p)
 )
+SEEK_CALLBACK = CFUNCTYPE(
+    c_longlong, c_int, c_void_p, c_longlong, c_int
+)
 OPEN_CALLBACK = CFUNCTYPE(c_int, c_void_p, c_void_p)
 CLOSE_CALLBACK = CFUNCTYPE(c_int, c_void_p, c_void_p)
 VOID_CB = lambda *_: ARCHIVE_OK
@@ -226,6 +229,8 @@ def get_write_filter_function(filter_name):
         logger.info(str(e))
         READ_FILTERS.remove(f_name)
 
+ffi('read_set_seek_callback', [c_archive_p, SEEK_CALLBACK], c_int, check_int)
+
 ffi('read_open',
     [c_archive_p, c_void_p, OPEN_CALLBACK, READ_CALLBACK, CLOSE_CALLBACK],
     c_int, check_int)

libarchive/read.py

@@ -6,7 +6,7 @@
 
 from . import ffi
 from .ffi import (ARCHIVE_EOF, OPEN_CALLBACK, READ_CALLBACK, CLOSE_CALLBACK,
-                  VOID_CB, page_size)
+                  SEEK_CALLBACK, VOID_CB, page_size)
 from .entry import ArchiveEntry, new_archive_entry
 
 
@@ -120,3 +120,38 @@ def read_func(archive_p, context, ptrptr):
     with new_archive_read(format_name, filter_name) as archive_p:
         ffi.read_open(archive_p, None, open_cb, read_cb, close_cb)
         yield ArchiveRead(archive_p)
+
+
+@contextmanager
+def seekable_stream_reader(stream, format_name='all', filter_name='all',
+                           block_size=page_size):
+    """Read an archive from a seekable stream.
+
+    The `stream` object must support the standard `readinto`, 'seek' and
+    'tell' methods.
+    """
+    buf = create_string_buffer(block_size)
+    buf_p = cast(buf, c_void_p)
+
+    def read_func(archive_p, context, ptrptr):
+        # readinto the buffer, returns number of bytes read
+        length = stream.readinto(buf)
+        # write the address of the buffer into the pointer
+        ptrptr = cast(ptrptr, POINTER(c_void_p))
+        ptrptr[0] = buf_p
+        # tell libarchive how much data was written into the buffer
+        return length
+
+    def seek_func(archive_p, context, offset, whence):
+        stream.seek(offset, whence)
+        # tell libarchive the current position
+        return stream.tell()
+
+    open_cb = OPEN_CALLBACK(VOID_CB)
+    read_cb = READ_CALLBACK(read_func)
+    seek_cb = SEEK_CALLBACK(seek_func)
+    close_cb = CLOSE_CALLBACK(VOID_CB)
+    with new_archive_read(format_name, filter_name) as archive_p:
+        ffi.read_set_seek_callback(archive_p, seek_cb)
+        ffi.read_open(archive_p, None, open_cb, read_cb, close_cb)
+        yield ArchiveRead(archive_p)

tests/test_rwx.py

@@ -99,6 +99,22 @@ def test_custom_writer_and_stream_reader():
         check_archive(archive, tree)
 
 
+def test_custom_writer_and_seekable_stream_reader():
+    # Collect information on what should be in the archive
+    tree = treestat('libarchive')
+
+    # Create an archive of our libarchive/ directory
+    stream = io.BytesIO()
+    with libarchive.custom_writer(stream.write, '7zip') as archive:
+        archive.add_files('libarchive/')
+    stream.seek(0)
+
+    # Read the archive and check that the data is correct
+    with libarchive.seekable_stream_reader(stream, '7zip') as archive:
+        paths = [entry.name.rstrip('/') for entry in archive]
+        assert sorted(paths) == sorted(tree)
+
+
 @patch('libarchive.ffi.write_fail')
 def test_write_fail(write_fail_mock):
     buf = bytes(bytearray(1000000))