Add blake3 support

buchgr · Jan 2, 2024 · 302aa4e · 302aa4e
1 parent 282245d
commit 302aa4e
Show file tree

Hide file tree

Showing 4 changed files with 148 additions and 2 deletions.
diff --git a/cache/disk/disk_test.go b/cache/disk/disk_test.go
@@ -416,6 +416,35 @@ func TestCacheExistingFiles(t *testing.T) {
 			"raw/733e21b37cef883579a88183eed0d00cdeea0b59e1bcd77db6957f881c3a6b54",
 			"raw.v2/73/733e21b37cef883579a88183eed0d00cdeea0b59e1bcd77db6957f881c3a6b54-123456789",
 		},
+
+		{
+			pb.DigestFunction_BLAKE3,
+			"hej",
+			"70514ff271364b5a63dc74edd2860b0bc54bdc3cc6b95bf4ea95554fdd08ba16",
+			"cas/blake3/70514ff271364b5a63dc74edd2860b0bc54bdc3cc6b95bf4ea95554fdd08ba16",
+			"cas.v2/blake3/79/70514ff271364b5a63dc74edd2860b0bc54bdc3cc6b95bf4ea95554fdd08ba16-3-123456789",
+		},
+		{
+			pb.DigestFunction_BLAKE3,
+			"världen",
+			"a01628acbe112a59c4312ba1bcd8628170fe90284a84848da598ffdf9ca04d17",
+			"cas/blake3/a01628acbe112a59c4312ba1bcd8628170fe90284a84848da598ffdf9ca04d17",
+			"cas.v2/blake3/a0/a01628acbe112a59c4312ba1bcd8628170fe90284a84848da598ffdf9ca04d17-8-123456789",
+		},
+		{
+			pb.DigestFunction_BLAKE3,
+			"foo",
+			"04e0bb39f30b1a3feb89f536c93be15055482df748674b00d26e5a75777702e9",
+			"ac/blake3/04e0bb39f30b1a3feb89f536c93be15055482df748674b00d26e5a75777702e9",
+			"ac.v2/blake3/04/04e0bb39f30b1a3feb89f536c93be15055482df748674b00d26e5a75777702e9-123456789",
+		},
+		{
+			pb.DigestFunction_BLAKE3,
+			"bar",
+			"f2e897eed7d206cd855d441598fa521abc75aa96953e97c030c9612c30c1293d",
+			"raw/blake3/f2e897eed7d206cd855d441598fa521abc75aa96953e97c030c9612c30c1293d",
+			"raw.v2/blake3/f2/f2e897eed7d206cd855d441598fa521abc75aa96953e97c030c9612c30c1293d-123456789",
+		},
 	}
 
 	var err error
@@ -468,8 +497,8 @@ func TestCacheExistingFiles(t *testing.T) {
 		origOnEvict(key, value)
 	}
 
