Add blake3 hashing (#538)

Because BLAKE3 natively supports parallelism without changing the final
hash, sharding is bypassed. This is much more useful than getting
different file hashes depending on which hashing method you used.

The BLAKE3 hashing is done by memory mapping the file, and defaults to
the max number of workers which is the number of logical CPU cores.
This is a good default and the most performant setup. It is also what
the standard BLAKE3 CLI tool (b3sum) does. It is implemented in Rust
and so will be true parallelism rather than the thread concurrency
implemented for other hashing algorithms, so the speed up should be
quite large. But it will likely be slower on HDDs than having no
parallelism. I think this is the right default, but the HDD concern
is documented.

Resolves: #530

Signed-off-by: makeworld <[email protected]>

Author: makew0rld

CHANGELOG.md

@@ -29,6 +29,7 @@ All versions prior to 1.0.0 are untracked.
 - Implemented a new, minimal container image. This variant excludes optional dependencies (like OTel and PKCS#11) to reduce footprint, focusing solely on core signing and verification mechanisms.
 - The library now requires at least v4.0.0 of `sigstore` due to breaking changes in that library
 - Added support for signing and verifying using private Sigstore instances (`--trust_config`)
+- Added support for BLAKE3 hashing
 
 ## [1.0.1] - 2024-04-18
 

benchmarks/exp_hash.py

@@ -48,7 +48,7 @@ def build_parser() -> argparse.ArgumentParser:
         help="hash methods to benchmark",
         nargs="+",
         type=str,
-        default=["sha256", "blake2"],
+        default=["sha256", "blake2", "blake3"],
     )
 
     parser.add_argument(
@@ -74,6 +74,8 @@ def _get_hasher(hash_algorithm: str) -> hashing.StreamingHashEngine:
         return memory.SHA256()
     if hash_algorithm == "blake2":
         return memory.BLAKE2()
+    if hash_algorithm == "blake3":
+        return memory.BLAKE3()
 
     raise ValueError(f"Cannot convert {hash_algorithm} to a hash engine")
 

benchmarks/serialize.py

@@ -46,6 +46,8 @@ def get_hash_engine_factory(
         return memory.SHA256
     if hash_algorithm == "blake2":
         return memory.BLAKE2
+    if hash_algorithm == "blake3":
+        return memory.BLAKE3
 
     raise ValueError(f"Cannot convert {hash_algorithm} to a hash engine")
 
@@ -146,7 +148,7 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--hash_method",
         help="hash method to use (default: sha256)",
-        choices=["sha256", "blake2"],
+        choices=["sha256", "blake2", "blake3"],
         default="sha256",
     )
     parser.add_argument(
@@ -180,7 +182,7 @@ def build_parser() -> argparse.ArgumentParser:
         "--merge_hasher",
         help="hasher to use to merge individual hashes "
         "when skipping manifest creation (default: sha256)",
-        choices=["sha256", "blake2"],
+        choices=["sha256", "blake2", "blake3"],
         default="sha256",
     )
 

pyproject.toml

@@ -28,6 +28,7 @@ classifiers = [
 ]
 dependencies = [
   "asn1crypto",
+  "blake3",
   "click",
   "cryptography",
   "in-toto-attestation",

src/model_signing/_hashing/io.py

@@ -38,6 +38,7 @@
 import pathlib
 from typing import Optional
 
+import blake3
 from typing_extensions import override
 
 from model_signing._hashing import hashing
@@ -135,6 +136,61 @@ def digest_size(self) -> int:
         return self._content_hasher.digest_size
 
 
+class Blake3FileHasher(FileHasher):
+    """Simple file hash engine that uses BLAKE3 in parallel.
+
+    This hash engine uses the fastest BLAKE3 settings, by using memory mapping
+    and multiple workers. This will greatly increase speed on SSDs, but may
+    not perform well on HDDs. For HDDs, you can set max_threads to 1.
+    """
+
+    def __init__(
+        self,
+        file: pathlib.Path,
+        *,
+        max_threads: int = blake3.blake3.AUTO,
+        digest_name_override: Optional[str] = None,
+    ):
+        """Initializes an instance to hash a file.
+
+        Args:
+            file: The file to hash. Use `set_file` to reset it.
+            max_threads: how many BLAKE3 workers to use. Defaults to number of
+              logical cores.
+            digest_name_override: Optional string to allow overriding the
+              `digest_name` property to support shorter, standardized names.
+        """
+        self._file = file
+        self._digest_name_override = digest_name_override
+        self._blake3 = blake3.blake3(max_threads=max_threads)
+
+    def set_file(self, file: pathlib.Path) -> None:
+        """Redefines the file to be hashed in `compute`.
+
+        Args:
+            file: The new file to be hashed.
+        """
+        self._file = file
+
+    @property
+    @override
+    def digest_name(self) -> str:
+        if self._digest_name_override is not None:
+            return self._digest_name_override
+        return "blake3"
+
+    @override
+    def compute(self) -> hashing.Digest:
+        self._blake3.reset()
+        self._blake3.update_mmap(self._file)
+        return hashing.Digest(self.digest_name, self._blake3.digest())
+
+    @property
+    @override
+    def digest_size(self) -> int:
+        return 32
+
+
 class ShardedFileHasher(SimpleFileHasher):
     """File hash engine that hashes a portion (shard) of the file.
 

src/model_signing/_hashing/memory.py

@@ -37,6 +37,7 @@
 
 import hashlib
 
+import blake3
 from typing_extensions import override
 
 from model_signing._hashing import hashing
@@ -108,3 +109,37 @@ def digest_name(self) -> str:
     @override
     def digest_size(self) -> int:
         return self._hasher.digest_size
+
+
+class BLAKE3(hashing.StreamingHashEngine):
+    """A wrapper around `blake3.blake3`."""
+
+    def __init__(self, initial_data: bytes = b""):
+        """Initializes an instance of a BLAKE3 hash engine.
+
+        Args:
+            initial_data: Optional initial data to hash.
+        """
+        self._hasher = blake3.blake3(initial_data)
+
+    @override
+    def update(self, data: bytes) -> None:
+        self._hasher.update(data)
+
+    @override
+    def reset(self, data: bytes = b"") -> None:
+        self._hasher = blake3.blake3(data)
+
+    @override
+    def compute(self) -> hashing.Digest:
+        return hashing.Digest(self.digest_name, self._hasher.digest())
+
+    @property
+    @override
+    def digest_name(self) -> str:
+        return "blake3"
+
+    @property
+    @override
+    def digest_size(self) -> int:
+        return self._hasher.digest_size

src/model_signing/_serialization/file.py

@@ -66,6 +66,7 @@ def __init__(
         self._serialization_description = manifest._FileSerialization(
             hasher.digest_name, self._allow_symlinks, self._ignore_paths
         )
+        self._is_blake3 = hasher.digest_name == "blake3"
 
     def set_allow_symlinks(self, allow_symlinks: bool) -> None:
         """Set whether following symlinks is allowed."""
@@ -120,7 +121,8 @@ def serialize(
 
         manifest_items = []
         with concurrent.futures.ThreadPoolExecutor(
-            max_workers=self._max_workers
+            # blake3 parallelizes internally
+            max_workers=1 if self._is_blake3 else self._max_workers
         ) as tpe:
             futures = [
                 tpe.submit(self._compute_hash, model_path, path)

src/model_signing/hashing.py

@@ -53,6 +53,8 @@
 import sys
 from typing import Literal, Optional, Union
 
+import blake3
+
 from model_signing import manifest
 from model_signing._hashing import hashing
 from model_signing._hashing import io
@@ -124,8 +126,8 @@ class Config:
     default, we hash at file level granularity.
 
     This configuration class also supports configuring the hash method used to
-    generate the hash for every object in the model. We currently support SHA256
-    and BLAKE2, with SHA256 being the default.
+    generate the hash for every object in the model. We currently support
+    SHA256, BLAKE2, and BLAKE3, with SHA256 being the default.
 
     This configuration class also supports configuring which paths from the
     model directory should be ignored. These are files that doesn't impact the
@@ -183,7 +185,8 @@ def hash(
         )
 
     def _build_stream_hasher(
-        self, hashing_algorithm: Literal["sha256", "blake2"] = "sha256"
+        self,
+        hashing_algorithm: Literal["sha256", "blake2", "blake3"] = "sha256",
     ) -> hashing.StreamingHashEngine:
         """Builds a streaming hasher from a constant string.
 
@@ -198,28 +201,37 @@ def _build_stream_hasher(
             return memory.SHA256()
         if hashing_algorithm == "blake2":
             return memory.BLAKE2()
+        if hashing_algorithm == "blake3":
+            return memory.BLAKE3()
 
         raise ValueError(f"Unsupported hashing method {hashing_algorithm}")
 
     def _build_file_hasher_factory(
         self,
-        hashing_algorithm: Literal["sha256", "blake2"] = "sha256",
+        hashing_algorithm: Literal["sha256", "blake2", "blake3"] = "sha256",
         chunk_size: int = 1048576,
-    ) -> Callable[[pathlib.Path], io.SimpleFileHasher]:
+        max_workers: Optional[int] = None,
+    ) -> Callable[[pathlib.Path], io.FileHasher]:
         """Builds the hasher factory for a serialization by file.
 
         Args:
             hashing_algorithm: The hashing algorithm to use to hash a file.
             chunk_size: The amount of file to read at once. Default is 1MB. A
               special value of 0 signals to attempt to read everything in a
-              single call.
+              single call. This is ignored for BLAKE3.
+            max_workers: Maximum number of workers to use in parallel. Defaults
+              to the number of logical cores. Only relevant for BLAKE3.
 
         Returns:
             The hasher factory that should be used by the active serialization
             method.
         """
+        if max_workers is None:
+            max_workers = blake3.blake3.AUTO
 
-        def _factory(path: pathlib.Path) -> io.SimpleFileHasher:
+        def _factory(path: pathlib.Path) -> io.FileHasher:
+            if hashing_algorithm == "blake3":
+                return io.Blake3FileHasher(path, max_threads=max_workers)
             hasher = self._build_stream_hasher(hashing_algorithm)
             return io.SimpleFileHasher(path, hasher, chunk_size=chunk_size)
 
@@ -233,6 +245,9 @@ def _build_sharded_file_hasher_factory(
     ) -> Callable[[pathlib.Path, int, int], io.ShardedFileHasher]:
         """Builds the hasher factory for a serialization by file shards.
 
+        This is not recommended for BLAKE3 because it is not necessary. BLAKE3
+        already operates in parallel.
+
         Args:
             hashing_algorithm: The hashing algorithm to use to hash a shard.
             chunk_size: The amount of file to read at once. Default is 1MB. A
@@ -263,7 +278,7 @@ def _factory(
     def use_file_serialization(
         self,
         *,
-        hashing_algorithm: Literal["sha256", "blake2"] = "sha256",
+        hashing_algorithm: Literal["sha256", "blake2", "blake3"] = "sha256",
         chunk_size: int = 1048576,
         max_workers: Optional[int] = None,
         allow_symlinks: bool = False,
@@ -279,10 +294,13 @@ def use_file_serialization(
             hashing_algorithm: The hashing algorithm to use to hash a file.
             chunk_size: The amount of file to read at once. Default is 1MB. A
               special value of 0 signals to attempt to read everything in a
-              single call.
+              single call. Ignored for BLAKE3.
             max_workers: Maximum number of workers to use in parallel. Default
               is to defer to the `concurrent.futures` library to select the best
-              value for the current machine.
+              value for the current machine, or the number of logical cores
+              when doing BLAKE3 hashing. When reading files off of slower
+              hardware like an HDD rather than an SSD, and using BLAKE3,
+              setting max_workers to 1 may improve performance.
             allow_symlinks: Controls whether symbolic links are included. If a
               symlink is present but the flag is `False` (default) the
               serialization would raise an error.
@@ -291,7 +309,9 @@ def use_file_serialization(
             The new hashing configuration with the new serialization method.
         """
         self._serializer = file.Serializer(
-            self._build_file_hasher_factory(hashing_algorithm, chunk_size),
+            self._build_file_hasher_factory(
+                hashing_algorithm, chunk_size, max_workers
+            ),
             max_workers=max_workers,
             allow_symlinks=allow_symlinks,
             ignore_paths=ignore_paths,
@@ -301,7 +321,7 @@ def use_file_serialization(
     def use_shard_serialization(
         self,
         *,
-        hashing_algorithm: Literal["sha256", "blake2"] = "sha256",
+        hashing_algorithm: Literal["sha256", "blake2", "blake3"] = "sha256",
         chunk_size: int = 1048576,
         shard_size: int = 1_000_000_000,
         max_workers: Optional[int] = None,
@@ -310,6 +330,10 @@ def use_shard_serialization(
     ) -> Self:
         """Configures serialization to build a manifest of (shard, hash) pairs.
 
+        For BLAKE3 this is equivalent to not sharding. Sharding is bypassed
+        because BLAKE3 already operates in parallel. This means the chunk_size
+        and shard_size args are ignored.
+
         The serialization method in this configuration is changed to one where
         every file in the model is sharded in equal sized shards, every shard is
         paired with its digest and a manifest containing all these pairings is
@@ -332,6 +356,15 @@ def use_shard_serialization(
         Returns:
             The new hashing configuration with the new serialization method.
         """
+        if hashing_algorithm == "blake3":
+            return self.use_file_serialization(
+                hashing_algorithm=hashing_algorithm,
+                chunk_size=chunk_size,
+                max_workers=max_workers,
+                allow_symlinks=allow_symlinks,
+                ignore_paths=ignore_paths,
+            )
+
         self._serializer = file_shard.Serializer(
             self._build_sharded_file_hasher_factory(
                 hashing_algorithm, chunk_size, shard_size