WS3: guard the unbounded tree paths (GetTree, AC Tree validation)

cache/disk/disk.go

@@ -93,6 +93,12 @@ type diskCache struct {
 	// observation, D15). Lazily initialized; see lruCaptureBudgetBytes.
 	lruCaptureSem *semaphore.Weighted
 
+	// Cap on total declared output-directory Tree bytes read into memory
+	// per ActionResult validation; see WithTreeValidationSizeLimit.
+	// Zero or negative disables the cap.
+	maxTreeValidationBytes int64
+	treeValidationExceeded func(declaredBytes int64)
+
 	mu  sync.Mutex
 	lru SizedLRU
 
@@ -1008,6 +1014,25 @@ func (c *diskCache) GetValidatedActionResult(ctx context.Context, hash string) (
 
 	pendingValidations := []*pb.Digest{}
 
+	// Pre-flight guard on validation memory (see
+	// WithTreeValidationSizeLimit): every referenced Tree blob is read
+	// wholly into memory below, and their sizes are declared in the
+	// ActionResult, so the total is knowable before allocating anything.
+	// Over-cap results are reported as a miss - semantically safe, the
+	// client rebuilds - and the trip is surfaced via the callback.
+	if c.maxTreeValidationBytes > 0 && len(result.OutputDirectories) > 0 {
+		var declaredTreeBytes int64
+		for _, d := range result.OutputDirectories {
+			declaredTreeBytes += d.TreeDigest.SizeBytes
+		}
+		if declaredTreeBytes > c.maxTreeValidationBytes {
+			if c.treeValidationExceeded != nil {
+				c.treeValidationExceeded(declaredTreeBytes)
+			}
+			return nil, nil, nil // aka "not found"
+		}
+	}
+
 	// treeLeafHashes collects output_directories Tree blob hashes for the LRU
 	// closure. The Tree blobs are recorded as closure leaves (the sweep must
 	// keep them) but are deliberately kept OUT of pendingValidations: their

cache/disk/options.go

@@ -75,6 +75,25 @@ func WithMaxEntries(n int64) Option {
 	}
 }
 
+// WithTreeValidationSizeLimit caps the total declared bytes of the
+// output-directory Tree blobs that GetValidatedActionResult will read into
+// memory while validating one ActionResult. Validation buffers each
+// referenced Tree blob wholly (read + unmarshal), so without a cap a single
+// action with a huge output directory pins an arbitrarily large allocation
+// per concurrent GetActionResult. An over-cap ActionResult is reported as a
+// cache miss - always semantically safe, the client rebuilds - rather than
+// an error. exceeded, when non-nil, is invoked once per capped validation
+// with the total declared Tree bytes; a trip means a client was just forced
+// to rebuild, so it should be wired to a visible metric. maxBytes <= 0
+// disables the cap.
+func WithTreeValidationSizeLimit(maxBytes int64, exceeded func(declaredBytes int64)) Option {
+	return func(c *CacheConfig) error {
+		c.diskCache.maxTreeValidationBytes = maxBytes
+		c.diskCache.treeValidationExceeded = exceeded
+		return nil
+	}
+}
+
 func WithMaxBlobSize(size int64) Option {
 	return func(c *CacheConfig) error {
 		if size <= 0 {

cache/disk/tree_validation_guard_test.go

@@ -0,0 +1,176 @@
+package disk
+
+// Tests for WithTreeValidationSizeLimit: GetValidatedActionResult reads every
+// referenced output-directory Tree blob wholly into memory, so a cap on the
+// total declared Tree bytes is its only per-request memory bound. Over-cap
+// results must be reported as a miss (semantically safe: the client
+// rebuilds), never as an error.
+
+import (
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"os"
+	"testing"
+
+	pb "github.com/buchgr/bazel-remote/v2/genproto/build/bazel/remote/execution/v2"
+
+	"github.com/buchgr/bazel-remote/v2/cache"
+	testutils "github.com/buchgr/bazel-remote/v2/utils"
+
+	"google.golang.org/protobuf/proto"
+)
+
+// putTreeBackedActionResult stores a CAS directory tree and an ActionResult
+// referencing it as an output directory, returning the AC hash and the
+// declared Tree blob size.
+func putTreeBackedActionResult(ctx context.Context, t *testing.T, c *diskCache) (string, int64) {
+	t.Helper()
+
+	fileData := []byte("tree validation guard file contents")
+	fileHash := sha256.Sum256(fileData)
+	fileHashStr := hex.EncodeToString(fileHash[:])
+	if err := c.Put(ctx, cache.CAS, fileHashStr, int64(len(fileData)),
+		bytes.NewReader(fileData)); err != nil {
+		t.Fatal(err)
+	}
+
+	rootDir := pb.Directory{
+		Files: []*pb.FileNode{
+			{
+				Name: "file.txt",
+				Digest: &pb.Digest{
+					Hash:      fileHashStr,
+					SizeBytes: int64(len(fileData)),
+				},
+			},
+		},
+	}
+
+	tree := pb.Tree{Root: &rootDir}
+	treeData, err := proto.Marshal(&tree)
+	if err != nil {
+		t.Fatal(err)
+	}
+	treeHash := sha256.Sum256(treeData)
+	treeHashStr := hex.EncodeToString(treeHash[:])
+	if err := c.Put(ctx, cache.CAS, treeHashStr, int64(len(treeData)),
+		bytes.NewReader(treeData)); err != nil {
+		t.Fatal(err)
+	}
+
+	ar := pb.ActionResult{
+		OutputFiles: []*pb.OutputFile{
+			{
+				Path: "file.txt",
+				Digest: &pb.Digest{
+					Hash:      fileHashStr,
+					SizeBytes: int64(len(fileData)),
+				},
+			},
+		},
+		OutputDirectories: []*pb.OutputDirectory{
+			{
+				Path: "out",
+				TreeDigest: &pb.Digest{
+					Hash:      treeHashStr,
+					SizeBytes: int64(len(treeData)),
+				},
+			},
+		},
+	}
+	arData, err := proto.Marshal(&ar)
+	if err != nil {
+		t.Fatal(err)
+	}
+	arHash := sha256.Sum256([]byte("tree validation guard action"))
+	arHashStr := hex.EncodeToString(arHash[:])
+	if err := c.Put(ctx, cache.AC, arHashStr, int64(len(arData)),
+		bytes.NewReader(arData)); err != nil {
+		t.Fatal(err)
+	}
+
+	return arHashStr, int64(len(treeData))
+}
+
+func newTreeValidationCache(t *testing.T, opts ...Option) *diskCache {
+	t.Helper()
+	cacheDir := testutils.TempDir(t)
+	t.Cleanup(func() { _ = os.RemoveAll(cacheDir) })
+
+	opts = append([]Option{WithAccessLogger(testutils.NewSilentLogger())}, opts...)
+	c, err := New(cacheDir, 1024*32, opts...)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return c.(*diskCache)
+}
+
+func TestTreeValidationSizeLimitAllowsUnderCap(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	var trips []int64
+	// The limit is set after the blobs exist, via a second cache below;
+	// here declare a generous cap up-front.
+	c := newTreeValidationCache(t,
+		WithTreeValidationSizeLimit(1024*1024, func(declared int64) {
+			trips = append(trips, declared)
+		}))
+	arHash, _ := putTreeBackedActionResult(ctx, t, c)
+
+	result, data, err := c.GetValidatedActionResult(ctx, arHash)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if result == nil || data == nil {
+		t.Fatal("expected a validated hit under the cap")
+	}
+	if len(trips) != 0 {
+		t.Fatalf("expected no guard trips, got %v", trips)
+	}
+}
+
+func TestTreeValidationSizeLimitReportsMissOverCap(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// Build the blobs first with no cap, then flip the cap below the
+	// declared Tree size on the same cache to prove the guard alone flips
+	// hit to miss.
+	c := newTreeValidationCache(t)
+	arHash, treeBytes := putTreeBackedActionResult(ctx, t, c)
+
+	result, data, err := c.GetValidatedActionResult(ctx, arHash)
+	if err != nil || result == nil || data == nil {
+		t.Fatalf("expected a hit before the cap: result=%v data=%v err=%v", result, data, err)
+	}
+
+	var trips []int64
+	c.maxTreeValidationBytes = treeBytes - 1
+	c.treeValidationExceeded = func(declared int64) {
+		trips = append(trips, declared)
+	}
+
+	result, data, err = c.GetValidatedActionResult(ctx, arHash)
+	if err != nil {
+		t.Fatalf("over-cap validation must be a miss, not an error: %v", err)
+	}
+	if result != nil || data != nil {
+		t.Fatal("expected a miss over the cap")
+	}
+	if len(trips) != 1 || trips[0] != treeBytes {
+		t.Fatalf("expected one trip with declared bytes %d, got %v", treeBytes, trips)
+	}
+
+	// An exact cap admits the result again.
+	c.maxTreeValidationBytes = treeBytes
+	result, data, err = c.GetValidatedActionResult(ctx, arHash)
+	if err != nil || result == nil || data == nil {
+		t.Fatalf("expected a hit at the exact cap: result=%v data=%v err=%v", result, data, err)
+	}
+	if len(trips) != 1 {
+		t.Fatalf("expected no further trips, got %v", trips)
+	}
+}

server/grpc.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"net/http"
 
+	"golang.org/x/sync/semaphore"
 	"google.golang.org/genproto/googleapis/bytestream"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -49,6 +50,10 @@ type grpcServer struct {
 	runtimeMetrics         RuntimeMetrics
 	readLimiter            *readLimiter
 	sourceBuffers          *sourceBufferPool
+	// GetTree guards; see GetTreeLimits.
+	getTreeSem              *semaphore.Weighted
+	getTreeMaxResponseBytes int64
+	getTreeMetrics          GetTreeMetrics
 	// writePayloadConsumed, when set, is invoked once per fully-consumed
 	// ByteStream WriteRequest; see WithWritePayloadConsumed.
 	writePayloadConsumed func(*bytestream.WriteRequest)

server/grpc_cas.go

@@ -7,6 +7,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"io"
+	"math"
 
 	"google.golang.org/genproto/googleapis/rpc/code"
 	"google.golang.org/genproto/googleapis/rpc/status"
@@ -349,6 +350,21 @@ func (s *grpcServer) GetTree(in *pb.GetTreeRequest,
 		return errNilDigest
 	}
 
+	// Fail-fast concurrency guard (see GetTreeLimits): GetTree materializes
+	// the whole tree in memory, so its memory bound is this slot count times
+	// the response byte cap. No queueing - a saturated caller retries.
+	if s.getTreeSem != nil {
+		if !s.getTreeSem.TryAcquire(1) {
+			if s.getTreeMetrics != nil {
+				s.getTreeMetrics.GetTreeDenied(GetTreeDeniedSaturated)
+			}
+			s.accessLogger.Printf("%s %s DENIED: concurrency limit", errorPrefix, in.RootDigest.Hash)
+			return grpc_status.Error(codes.ResourceExhausted,
+				"too many concurrent GetTree requests, please retry")
+		}
+		defer s.getTreeSem.Release(1)
+	}
+
 	err := s.validateHash(in.RootDigest.Hash, in.RootDigest.SizeBytes, errorPrefix)
 	if err != nil {
 		return err
@@ -365,14 +381,30 @@ func (s *grpcServer) GetTree(in *pb.GetTreeRequest,
 		return grpc_status.Error(codes.Unknown, err.Error())
 	}
 
+	// Running response byte budget (see GetTreeLimits). The response size is
+	// discovered directory-by-directory during the walk, so this is a
+	// mid-traversal check on accumulated serialized bytes, not an up-front
+	// reservation. Zero (disabled) means an effectively unlimited budget.
+	budget := s.getTreeMaxResponseBytes
+	if budget <= 0 {
+		budget = math.MaxInt64
+	}
+	budget -= int64(len(data))
+	if budget < 0 {
+		return s.getTreeOverBudget(errorPrefix, in.RootDigest.Hash)
+	}
+
 	dir := pb.Directory{}
 	err = proto.Unmarshal(data, &dir)
 	if err != nil {
 		s.errorLogger.Printf("%s %s %s", errorPrefix, in.RootDigest.Hash, err)
 		return grpc_status.Error(codes.DataLoss, err.Error())
 	}
 
-	err = s.fillDirectories(stream.Context(), &resp, &dir, errorPrefix)
+	err = s.fillDirectories(stream.Context(), &resp, &dir, &budget, errorPrefix)
+	if err == errGetTreeOverBudget {
+		return s.getTreeOverBudget(errorPrefix, in.RootDigest.Hash)
+	}
 	if err != nil {
 		return err
 	}
@@ -388,9 +420,27 @@ func (s *grpcServer) GetTree(in *pb.GetTreeRequest,
 	return nil
 }
 
+// errGetTreeOverBudget aborts the traversal when the accumulated response
+// exceeds the configured byte cap. Mapped to ResourceExhausted in GetTree.
+var errGetTreeOverBudget = errors.New("GetTree response byte budget exceeded")
+
+func (s *grpcServer) getTreeOverBudget(errorPrefix, rootHash string) error {
+	if s.getTreeMetrics != nil {
+		s.getTreeMetrics.GetTreeDenied(GetTreeDeniedResponseBytes)
+	}
+	s.accessLogger.Printf("%s %s DENIED: response over %d byte limit",
+		errorPrefix, rootHash, s.getTreeMaxResponseBytes)
+	return grpc_status.Errorf(codes.ResourceExhausted,
+		"tree exceeds the server's %d byte GetTree response limit",
+		s.getTreeMaxResponseBytes)
+}
+
 // Attempt to populate `resp`. Return errors for invalid requests, but
-// otherwise attempt to return as many blobs as possible.
-func (s *grpcServer) fillDirectories(ctx context.Context, resp *pb.GetTreeResponse, dir *pb.Directory, errorPrefix string) error {
+// otherwise attempt to return as many blobs as possible. budget is the
+// remaining serialized-byte allowance for the accumulated response; it is
+// decremented as the walk discovers directories, and crossing it aborts
+// with errGetTreeOverBudget.
+func (s *grpcServer) fillDirectories(ctx context.Context, resp *pb.GetTreeResponse, dir *pb.Directory, budget *int64, errorPrefix string) error {
 
 	// Add this dir.
 	resp.Directories = append(resp.Directories, dir)
@@ -403,28 +453,38 @@ func (s *grpcServer) fillDirectories(ctx context.Context, resp *pb.GetTreeRespon
 			return err
 		}
 
+		// Check the declared size before fetching, so the guard also
+		// bounds the transient getBlobData allocation.
+		*budget -= dirNode.Digest.SizeBytes
+		if *budget < 0 {
+			return errGetTreeOverBudget
+		}
+
 		data, err := s.getBlobData(ctx, dirNode.Digest.Hash, dirNode.Digest.SizeBytes)
 		if err == errBlobNotFound {
 			s.accessLogger.Printf("GRPC GETTREEREQUEST BLOB %s NOT FOUND",
 				dirNode.Digest.Hash)
+			*budget += dirNode.Digest.SizeBytes
 			continue
 		}
 		if err != nil {
 			s.accessLogger.Printf("GRPC GETTREEREQUEST BLOB %s ERR: %v", err)
+			*budget += dirNode.Digest.SizeBytes
 			continue
 		}
 
 		dirMsg := pb.Directory{}
 		err = proto.Unmarshal(data, &dirMsg)
 		if err != nil {
 			s.accessLogger.Printf("GRPC GETTREEREQUEST BAD BLOB: %v", err)
+			*budget += dirNode.Digest.SizeBytes
 			continue
 		}
 
 		s.accessLogger.Printf("GRPC GETTREEREQUEST BLOB %s ADDED OK",
 			dirNode.Digest.Hash)
 
-		err = s.fillDirectories(ctx, resp, &dirMsg, errorPrefix)
+		err = s.fillDirectories(ctx, resp, &dirMsg, budget, errorPrefix)
 		if err != nil {
 			return err
 		}