-	if testCache.lru.Len() != 4 {
-		t.Fatal("expected four items in the cache, found", testCache.lru.Len())
+	if testCache.lru.Len() != len(items) {
+		t.Fatalf("expected %d items in the cache, found %d", len(items), testCache.lru.Len())
 	}
 
 	// Adding new blobs should eventually evict the oldest (items[0]).

diff --git a/cache/hashing/BUILD.bazel b/cache/hashing/BUILD.bazel
@@ -3,6 +3,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 go_library(
     name = "go_default_library",
     srcs = [
+        "blake3.go",
         "hashing.go",
         "sha256.go",
     ],

diff --git a/cache/hashing/blake3.go b/cache/hashing/blake3.go
@@ -0,0 +1,66 @@
+package hashing
+
+import (
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"hash"
+	"regexp"
+
+	pb "github.com/buchgr/bazel-remote/v2/genproto/build/bazel/remote/execution/v2"
+	"lukechampine.com/blake3"
+)
+
+func init() {
+	hasher := &b3Hasher{}
+	register(hasher.DigestFunction(), hasher)
+}
+
+var b3Regex = regexp.MustCompile("^[a-f0-9]{64}$")
+
+type b3Hasher struct{}
+
+func (d *b3Hasher) New() hash.Hash {
+	return blake3.New(d.Size(), nil)
+}
+
+func (d *b3Hasher) Hash(data []byte) string {
+	sum := blake3.Sum256(data)
+	return hex.EncodeToString(sum[:])
+}
+
+func (d *b3Hasher) DigestFunction() pb.DigestFunction_Value {
+	return pb.DigestFunction_BLAKE3
+}
+
+func (d *b3Hasher) Dir() string {
+	return "blake3"
+}
+
+func (d *b3Hasher) Empty() string {
+	return "af1349b9f5f9a1a6a0404dea36dcc9499bcb25c9adc112b7cc9a93cae41f3262"
+}
+
+func (d *b3Hasher) Size() int {
+	return 32
+}
+
+func (d *b3Hasher) Validate(value string) error {
+	if d.Size()*2 != len(value) {
+		return fmt.Errorf("Invalid blake3 hash length %d: expected %d", len(value), d.Size())
+	}
+	if !b3Regex.MatchString(value) {
+		return errors.New("Malformed blake3 hash " + value)
+	}
+	return nil
+}
+
+func (d *b3Hasher) ValidateDigest(hash string, size int64) error {
+	if size == int64(0) {
+		if hash == d.Empty() {
+			return nil
+		}
+		return fmt.Errorf("Invalid zero-length %s hash", d.DigestFunction())
+	}
+	return d.Validate(hash)
+}
diff --git a/utils/resourcename/resourcename_test.go b/utils/resourcename/resourcename_test.go
@@ -183,6 +183,24 @@ func TestParseReadResource(t *testing.T) {
 			resourceName: "pretenduuid/compressed-blobs/IDENTITY/0123456789012345678901234567890123456789012345678901234567890123/42",
 			expectError:  true,
 		},
+
+		// Contains digest function
+		{
+			"blobs/blake3/0123456789012345678901234567890123456789012345678901234567890123/42",
+			"0123456789012345678901234567890123456789012345678901234567890123",
+			42,
+			casblob.Identity,
+			pb.DigestFunction_BLAKE3,
+			false,
+		},
+		{
+			"compressed-blobs/zstd/blake3/0123456789012345678901234567890123456789012345678901234567890123/42",
+			"0123456789012345678901234567890123456789012345678901234567890123",
+			42,
+			casblob.Zstandard,
+			pb.DigestFunction_BLAKE3,
+			false,
+		},
 	}
 
 	for _, tc := range tcs {
@@ -395,6 +413,38 @@ func TestParseWriteResource(t *testing.T) {
 		},
 
 		// Contains digest function
+		{
+			"uploads/pretenduuid/blobs/blake3/0123456789012345678901234567890123456789012345678901234567890123/42",
+			"0123456789012345678901234567890123456789012345678901234567890123",
+			42,
+			casblob.Identity,
+			pb.DigestFunction_BLAKE3,
+			false,
+		},
+		{
+			"uploads/pretenduuid/compressed-blobs/zstd/blake3/0123456789012345678901234567890123456789012345678901234567890123/42",
+			"0123456789012345678901234567890123456789012345678901234567890123",
+			42,
+			casblob.Zstandard,
+			pb.DigestFunction_BLAKE3,
+			false,
+		},
+		{
+			"uploads/pretenduuid/blobs/blake3/0123456789012345678901234567890123456789012345678901234567890123/42/metadata",
+			"0123456789012345678901234567890123456789012345678901234567890123",
+			42,
+			casblob.Identity,
+			pb.DigestFunction_BLAKE3,
+			false,
+		},
+		{
+			"uploads/pretenduuid/compressed-blobs/zstd/blake3/0123456789012345678901234567890123456789012345678901234567890123/42/metadata",
+			"0123456789012345678901234567890123456789012345678901234567890123",
+			42,
+			casblob.Zstandard,
+			pb.DigestFunction_BLAKE3,
+			false,
+		},
 		{
 			resourceName: "uploads/pretenduuid/blobs/foo/0123456789012345678901234567890123456789012345678901234567890123/42",
 			expectError:  true,