[feature] Add option for computing hashes and storing them as xattrs in the index

Author: mxmlnkn

core/ratarmountcore/mountsource/SQLiteIndexMountSource.py

@@ -2,16 +2,21 @@
 
 import builtins
 import json
+import logging
 import re
 import shutil
+import stat
 import tempfile
 from collections.abc import Iterable
 from typing import IO, Any, Callable, Optional, Union
 
+from ratarmountcore.hashing import compute_hashes
 from ratarmountcore.mountsource import FileInfo, MountSource
 from ratarmountcore.SQLiteIndex import SQLiteIndex
 from ratarmountcore.utils import RatarmountError, overrides
 
+logger = logging.getLogger(__name__)
+
 
 class SQLiteIndexMountSource(MountSource):
     def __init__(
@@ -21,6 +26,7 @@ def __init__(
         clearIndexCache: bool = False,
         checkMetadata: Optional[Callable[[dict[str, Any]], None]] = None,
         transform: Optional[tuple[str, str]] = None,
+        hashes: Optional[list[str]] = None,
         writeIndex: bool = False,
         verifyModificationTime: bool = False,
         indexMinimumFileCount: int = 1000,
@@ -48,6 +54,7 @@ def __init__(
         self.writeIndex = writeIndex
         self.verifyModificationTime = verifyModificationTime
         self.options = options
+        self.hashes = sorted(set(hashes or []))
 
         # Initialize index
         if index is None:
@@ -87,7 +94,7 @@ def __init__(
                 raise RatarmountError(f"Specified file {self.indexFilePath} is not a valid Ratarmount index.")
 
     def _store_default_metadata(self) -> None:
-        argumentsToSave = ['encoding', 'transformPattern']
+        argumentsToSave = ['encoding', 'transformPattern', 'hashes']
         argumentsMetadata = json.dumps(
             {argument: getattr(self, argument) for argument in argumentsToSave if hasattr(self, argument)}
         )
@@ -102,27 +109,63 @@ def _check_metadata(self, metadata: dict[str, Any]) -> None:
 
         if 'arguments' in metadata:
             SQLiteIndex.check_metadata_arguments(
-                json.loads(metadata['arguments']), self, argumentsToCheck=['encoding', 'transformPattern']
+                json.loads(metadata['arguments']), self, argumentsToCheck=['encoding', 'transformPattern', 'hashes']
             )
 
         if 'backendName' not in metadata:
             self.index.try_to_open_first_file(lambda path: self.open(self.lookup(path)))
 
+    def _compute_and_store_hashes(self) -> None:
+        if not self.hashes:
+            return
+
+        # Simply go over all file rows instead of expensive and complicated recursive tree traversal.
+        rows = self.index.get_connection().execute(
+            f'SELECT * {SQLiteIndex.FROM_REGULAR_FILES} AND NOT isgenerated ORDER BY "offsetheader" ASC;'
+        )
+        xattrs: list[tuple[int, str, bytes]] = []
+        for row in rows:
+            fileInfo = self.index._row_to_file_info(row)  # pylint: disable=protected-access
+            if not stat.S_ISREG(fileInfo.mode) or not fileInfo.userdata:
+                continue
+            userData = fileInfo.userdata[-1]
+            if userData.isgenerated:
+                continue
+
+            try:
+                with self.open(fileInfo) as fileObject:
+                    computed = compute_hashes(fileObject, fileInfo.size, self.hashes)
+            except Exception as exception:
+                logger.warning(
+                    "Failed to compute hashes for indexed file with offsetheader=%s: %s",
+                    userData.offsetheader,
+                    exception,
+                )
+                continue
+
+            xattrs += [
+                (userData.offsetheader, f"user.hash.{name}", value.encode('utf-8')) for name, value in computed.items()
+            ]
+            if len(xattrs) >= 1000:
+                self.index.setxattrs(xattrs)
+                xattrs.clear()
+
+        if xattrs:
+            self.index.setxattrs(xattrs)
+
     def _finalize_index(
         self,
         create_index: Callable[[], None],
         *,  # force all parameters after to be keyword-only
         store_metadata: Optional[Callable[[], None]] = None,
         isFileObject: Optional[bool] = None,
     ):
-        """
-        metadata
-            Should either be a list of attributes on 'self' that should be stored or a callable that stores
-            metadata by calling self.index.store_metadata. If it is None a default selection of attributes
-            will be saved.
-        """
+        def create_index_and_post_process():
+            create_index()
+            self._compute_and_store_hashes()
+
         self.index.finalize_index(
-            create_index=create_index,
+            create_index=create_index_and_post_process,
             store_metadata=store_metadata if callable(store_metadata) else self._store_default_metadata,
             isFileObject=isFileObject,
             writeIndex=self.writeIndex,

core/ratarmountcore/mountsource/formats/tar.py

@@ -511,6 +511,7 @@ class SQLiteIndexedTar(SQLiteIndexMountSource):
         'transformRecursiveMountPoint',
         'transformPattern',
         'ignoreZeros',
+        'hashes',
     )
 
     # fmt: off
@@ -793,7 +794,7 @@ def _get_archive_path(self) -> Optional[str]:
         return None if self.tarFileName == '<file object>' else self.tarFileName
 
     def _store_metadata(self) -> None:
-        argumentsToSave = list(self.INDEX_ARGUMENTS_TO_CHECK) + ['gzipSeekPointSpacing']
+        argumentsToSave = [*list(self.INDEX_ARGUMENTS_TO_CHECK), 'gzipSeekPointSpacing']
         argumentsMetadata = json.dumps({argument: getattr(self, argument) for argument in argumentsToSave})
         # The second argument must be a path to a file to call os.stat with, not simply a file name.
         self.index.store_metadata(argumentsMetadata, "" if self.isFileObject else self.tarFileName)

core/tests/test_SQLiteIndexedTar.py

@@ -4,6 +4,7 @@
 
 import bz2
 import concurrent.futures
+import hashlib
 import io
 import os
 import stat
@@ -81,6 +82,9 @@ def test_tar_bz2_with_parallelization(parallelization):
             info = file.lookup('/mimi/00105.tar')
             assert info.userdata[0].offset == 1248256
 
+            # There should especially be no hash xattrs without explicit options being set for that.
+            assert not file.list_xattr(info)
+
     @staticmethod
     def test_recursive_tar_bz2_with_parallelization(parallelization):
         with (
@@ -661,3 +665,63 @@ def test_appending_to_large_archive(parallelization, tmpdir):
             assert indexedTar.exists("/bar")
             assert indexedTar.exists("/folder")
             assert indexedTar.exists("/bar2")
+
+
+@pytest.mark.parallel
+@pytest.mark.parametrize("parallelization", [1])
+class TestSQLiteIndexedTarXattrs:
+    @staticmethod
+    def test_hash_xattrs_for_regular_file(parallelization):
+        with (
+            copy_test_file("single-file.tar") as path,
+            SQLiteIndexedTar(
+                path, writeIndex=False, parallelization=parallelization, hashes=['sha256', 'sha3_256', 'smplayer']
+            ) as indexedTar,
+        ):
+            fileInfo = indexedTar.lookup('/bar')
+            assert fileInfo
+            keys = indexedTar.list_xattr(fileInfo)
+
+            assert 'user.hash.sha256' in keys
+            assert 'user.hash.sha3_256' in keys
+            assert 'user.hash.smplayer' in keys
+
+            data = indexedTar.open(fileInfo).read()
+            assert indexedTar.get_xattr(fileInfo, 'user.hash.sha256') == hashlib.sha256(data).hexdigest().encode(
+                'utf-8'
+            )
+            assert indexedTar.get_xattr(fileInfo, 'user.hash.sha3_256') == hashlib.sha3_256(data).hexdigest().encode(
+                'utf-8'
+            )
+            assert indexedTar.get_xattr(fileInfo, 'user.hash.smplayer')
+
+    @staticmethod
+    def test_hashes_persist_in_written_sidecar_index(parallelization, tmpdir):
+        with copy_test_file('single-file.tar') as tarPath:
+            indexFilePath = os.path.join(tmpdir, 'single-file.tar.index')
+
+            # Open and force index creation.
+            with SQLiteIndexedTar(
+                tarFileName=tarPath,
+                writeIndex=True,
+                clearIndexCache=True,
+                indexFilePath=indexFilePath,
+                parallelization=parallelization,
+                hashes=['sha256'],
+            ) as indexedTar:
+                fileInfo = indexedTar.lookup('/bar')
+                assert fileInfo
+                assert indexedTar.get_xattr(fileInfo, 'user.hash.sha256')
+
+            # Load existing index with stored xattrs.
+            with SQLiteIndexedTar(
+                tarFileName=tarPath,
+                writeIndex=False,
+                clearIndexCache=False,
+                indexFilePath=indexFilePath,
+                parallelization=parallelization,
+            ) as indexedTar:
+                fileInfo = indexedTar.lookup('/bar')
+                assert fileInfo
+                digest = indexedTar.get_xattr(fileInfo, 'user.hash.sha256')
+                assert digest

core/tests/test_ZipMountSource.py

@@ -4,6 +4,7 @@
 # pylint: disable=wrong-import-position
 # pylint: disable=protected-access
 
+import hashlib
 import os
 import shutil
 import stat
@@ -121,6 +122,21 @@ def test_transform():
                 with mountSource.open(mountSource.lookup(filePath)) as file:
                     assert file.read() == b'iriya\n'
 
+    @staticmethod
+    def test_hash_xattrs_for_regular_file():
+        with ZipMountSource(find_test_file('folder-symlink.zip'), hashes=['sha256']) as mountSource:
+            path = '/foo/fighter/ufo'
+            fileInfo = mountSource.lookup(path)
+            assert fileInfo, path
+
+            keys = mountSource.list_xattr(fileInfo)
+            assert 'user.hash.sha256' in keys
+            assert 'user.hash.crc32' not in keys
+
+            with mountSource.open(fileInfo) as file:
+                expected = hashlib.sha256(file.read()).hexdigest().encode('utf-8')
+            assert mountSource.get_xattr(fileInfo, 'user.hash.sha256') == expected
+
 
 def benchmark_fast_zipfile_decryption():
     with tempfile.TemporaryDirectory(suffix=".ratarmount-benchmark") as folder:

ratarmount/CLIHelpers.py

@@ -252,6 +252,7 @@ def parsed_args_to_options(args) -> dict[str, Any]:
         'logFile'                      : args.log_file,
         'enableFileVersions'           : args.file_versions,
         'controlInterface'             : args.control_interface,
+        'hashes'                       : args.hashes,
         'writeIndex'                   : True,
         'mount'                        : args.mount,
     }

ratarmount/cli.py

@@ -16,6 +16,7 @@
 import textwrap
 from typing import Any, Optional
 
+from ratarmountcore.hashing import HASH_REGISTRY
 from ratarmountcore.utils import RatarmountError, get_xdg_cache_home
 
 with contextlib.suppress(ImportError):
@@ -303,6 +304,12 @@ def create_parser(useColor: Optional[bool] = True) -> argparse.ArgumentParser:
              '   --index-folder ~/.ratarmount : will only test ~/.ratarmount as a storage location and nothing else. '
              'Instead, it will first try ~/.ratarmount and the folder "foo,9000". ')
 
+    indexGroup.add_argument(
+        '--hashes', type=str, action=AppendCommaSeparatedOverwriteDefault, default=None,
+        help='Comma-separated list of hashes to compute for indexed files. They are stored as xattrs in the index '
+             'with key user.hash.<hash_name>. Disabled by default because it is expensive. '
+             f'Available algorithms: {", ".join(sorted(HASH_REGISTRY.keys()))}.')
+
     # Recursion Options
 
     # TODO The recursion depth is only heeded by AutoMountLayer but not by SQLiteIndexedTar.

tests/ratarmount-help.txt

@@ -21,6 +21,7 @@
 from ratarmountcore.utils import ceil_div
 
 from ratarmount.cli import cli as ratarmountcli
+from ratarmount.cli import create_parser
 
 try:
     import ext4
@@ -441,3 +442,14 @@ def test_recursive_extensions_patterns(recursiveExtensions, pathInArchive, shoul
         with RunRatarmount(mountPoint, args):
             path = Path(mountPoint) / pathInArchive
             assert path.is_file() == shouldExist
+
+
+def test_hashes_cli_defaults_and_parsing():
+    parser = create_parser(useColor=False)
+
+    def parse(args: list[str]):
+        return parser.parse_args(args).hashes
+
+    assert parse(['--hashes', 'sha256', '--hashes', 'smplayer', 'single-file.tar']) == ['sha256', 'smplayer']
+    assert parse(['--hashes', 'sha256,smplayer', 'single-file.tar']) == ['sha256', 'smplayer']
+    assert parse(['single-file.tar']) is None