Also inline 'addfile' function of tarfile

Author: Joost Jansen

acquire/outputs/tar.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
+import copy
 import io
 import shutil
 import tarfile
 from typing import TYPE_CHECKING, BinaryIO
-from unittest.mock import patch
 
 from acquire.crypt import EncryptedStream
 from acquire.outputs.base import Output
@@ -17,43 +17,6 @@
 TAR_COMPRESSION_METHODS = {"gzip": "gz", "bzip2": "bz2", "xz": "xz"}
 
 
-def copyfileobj(
-    src: BinaryIO, dst: BinaryIO, length: int | None = None, _: Exception = OSError, bufsize: int | None = None
-) -> None:
-    """
-    Patched version of the copyfileobj function from the Python stdlib (tarfile.py),
-    to handle cases where the source file is actually shorter than expected.
-    By patching the missing bytes with zeroes, we avoid raising an exception
-    and potentially corrupting the tar file.
-    """
-    bufsize = bufsize or 16 * 1024
-    if length == 0:
-        return
-    if length is None:
-        shutil.copyfileobj(src, dst, bufsize)
-        return
-
-    blocks, remainder = divmod(length, bufsize)
-    for _ in range(blocks):
-        # Already prevents "long reads" because it reads at max bufsize bytes at a time
-        buf = src.read(bufsize)
-        if len(buf) < bufsize:
-            # raise exception("unexpected end of data")
-            # PATCH; instead of raising an exception, pad the data to the desired length
-            buf += b"\x00" * (bufsize - len(buf))
-        dst.write(buf)
-
-    if remainder != 0:
-        # Already prevents "long reads" because it reads at max bufsize bytes at a time
-        buf = src.read(remainder)
-        if len(buf) < remainder:
-            # raise exception("unexpected end of data")
-            # PATCH; instead of raising an exception, pad the data to the desired length
-            buf += b"\x00" * (remainder - len(buf))
-        dst.write(buf)
-    return
-
-
 class TarOutput(Output):
     """Tar archive acquire output format. Output can be compressed and/or encrypted.
 
@@ -139,8 +102,54 @@ def write(
             if stat:
                 info.mtime = stat.st_mtime
 
-        with patch("tarfile.copyfileobj", copyfileobj):
-            self.tar.addfile(info, fh)
+        # Inline version of Python stdlib's tarfile.addfile & tarfile.copyfileobj,
+        # to allow for padding and more control over the tar file writing.
+        self.tar._check("awx")
+
+        if fh is None and info.isreg() and info.size != 0:
+            raise ValueError("fileobj not provided for non zero-size regular file")
+
+        info = copy.copy(info)
+
+        buf = info.tobuf(self.tar.format, self.tar.encoding, self.tar.errors)
+        self.tar.fileobj.write(buf)
+        self.tar.offset += len(buf)
+        bufsize = self.tar.copybufsize
+        if fh is not None:
+            bufsize = bufsize or 16 * 1024
+
+            if info.size == 0:
+                return
+            if info.size is None:
+                shutil.copyfileobj(fh, self.tar.fileobj, bufsize)
+                return
+
+            blocks, remainder = divmod(info.size, bufsize)
+            for _ in range(blocks):
+                # Prevents "long reads" because it reads at max bufsize bytes at a time
+                buf = fh.read(bufsize)
+                if len(buf) < bufsize:
+                    # raise exception("unexpected end of data")
+                    # PATCH; instead of raising an exception, pad the data to the desired length
+                    buf += tarfile.NUL * (bufsize - len(buf))
+                self.tar.fileobj.write(buf)
+
+            if remainder != 0:
+                # Prevents "long reads" because it reads at max bufsize bytes at a time
+                buf = fh.read(remainder)
+                if len(buf) < remainder:
+                    # raise exception("unexpected end of data")
+                    # PATCH; instead of raising an exception, pad the data to the desired length
+                    buf += tarfile.NUL * (remainder - len(buf))
+                self.tar.fileobj.write(buf)
+
+            blocks, remainder = divmod(info.size, tarfile.BLOCKSIZE)
+            if remainder > 0:
+                self.tar.fileobj.write(tarfile.NUL * (tarfile.BLOCKSIZE - remainder))
+                blocks += 1
+            self.tar.offset += blocks * tarfile.BLOCKSIZE
+
+        self.tar.members.append(info)
 
     def close(self) -> None:
         """Closes the tar file."""

tests/test_outputs_tar.py

@@ -82,7 +82,7 @@ def read(self, size: int) -> bytes:
 
     content = b"some text"
 
-    content_padded = content[:-5] + b"\x00" * 5
+    content_padded = content[:-5] + tarfile.NUL * 5
     file = ShrinkingFile(content)
 
     tar_output = TarOutput(tmp_path / "race.tar", encrypt=True, public_key=public_key)