Add a helper function for traversing topologically sorted trees

internal/mock/BUILD.bazel

@@ -6,6 +6,7 @@ gomock(
     out = "aliases.go",
     interfaces = [
         "AEAD",
+        "IntTreeDirectoryVisitor",
         "ReadCloser",
         "ResponseWriter",
         "RoundTripper",

internal/mock/aliases/BUILD.bazel

@@ -5,4 +5,5 @@ go_library(
     srcs = ["aliases.go"],
     importpath = "github.com/buildbarn/bb-storage/internal/mock/aliases",
     visibility = ["//:__subpackages__"],
+    deps = ["//pkg/blobstore"],
 )

internal/mock/aliases/aliases.go

@@ -4,6 +4,8 @@ import (
 	"crypto/cipher"
 	"io"
 	"net/http"
+
+	"github.com/buildbarn/bb-storage/pkg/blobstore"
 )
 
 // This file contains aliases for some of the interfaces provided by the
@@ -15,6 +17,10 @@ import (
 // AEAD is an alias of cipher.AEAD.
 type AEAD = cipher.AEAD
 
+// IntTreeDirectoryVisitor is a TreeDirectoryVisitor that takes integer
+// arguments.
+type IntTreeDirectoryVisitor = blobstore.TreeDirectoryVisitor[int]
+
 // ReadCloser is an alias of io.ReadCloser.
 type ReadCloser = io.ReadCloser
 

pkg/blobstore/BUILD.bazel

@@ -23,6 +23,7 @@ go_library(
         "read_canarying_blob_access.go",
         "reference_expanding_blob_access.go",
         "validation_caching_read_buffer_factory.go",
+        "visit_topologically_sorted_tree.go",
         "zip_reading_blob_access.go",
         "zip_writing_blob_access.go",
     ],
@@ -68,6 +69,7 @@ go_test(
         "read_canarying_blob_access_test.go",
         "reference_expanding_blob_access_test.go",
         "validation_caching_read_buffer_factory_test.go",
+        "visit_topologically_sorted_tree_test.go",
         "zip_reading_blob_access_test.go",
         "zip_writing_blob_access_test.go",
     ],
@@ -87,6 +89,8 @@ go_test(
         "@com_github_stretchr_testify//require",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//status",
+        "@org_golang_google_protobuf//encoding/protowire",
+        "@org_golang_google_protobuf//proto",
         "@org_golang_google_protobuf//types/known/timestamppb",
     ],
 )

pkg/blobstore/visit_topologically_sorted_tree.go

@@ -0,0 +1,105 @@
+package blobstore
+
+import (
+	"io"
+
+	remoteexecution "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
+	"github.com/buildbarn/bb-storage/pkg/blobstore/buffer"
+	"github.com/buildbarn/bb-storage/pkg/digest"
+	"github.com/buildbarn/bb-storage/pkg/util"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/encoding/protowire"
+	"google.golang.org/protobuf/proto"
+)
+
+// TreeDirectoryVisitor is a callback type that is invoked by
+// VisitTopologicallySortedTree for each Directory message contained in
+// the REv2 Tree object.
+type TreeDirectoryVisitor[TArgument any] func(d *remoteexecution.Directory, argument *TArgument, childArguments []*TArgument) error
+
+// VisitTopologicallySortedTree iterates over all Directory messages
+// contained in an REv2 Tree object. For each directory, the visitor is
+// capable of tracking state, which it can also build up when parent
+// directories are visited.
+//
+// This function expects the REv2 Tree object to be topologically
+// sorted, with parents being stored before children. As a result,
+// directories are also visited in topological order.
+func VisitTopologicallySortedTree[TArgument any](r io.Reader, digestFunction digest.Function, maximumDirectorySizeBytes int, rootArgument *TArgument, visitDirectory TreeDirectoryVisitor[TArgument]) error {
+	expectedFieldNumber := TreeRootFieldNumber
+	expectedDirectories := map[digest.Digest]*TArgument{}
+	if err := util.VisitProtoBytesFields(r, func(fieldNumber protowire.Number, offsetBytes, sizeBytes int64, fieldReader io.Reader) error {
+		if fieldNumber != expectedFieldNumber {
+			return status.Errorf(codes.InvalidArgument, "Expected field number %d", expectedFieldNumber)
+		}
+		expectedFieldNumber = TreeChildrenFieldNumber
+
+		var directoryMessage proto.Message
+		var errUnmarshal error
+		var argument *TArgument
+		switch fieldNumber {
+		case TreeRootFieldNumber:
+			directoryMessage, errUnmarshal = buffer.NewProtoBufferFromReader(
+				&remoteexecution.Directory{},
+				io.NopCloser(fieldReader),
+				buffer.UserProvided,
+			).ToProto(&remoteexecution.Directory{}, maximumDirectorySizeBytes)
+			argument = rootArgument
+		case TreeChildrenFieldNumber:
+			b1, b2 := buffer.NewProtoBufferFromReader(
+				&remoteexecution.Directory{},
+				io.NopCloser(fieldReader),
+				buffer.UserProvided,
+			).CloneCopy(maximumDirectorySizeBytes)
+
+			digestGenerator := digestFunction.NewGenerator(sizeBytes)
+			if err := b1.IntoWriter(digestGenerator); err != nil {
+				b2.Discard()
+				return err
+			}
+			directoryDigest := digestGenerator.Sum()
+
+			var ok bool
+			argument, ok = expectedDirectories[directoryDigest]
+			if !ok {
+				b2.Discard()
+				return status.Errorf(codes.InvalidArgument, "Directory has digest %#v, which was not expected", directoryDigest.String())
+			}
+			delete(expectedDirectories, directoryDigest)
+
+			directoryMessage, errUnmarshal = b2.ToProto(&remoteexecution.Directory{}, maximumDirectorySizeBytes)
+		}
+		if errUnmarshal != nil {
+			return errUnmarshal
+		}
+
+		directory := directoryMessage.(*remoteexecution.Directory)
+		childArguments := make([]*TArgument, 0, len(directory.Directories))
+		for _, childDirectory := range directory.Directories {
+			digest, err := digestFunction.NewDigestFromProto(childDirectory.Digest)
+			if err != nil {
+				return util.StatusWrapf(err, "Invalid digest for child directory %#v", childDirectory.Name)
+			}
+			childArgument, ok := expectedDirectories[digest]
+			if !ok {
+				childArgument = new(TArgument)
+				expectedDirectories[digest] = childArgument
+			}
+			childArguments = append(childArguments, childArgument)
+		}
+
+		return visitDirectory(directory, argument, childArguments)
+	}); err != nil {
+		return err
+	}
+
+	if expectedFieldNumber == TreeRootFieldNumber {
+		return status.Error(codes.InvalidArgument, "Tree does not contain any directories")
+	}
+	if len(expectedDirectories) > 0 {
+		return status.Errorf(codes.InvalidArgument, "At least %d more directories were expected", len(expectedDirectories))
+	}
+	return nil
+}