Repository: buchgr/bazel-remote
Pull request #526 (2022-02-15 17:06:18)

Add a max proxy blob size option
When the proxy server is only connected with a slow network link, it's sometimes beneficial not to download large blobs and recompile them locally instead.
With the new maxProxyBlobSize setting, the remote blob is not read when it is larger.


Patch:
diff --git a/README.md b/README.md
index 24f9cc0..05f0659 100644
--- a/README.md
+++ b/README.md
@@ -221,6 +221,11 @@ OPTIONS:
       preexisting blobs in the cache. (default: 9223372036854775807)
       [$BAZEL_REMOTE_MAX_BLOB_SIZE]
 
+   --max_proxy_blob_size value The maximum logical/uncompressed blob size that will
+      be downloaded from proxies. Note that this limit is not applied to
+      preexisting blobs in the cache. (default: 9223372036854775807)
+      [$BAZEL_REMOTE_MAX_PROXY_BLOB_SIZE]
+
    --num_uploaders value When using proxy backends, sets the number of
       Goroutines to process parallel uploads to backend. (default: 100)
       [$BAZEL_REMOTE_NUM_UPLOADERS]
@@ -462,13 +467,26 @@ If you want the docker container to run in the background pass the `-d` flag rig
 You can adjust the maximum cache size by appending `--max_size=N`, where N is
 the maximum size in Gibibytes.
 
-### Kubernetes note
-
-Don't name your deployment `bazel-remote`!
+### Kubernetes notes
+
+* Don't name your deployment `bazel-remote`!
+  
+  Kubernetes sets some environment variables based on this name, which conflict
+  with the `BAZEL_REMOTE_*` environment variables that bazel-remote tries to
+  parse.
+
+* bazel-remote supports the `/grpc.health.v1.Health/Check` service, which you can
+  configure like so:
+  ```
+  alb.ingress.kubernetes.io/backend-protocol: HTTP
+  alb.ingress.kubernetes.io/backend-protocol-version: GRPC
+  alb.ingress.kubernetes.io/healthcheck-path: /grpc.health.v1.Health/Check
+  alb.ingress.kubernetes.io/healthcheck-port: 9092
+  alb.ingress.kubernetes.io/listen-ports: [{"HTTPS": 9092}]
+  alb.ingress.kubernetes.io/success-codes: 0
+  alb.ingress.kubernetes.io/target-type: ip
+  ```
 
-Kubernetes sets some environment variables based on this name, which conflict
-with the `BAZEL_REMOTE_*` environment variables that bazel-remote tries to
-parse.
 
 ### Build your own
 
diff --git a/cache/disk/disk.go b/cache/disk/disk.go
index 6ad7222..d460f53 100644
--- a/cache/disk/disk.go
+++ b/cache/disk/disk.go
@@ -69,12 +69,13 @@ type lruItem struct {
 // diskCache is a filesystem-based LRU cache, with an optional backend proxy.
 // It is safe for concurrent use.
 type diskCache struct {
-	dir           string
-	proxy         cache.Proxy
-	storageMode   casblob.CompressionType
-	maxBlobSize   int64
-	accessLogger  *log.Logger
-	containsQueue chan proxyCheck
+	dir              string
+	proxy            cache.Proxy
+	storageMode      casblob.CompressionType
+	maxBlobSize      int64
+	maxProxyBlobSize int64
+	accessLogger     *log.Logger
+	containsQueue    chan proxyCheck
 
 	// Limit the number of simultaneous file removals.
 	fileRemovalSem *semaphore.Weighted
@@ -123,8 +124,9 @@ func New(dir string, maxSizeBytes int64, opts ...Option) (Cache, error) {
 		dir: dir,
 
 		// Not using config here, to avoid test import cycles.
-		storageMode: casblob.Zstandard,
-		maxBlobSize: math.MaxInt64,
+		storageMode:      casblob.Zstandard,
+		maxBlobSize:      math.MaxInt64,
+		maxProxyBlobSize: math.MaxInt64,
 
 		// Go defaults to a limit of 10,000 operating system threads.
 		// We probably don't need half of those for file removals at
@@ -815,7 +817,7 @@ func (c *diskCache) availableOrTryProxy(kind cache.EntryKind, hash string, size
 	}
 	err = nil
 
-	if c.proxy != nil {
+	if c.proxy != nil && size <= c.maxProxyBlobSize {
 		if size > 0 {
 			// If we know the size, attempt to reserve that much space.
 			if !locked {
@@ -944,6 +946,10 @@ func (c *diskCache) get(ctx context.Context, kind cache.EntryKind, hash string,
 	if r == nil {
 		return nil, -1, nil
 	}
+	if foundSize > c.maxProxyBlobSize {
+		r.Close()
+		return nil, -1, nil
+	}
 
 	if isSizeMismatch(size, foundSize) || foundSize < 0 {
 		return nil, -1, nil
@@ -1038,9 +1044,9 @@ func (c *diskCache) Contains(ctx context.Context, kind cache.EntryKind, hash str
 		return true, foundSize
 	}
 
-	if c.proxy != nil {
+	if c.proxy != nil && size <= c.maxProxyBlobSize {
 		exists, foundSize = c.proxy.Contains(ctx, kind, hash)
-		if exists && !isSizeMismatch(size, foundSize) {
+		if exists && foundSize <= c.maxProxyBlobSize && !isSizeMismatch(size, foundSize) {
 			return true, foundSize
 		}
 	}
diff --git a/cache/disk/options.go b/cache/disk/options.go
index abaa54d..bf103e0 100644
--- a/cache/disk/options.go
+++ b/cache/disk/options.go
@@ -57,6 +57,17 @@ func WithProxyBackend(proxy cache.Proxy) Option {
 	}
 }
 
+func WithProxyMaxBlobSize(maxProxyBlobSize int64) Option {
+	return func(c *CacheConfig) error {
+		if maxProxyBlobSize <= 0 {
+			return fmt.Errorf("Invalid MaxProxyBlobSize: %d", maxProxyBlobSize)
+		}
+
+		c.diskCache.maxProxyBlobSize = maxProxyBlobSize
+		return nil
+	}
+}
+
 func WithAccessLogger(logger *log.Logger) Option {
 	return func(c *CacheConfig) error {
 		c.diskCache.accessLogger = logger
diff --git a/config/config.go b/config/config.go
index 5cbfe11..55e64d0 100644
--- a/config/config.go
+++ b/config/config.go
@@ -62,6 +62,7 @@ type Config struct {
 	HTTPWriteTimeout            time.Duration             `yaml:"http_write_timeout"`
 	AccessLogLevel              string                    `yaml:"access_log_level"`
 	MaxBlobSize                 int64                     `yaml:"max_blob_size"`
+	MaxProxyBlobSize            int64                     `yaml:"max_proxy_blob_size"`
 
 	// Fields that are created by combinations of the flags above.
 	ProxyBackend cache.Proxy
@@ -112,7 +113,8 @@ func newFromArgs(dir string, maxSize int, storageMode string,
 	httpReadTimeout time.Duration,
 	httpWriteTimeout time.Duration,
 	accessLogLevel string,
-	maxBlobSize int64) (*Config, error) {
+	maxBlobSize int64,
+	maxProxyBlobSize int64) (*Config, error) {
 
 	c := Config{
 		HTTPAddress:                 httpAddress,
@@ -142,6 +144,7 @@ func newFromArgs(dir string, maxSize int, storageMode string,
 		HTTPWriteTimeout:            httpWriteTimeout,
 		AccessLogLevel:              accessLogLevel,
 		MaxBlobSize:                 maxBlobSize,
+		MaxProxyBlobSize:            maxProxyBlobSize,
 	}
 
 	err := validateConfig(&c)
@@ -176,6 +179,7 @@ func newFromYaml(data []byte) (*Config, error) {
 			NumUploaders:           100,
 			MaxQueuedUploads:       1000000,
 			MaxBlobSize:            math.MaxInt64,
+			MaxProxyBlobSize:       math.MaxInt64,
 			MetricsDurationBuckets: defaultDurationBuckets,
 			AccessLogLevel:         "all",
 		},
@@ -300,6 +304,10 @@ func validateConfig(c *Config) error {
 		return errors.New("The 'max_blob_size' flag/key must be a positive integer")
 	}
 
+	if c.MaxProxyBlobSize <= 0 {
+		return errors.New("The 'max_proxy_blob_size' flag/key must be a positive integer")
+	}
+
 	if c.GoogleCloudStorage != nil && c.HTTPBackend != nil && c.S3CloudStorage != nil {
 		return errors.New("One can specify at most one proxying backend")
 	}
@@ -457,5 +465,6 @@ func get(ctx *cli.Context) (*Config, error) {
 		ctx.Duration("http_write_timeout"),
 		ctx.String("access_log_level"),
 		ctx.Int64("max_blob_size"),
+		ctx.Int64("max_proxy_blob_size"),
 	)
 }
diff --git a/docker/push_to_dockerhub b/docker/push_to_dockerhub
index 7a156be..2f685c8 100755
--- a/docker/push_to_dockerhub
+++ b/docker/push_to_dockerhub
@@ -29,7 +29,7 @@ bazel run //:push_to_dockerhub_amd64
 bazel run //:push_to_dockerhub_arm64
 
 # Create the multiarch manifest.
-docker manifest rm buchgr/bazel-remote-cache:$tag
+docker manifest rm buchgr/bazel-remote-cache:$tag || true
 docker manifest create buchgr/bazel-remote-cache:$tag \
     --amend buchgr/bazel-remote-cache:tmp-amd64 \
     --amend buchgr/bazel-remote-cache:tmp-arm64
diff --git a/main.go b/main.go
index b04b122..2512bad 100644
--- a/main.go
+++ b/main.go
@@ -88,6 +88,7 @@ func run(ctx *cli.Context) error {
 	opts := []disk.Option{
 		disk.WithStorageMode(c.StorageMode),
 		disk.WithMaxBlobSize(c.MaxBlobSize),
+		disk.WithProxyMaxBlobSize(c.MaxProxyBlobSize),
 		disk.WithAccessLogger(c.AccessLogger),
 	}
 	if c.ProxyBackend != nil {
@@ -181,6 +182,13 @@ func run(ctx *cli.Context) error {
 
 			if c.TLSConfig != nil {
 				opts = append(opts, grpc.Creds(credentials.NewTLS(c.TLSConfig)))
+
+				if c.TLSCaFile != "" {
+					streamInterceptors = append(streamInterceptors,
+						server.GRPCmTLSStreamServerInterceptor(c.AllowUnauthenticatedReads))
+					unaryInterceptors = append(unaryInterceptors,
+						server.GRPCmTLSUnaryServerInterceptor(c.AllowUnauthenticatedReads))
+				}
 			}
 
 			if htpasswdSecrets != nil {
@@ -212,8 +220,6 @@ func run(ctx *cli.Context) error {
 			}
 			log.Println("experimental gRPC remote asset API:", remoteAssetStatus)
 
-			checkClientCertForWrites := c.AllowUnauthenticatedReads && c.TLSCaFile != ""
-
 			network := "tcp"
 			addr := c.GRPCAddress
 			if strings.HasPrefix(c.GRPCAddress, "unix://") {
@@ -228,7 +234,6 @@ func run(ctx *cli.Context) error {
 				validateAC,
 				c.EnableACKeyInstanceMangling,
 				enableRemoteAssetAPI,
-				checkClientCertForWrites,
 				diskCache, c.AccessLogger, c.ErrorLogger)
 			if err3 != nil {
 				log.Fatal(err3)
diff --git a/server/BUILD.bazel b/server/BUILD.bazel
index 94190f2..5212879 100644
--- a/server/BUILD.bazel
+++ b/server/BUILD.bazel
@@ -36,6 +36,8 @@ go_library(
         "@org_golang_google_grpc//codes:go_default_library",
         "@org_golang_google_grpc//credentials:go_default_library",
         "@org_golang_google_grpc//encoding/gzip:go_default_library",
+        "@org_golang_google_grpc//health:go_default_library",
+        "@org_golang_google_grpc//health/grpc_health_v1:go_default_library",
         "@org_golang_google_grpc//metadata:go_default_library",
         "@org_golang_google_grpc//peer:go_default_library",
         "@org_golang_google_grpc//status:go_default_library",
@@ -65,6 +67,7 @@ go_test(
         "@org_golang_google_grpc//:go_default_library",
         "@org_golang_google_grpc//codes:go_default_library",
         "@org_golang_google_grpc//credentials/insecure:go_default_library",
+        "@org_golang_google_grpc//health/grpc_health_v1:go_default_library",
         "@org_golang_google_grpc//status:go_default_library",
         "@org_golang_google_grpc//test/bufconn:go_default_library",
         "@org_golang_google_protobuf//proto:go_default_library",
diff --git a/server/grpc.go b/server/grpc.go
index 4966785..53342cf 100644
--- a/server/grpc.go
+++ b/server/grpc.go
@@ -11,6 +11,8 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 	_ "google.golang.org/grpc/encoding/gzip" // Register gzip support.
+	"google.golang.org/grpc/health"
+	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
 
@@ -31,13 +33,23 @@ const (
 	emptySha256   = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
 )
 
+const grpcHealthServiceName = "/grpc.health.v1.Health/Check"
+
 type grpcServer struct {
-	cache                    disk.Cache
-	accessLogger             cache.Logger
-	errorLogger              cache.Logger
-	depsCheck                bool
-	mangleACKeys             bool
-	checkClientCertForWrites bool
+	cache        disk.Cache
+	accessLogger cache.Logger
+	errorLogger  cache.Logger
+	depsCheck    bool
+	mangleACKeys bool
+}
+
+var readOnlyMethods = map[string]struct{}{
+	"/build.bazel.remote.execution.v2.ActionCache/GetActionResult":                {},
+	"/build.bazel.remote.execution.v2.ContentAddressableStorage/FindMissingBlobs": {},
+	"/build.bazel.remote.execution.v2.ContentAddressableStorage/BatchReadBlobs":   {},
+	"/build.bazel.remote.execution.v2.ContentAddressableStorage/GetTree":          {},
+	"/build.bazel.remote.execution.v2.Capabilities/GetCapabilities":               {},
+	"/google.bytestream.ByteStream/Read":                                          {},
 }
 
 // ListenAndServeGRPC creates a new gRPC server and listens on the given
@@ -48,7 +60,6 @@ func ListenAndServeGRPC(
 	validateACDeps bool,
 	mangleACKeys bool,
 	enableRemoteAssetAPI bool,
-	checkClientCertForWrites bool,
 	c disk.Cache, a cache.Logger, e cache.Logger) error {
 
 	listener, err := net.Listen(network, addr)
@@ -56,22 +67,20 @@ func ListenAndServeGRPC(
 		return err
 	}
 
-	return serveGRPC(listener, opts, validateACDeps, mangleACKeys, enableRemoteAssetAPI, checkClientCertForWrites, c, a, e)
+	return serveGRPC(listener, opts, validateACDeps, mangleACKeys, enableRemoteAssetAPI, c, a, e)
 }
 
 func serveGRPC(l net.Listener, opts []grpc.ServerOption,
 	validateACDepsCheck bool,
 	mangleACKeys bool,
 	enableRemoteAssetAPI bool,
-	checkClientCertForWrites bool,
 	c disk.Cache, a cache.Logger, e cache.Logger) error {
 
 	srv := grpc.NewServer(opts...)
 	s := &grpcServer{
 		cache: c, accessLogger: a, errorLogger: e,
-		depsCheck:                validateACDepsCheck,
-		mangleACKeys:             mangleACKeys,
-		checkClientCertForWrites: checkClientCertForWrites,
+		depsCheck:    validateACDepsCheck,
+		mangleACKeys: mangleACKeys,
 	}
 	pb.RegisterActionCacheServer(srv, s)
 	pb.RegisterCapabilitiesServer(srv, s)
@@ -80,6 +89,11 @@ func serveGRPC(l net.Listener, opts []grpc.ServerOption,
 	if enableRemoteAssetAPI {
 		asset.RegisterFetchServer(srv, s)
 	}
+
+	h := health.NewServer()
+	grpc_health_v1.RegisterHealthServer(srv, h)
+	h.SetServingStatus(grpcHealthServiceName, grpc_health_v1.HealthCheckResponse_SERVING)
+
 	return srv.Serve(l)
 }
 
@@ -145,6 +159,57 @@ func (s *grpcServer) validateHash(hash string, size int64, logPrefix string) err
 	return nil
 }
 
+// Return a grpc.StreamServerInterceptor that checks for mTLS/client cert
+// authentication, and optionally allows unauthenticated access to readonly
+// RPCs.
+func GRPCmTLSStreamServerInterceptor(allowUnauthenticatedReads bool) grpc.StreamServerInterceptor {
+
+	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+
+		if allowUnauthenticatedReads {
+			_, ro := readOnlyMethods[info.FullMethod]
+			if ro {
+				return handler(srv, ss)
+			}
+		}
+
+		err := checkGRPCClientCert(ss.Context())
+		if err != nil {
+			return err
+		}
+
+		return handler(srv, ss)
+	}
+}
+
+// Return a grpc.UnaryServerInterceptor that checks for mTLS/client cert
+// authentication, and optionally allows unauthenticated access to readonly
+// RPCs, and allows all clients access to the health service.
+func GRPCmTLSUnaryServerInterceptor(allowUnauthenticatedReads bool) grpc.UnaryServerInterceptor {
+
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+
+		// Always allow health service requests.
+		if info.FullMethod == grpcHealthServiceName {
+			return handler(ctx, req)
+		}
+
+		if allowUnauthenticatedReads {
+			_, ro := readOnlyMethods[info.FullMethod]
+			if ro {
+				return handler(ctx, req)
+			}
+		}
+
+		err := checkGRPCClientCert(ctx)
+		if err != nil {
+			return nil, err
+		}
+
+		return handler(ctx, req)
+	}
+}
+
 // Return a non-nil grpc error if a valid client certificate can't be
 // extracted from ctx.
 //
diff --git a/server/grpc_ac.go b/server/grpc_ac.go
index 9a0b543..266a1e4 100644
--- a/server/grpc_ac.go
+++ b/server/grpc_ac.go
@@ -204,13 +204,6 @@ func (s *grpcServer) maybeInline(ctx context.Context, inline bool, slice *[]byte
 func (s *grpcServer) UpdateActionResult(ctx context.Context,
 	req *pb.UpdateActionResultRequest) (*pb.ActionResult, error) {
 
-	if s.checkClientCertForWrites {
-		err := checkGRPCClientCert(ctx)
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	logPrefix := "GRPC AC PUT"
 	err := s.validateHash(req.ActionDigest.Hash, req.ActionDigest.SizeBytes, logPrefix)
 	if err != nil {
diff --git a/server/grpc_asset.go b/server/grpc_asset.go
index 06dfc83..600156e 100644
--- a/server/grpc_asset.go
+++ b/server/grpc_asset.go
@@ -95,15 +95,6 @@ func (s *grpcServer) FetchBlob(ctx context.Context, req *asset.FetchBlobRequest)
 
 	// Cache miss.
 
-	if s.checkClientCertForWrites {
-		err := checkGRPCClientCert(ctx)
-		if err != nil {
-			return &asset.FetchBlobResponse{
-				Status: &status.Status{Code: int32(codes.NotFound)},
-			}, nil
-		}
-	}
-
 	// See if we can download one of the URIs.
 
 	for _, uri := range req.GetUris() {
diff --git a/server/grpc_basic_auth.go b/server/grpc_basic_auth.go
index 1717914..8fa6064 100644
--- a/server/grpc_basic_auth.go
+++ b/server/grpc_basic_auth.go
@@ -21,15 +21,6 @@ var (
 		"access denied")
 )
 
-var readOnlyMethods = map[string]struct{}{
-	"/build.bazel.remote.execution.v2.ActionCache/GetActionResult":                {},
-	"/build.bazel.remote.execution.v2.ContentAddressableStorage/FindMissingBlobs": {},
-	"/build.bazel.remote.execution.v2.ContentAddressableStorage/BatchReadBlobs":   {},
-	"/build.bazel.remote.execution.v2.ContentAddressableStorage/GetTree":          {},
-	"/build.bazel.remote.execution.v2.Capabilities/GetCapabilities":               {},
-	"/google.bytestream.ByteStream/Read":                                          {},
-}
-
 // GrpcBasicAuth wraps an auth.SecretProvider, and provides gRPC interceptors
 // that verify that requests can be authenticated using HTTP basic auth.
 type GrpcBasicAuth struct {
@@ -46,10 +37,15 @@ func NewGrpcBasicAuth(secrets auth.SecretProvider, allowUnauthenticatedReadOnly
 	}
 }
 
-// StreamServerInterceptor returns a streaming server interceptor that
-// verifies that each request can be authenticated using HTTP basic auth.
+// StreamServerInterceptor verifies that each request can be authenticated
+// using HTTP basic auth, or is allowed without authentication.
 func (b *GrpcBasicAuth) StreamServerInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
 
+	// Always allow health service requests.
+	if info.FullMethod == grpcHealthServiceName {
+		return handler(srv, ss)
+	}
+
 	if b.allowUnauthenticatedReadOnly {
 		_, ro := readOnlyMethods[info.FullMethod]
 		if ro {
@@ -72,10 +68,15 @@ func (b *GrpcBasicAuth) StreamServerInterceptor(srv interface{}, ss grpc.ServerS
 	return handler(srv, ss)
 }
 
-// UnaryServerInterceptor returns a unary server interceptor that verifies
-// that each request can be authenticated using HTTP basic auth.
+// UnaryServerInterceptor verifies that each request can be authenticated
+// using HTTP basic auth, or is allowed without authenticated.
 func (b *GrpcBasicAuth) UnaryServerInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
 
+	// Always allow health service requests.
+	if info.FullMethod == grpcHealthServiceName {
+		return handler(ctx, req)
+	}
+
 	if b.allowUnauthenticatedReadOnly {
 		_, ro := readOnlyMethods[info.FullMethod]
 		if ro {
diff --git a/server/grpc_cas.go b/server/grpc_cas.go
index b65d786..534add1 100644
--- a/server/grpc_cas.go
+++ b/server/grpc_cas.go
@@ -48,13 +48,6 @@ func (s *grpcServer) FindMissingBlobs(ctx context.Context,
 func (s *grpcServer) BatchUpdateBlobs(ctx context.Context,
 	in *pb.BatchUpdateBlobsRequest) (*pb.BatchUpdateBlobsResponse, error) {
 
-	if s.checkClientCertForWrites {
-		err := checkGRPCClientCert(ctx)
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	resp := pb.BatchUpdateBlobsResponse{
 		Responses: make([]*pb.BatchUpdateBlobsResponse_Response,
 			0, len(in.Requests)),
diff --git a/utils/flags/flags.go b/utils/flags/flags.go
index 6395df9..aa183ab 100644
--- a/utils/flags/flags.go
+++ b/utils/flags/flags.go
@@ -158,6 +158,13 @@ func GetCliFlags() []cli.Flag {
 			DefaultText: strconv.FormatInt(math.MaxInt64, 10),
 			EnvVars:     []string{"BAZEL_REMOTE_MAX_BLOB_SIZE"},
 		},
+		&cli.Int64Flag{
+			Name:        "max_proxy_blob_size",
+			Value:       math.MaxInt64,
+			Usage:       "The maximum logical/uncompressed blob size that will be downloaded from proxies. Note that this limit is not applied to preexisting blobs in the cache.",
+			DefaultText: strconv.FormatInt(math.MaxInt64, 10),
+			EnvVars:     []string{"BAZEL_REMOTE_MAX_PROXY_BLOB_SIZE"},
+		},
 		&cli.IntFlag{
 			Name:    "num_uploaders",
 			Value:   100,
diff --git a/utils/grpcreadclient/BUILD.bazel b/utils/grpcreadclient/BUILD.bazel
index 9070e45..0e2d3e1 100644
--- a/utils/grpcreadclient/BUILD.bazel
+++ b/utils/grpcreadclient/BUILD.bazel
@@ -13,6 +13,7 @@ go_library(
         "@org_golang_google_grpc//codes:go_default_library",
         "@org_golang_google_grpc//credentials:go_default_library",
         "@org_golang_google_grpc//credentials/insecure:go_default_library",
+        "@org_golang_google_grpc//health/grpc_health_v1:go_default_library",
         "@org_golang_google_grpc//status:go_default_library",
     ],
 )
diff --git a/utils/grpcreadclient/grpcreadclient.go b/utils/grpcreadclient/grpcreadclient.go
index 11d37c9..98a451f 100644
--- a/utils/grpcreadclient/grpcreadclient.go
+++ b/utils/grpcreadclient/grpcreadclient.go
@@ -26,6 +26,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/status"
 )
 
@@ -444,6 +445,13 @@ func checkCacheWriteOps(conn *grpc.ClientConn, shouldWork bool) error {
 		return err
 	}
 
+	healthClient := grpc_health_v1.NewHealthClient(conn)
+
+	err = checkHealth(healthClient) // This should always work.
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -587,3 +595,24 @@ func checkBytestreamWrite(bsClient bytestream.ByteStreamClient, shouldWork bool)
 	fmt.Println("BytestreamWrite succeeded, as expected")
 	return nil
 }
+
+func checkHealth(healthClient grpc_health_v1.HealthClient) error {
+
+	const grpcHealthServiceName = "/grpc.health.v1.Health/Check"
+
+	req := grpc_health_v1.HealthCheckRequest{Service: grpcHealthServiceName}
+	resp, err := healthClient.Check(context.Background(), &req)
+	if err != nil {
+		return fmt.Errorf("Failed to check health status: %w", err)
+	}
+
+	if resp == nil {
+		return fmt.Errorf("Expected non-nil *HealthCheckResponse")
+	}
+
+	if resp.Status != grpc_health_v1.HealthCheckResponse_SERVING {
+		return fmt.Errorf("Expected health check to return SERVING status, got: %s", resp.Status.String())
+	}
+
+	return nil
+}


Test patch:
diff --git a/.bazelci/basic-auth-tests.sh b/.bazelci/basic-auth-tests.sh
index beb84ed..34e12ff 100755
--- a/.bazelci/basic-auth-tests.sh
+++ b/.bazelci/basic-auth-tests.sh
@@ -39,6 +39,8 @@ for i in $(seq 1 20)
 do
 	sleep 1
 
+	ps -p $server_pid > /dev/null || break
+
 	if wget --inet4-only -d -O - --timeout=2 \
 		--http-user "$USER" --http-password "$PASS" \
 		"http://localhost:$HTTP_PORT/status"
@@ -145,6 +147,8 @@ for i in $(seq 1 20)
 do
 	sleep 1
 
+	ps -p $server_pid > /dev/null || break
+
 	if wget --inet4-only -d -O - --timeout=2 \
 		--http-user "$USER" --http-password "$PASS" \
 		"http://localhost:$HTTP_PORT/status"
diff --git a/.bazelci/system-test.sh b/.bazelci/system-test.sh
index fe9ed7e..381ed17 100755
--- a/.bazelci/system-test.sh
+++ b/.bazelci/system-test.sh
@@ -50,6 +50,8 @@ wait_for_startup() {
 	do
 		sleep 1
 
+		ps -p $server_pid > /dev/null || break
+
 		if wget --inet4-only -d -O - "http://127.0.0.1:$HTTP_PORT/status"
 		then
 			return
diff --git a/.bazelci/tls-tests.sh b/.bazelci/tls-tests.sh
index bf475c8..e1a9026 100755
--- a/.bazelci/tls-tests.sh
+++ b/.bazelci/tls-tests.sh
@@ -94,6 +94,8 @@ for i in $(seq 1 20)
 do
 	sleep 1
 
+	ps -p $server_pid > /dev/null || break
+
 	if wget --inet4-only -d -O - --ca-certificate=$tmpdir/server.crt \
 		--certificate=$tmpdir/client.crt \
 		--private-key=$tmpdir/client.pem \
@@ -186,6 +188,8 @@ for i in $(seq 1 20)
 do
 	sleep 1
 
+	ps -p $server_pid > /dev/null || break
+
 	if wget --inet4-only -d -O - --ca-certificate=$tmpdir/server.crt \
 		--certificate=$tmpdir/client.crt \
 		--private-key=$tmpdir/client.pem \
diff --git a/cache/disk/disk_test.go b/cache/disk/disk_test.go
index 5a85810..13abe76 100644
--- a/cache/disk/disk_test.go
+++ b/cache/disk/disk_test.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"math"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -913,8 +914,24 @@ func TestHttpProxyBackend(t *testing.T) {
 		t.Fatal("Expected testCache to be empty")
 	}
 
-	// Add the proxy backend and check that we can Get the item.
+	// Add the proxy backend
 	testCache.proxy = proxy
+	testCache.maxProxyBlobSize = blobSize - 1
+	found, _ = testCache.Contains(ctx, cache.CAS, casHash, blobSize)
+	if found {
+		t.Fatalf("Expected the cache to not contain %s (via the proxy)", casHash)
+	}
+
+	r, _, err = testCache.Get(ctx, cache.CAS, casHash, blobSize, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if r != nil {
+		t.Fatal("Expected the Get to fail")
+	}
+
+	// Set a larger max proxy blob size and check that we can Get the item.
+	testCache.maxProxyBlobSize = math.MaxInt64
 
 	found, _ = testCache.Contains(ctx, cache.CAS, casHash, blobSize)
 	if !found {
diff --git a/config/config_test.go b/config/config_test.go
index c0b1df2..9758aff 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -51,6 +51,7 @@ access_log_level: none
 		NumUploaders:                100,
 		MaxQueuedUploads:            1000000,
 		MaxBlobSize:                 math.MaxInt64,
+		MaxProxyBlobSize:            math.MaxInt64,
 		MetricsDurationBuckets:      []float64{.5, 1, 2.5, 5, 10, 20, 40, 80, 160, 320},
 		AccessLogLevel:              "none",
 	}
@@ -90,6 +91,7 @@ gcs_proxy:
 		NumUploaders:           100,
 		MaxQueuedUploads:       1000000,
 		MaxBlobSize:            math.MaxInt64,
+		MaxProxyBlobSize:       math.MaxInt64,
 		MetricsDurationBuckets: []float64{.5, 1, 2.5, 5, 10, 20, 40, 80, 160, 320},
 		AccessLogLevel:         "all",
 	}
@@ -126,6 +128,7 @@ http_proxy:
 		NumUploaders:           100,
 		MaxQueuedUploads:       1000000,
 		MaxBlobSize:            math.MaxInt64,
+		MaxProxyBlobSize:       math.MaxInt64,
 		MetricsDurationBuckets: []float64{.5, 1, 2.5, 5, 10, 20, 40, 80, 160, 320},
 		AccessLogLevel:         "all",
 	}
@@ -199,6 +202,7 @@ s3_proxy:
 		NumUploaders:           100,
 		MaxQueuedUploads:       1000000,
 		MaxBlobSize:            math.MaxInt64,
+		MaxProxyBlobSize:       math.MaxInt64,
 		MetricsDurationBuckets: []float64{.5, 1, 2.5, 5, 10, 20, 40, 80, 160, 320},
 		AccessLogLevel:         "all",
 	}
@@ -229,6 +233,7 @@ profile_address: :7070
 		NumUploaders:           100,
 		MaxQueuedUploads:       1000000,
 		MaxBlobSize:            math.MaxInt64,
+		MaxProxyBlobSize:       math.MaxInt64,
 		MetricsDurationBuckets: []float64{.5, 1, 2.5, 5, 10, 20, 40, 80, 160, 320},
 		AccessLogLevel:         "all",
 	}
@@ -273,6 +278,7 @@ endpoint_metrics_duration_buckets: [.005, .1, 5]
 		NumUploaders:           100,
 		MaxQueuedUploads:       1000000,
 		MaxBlobSize:            math.MaxInt64,
+		MaxProxyBlobSize:       math.MaxInt64,
 		MetricsDurationBuckets: []float64{0.005, 0.1, 5},
 		AccessLogLevel:         "all",
 	}
@@ -287,6 +293,7 @@ func TestMetricsDurationBucketsNoDuplicates(t *testing.T) {
 		HTTPAddress:            "localhost:8080",
 		MaxSize:                42,
 		MaxBlobSize:            200,
+		MaxProxyBlobSize:       math.MaxInt64,
 		Dir:                    "/opt/cache-dir",
 		StorageMode:            "uncompressed",
 		MetricsDurationBuckets: []float64{1, 2, 3, 3},
@@ -398,6 +405,7 @@ storage_mode: zstd
 		NumUploaders:           100,
 		MaxQueuedUploads:       1000000,
 		MaxBlobSize:            math.MaxInt64,
+		MaxProxyBlobSize:       math.MaxInt64,
 		MetricsDurationBuckets: []float64{.5, 1, 2.5, 5, 10, 20, 40, 80, 160, 320},
 		AccessLogLevel:         "all",
 	}
@@ -428,6 +436,7 @@ storage_mode: zstd
 		NumUploaders:           100,
 		MaxQueuedUploads:       1000000,
 		MaxBlobSize:            math.MaxInt64,
+		MaxProxyBlobSize:       math.MaxInt64,
 		MetricsDurationBuckets: []float64{.5, 1, 2.5, 5, 10, 20, 40, 80, 160, 320},
 		AccessLogLevel:         "all",
 	}
diff --git a/server/grpc_bytestream.go b/server/grpc_bytestream.go
index 5d73bd3..f8ea6ea 100644
--- a/server/grpc_bytestream.go
+++ b/server/grpc_bytestream.go
@@ -348,13 +348,6 @@ var errDecoderPoolFail error = errors.New("failed to get DecoderWrapper from poo
 
 func (s *grpcServer) Write(srv bytestream.ByteStream_WriteServer) error {
 
-	if s.checkClientCertForWrites {
-		err := checkGRPCClientCert(srv.Context())
-		if err != nil {
-			return err
-		}
-	}
-
 	var resp bytestream.WriteResponse
 	pr, pw := io.Pipe()
 
@@ -415,7 +408,11 @@ func (s *grpcServer) Write(srv bytestream.ByteStream_WriteServer) error {
 				exists, _ := s.cache.Contains(srv.Context(), cache.CAS, hash, size)
 				if exists {
 					// Blob already exists, return without writing anything.
-					resp.CommittedSize = size
+					if cmp == casblob.Identity {
+						resp.CommittedSize = size
+					} else {
+						resp.CommittedSize = -1
+					}
 					putResult <- io.EOF
 					return
 				}
@@ -583,7 +580,21 @@ func (s *grpcServer) Write(srv bytestream.ByteStream_WriteServer) error {
 	return nil
 }
 
-func (s *grpcServer) QueryWriteStatus(context.Context, *bytestream.QueryWriteStatusRequest) (*bytestream.QueryWriteStatusResponse, error) {
-	return nil, status.Error(codes.Unimplemented,
-		"QueryWriteStatus is not implemented")
+func (s *grpcServer) QueryWriteStatus(ctx context.Context, req *bytestream.QueryWriteStatusRequest) (*bytestream.QueryWriteStatusResponse, error) {
+
+	hash, size, _, err := s.parseWriteResource(req.ResourceName)
+	if err != nil {
+		return nil, err
+	}
+
+	// We don't support partial writes, so the status will either be fully written
+	// and complete, or 0 written and incomplete.
+
+	exists, _ := s.cache.Contains(ctx, cache.CAS, hash, size)
+
+	if !exists {
+		return &bytestream.QueryWriteStatusResponse{CommittedSize: 0, Complete: false}, nil
+	}
+
+	return &bytestream.QueryWriteStatusResponse{CommittedSize: size, Complete: true}, nil
 }
diff --git a/server/grpc_test.go b/server/grpc_test.go
index eaf89fc..79b27f1 100644
--- a/server/grpc_test.go
+++ b/server/grpc_test.go
@@ -16,6 +16,8 @@ import (
 
 	asset "github.com/buchgr/bazel-remote/genproto/build/bazel/remote/asset/v1"
 	pb "github.com/buchgr/bazel-remote/genproto/build/bazel/remote/execution/v2"
+	"google.golang.org/grpc/health/grpc_health_v1"
+
 	"github.com/google/uuid"
 	"google.golang.org/protobuf/proto"
 
@@ -44,12 +46,14 @@ const bufSize = 1024 * 1024
 var (
 	listener *bufconn.Listener
 
-	acClient    pb.ActionCacheClient
-	casClient   pb.ContentAddressableStorageClient
-	bsClient    bytestream.ByteStreamClient
-	assetClient asset.FetchClient
-	ctx         = context.Background()
-	diskCache   disk.Cache
+	acClient     pb.ActionCacheClient
+	casClient    pb.ContentAddressableStorageClient
+	bsClient     bytestream.ByteStreamClient
+	assetClient  asset.FetchClient
+	healthClient grpc_health_v1.HealthClient
+
+	ctx       = context.Background()
+	diskCache disk.Cache
 
 	badDigestTestCases = []badDigest{
 		{digest: &pb.Digest{Hash: ""}, reason: "empty hash"},
@@ -91,7 +95,6 @@ func TestMain(m *testing.M) {
 	validateAC := true
 	mangleACKeys := false
 	enableRemoteAssetAPI := true
-	allowUnauthenticatedReads := false
 
 	go func() {
 		err2 := serveGRPC(
@@ -100,7 +103,6 @@ func TestMain(m *testing.M) {
 			validateAC,
 			mangleACKeys,
 			enableRemoteAssetAPI,
-			allowUnauthenticatedReads,
 			diskCache, accessLogger, errorLogger)
 		if err2 != nil {
 			fmt.Println(err2)
@@ -120,6 +122,7 @@ func TestMain(m *testing.M) {
 	acClient = pb.NewActionCacheClient(conn)
 	bsClient = bytestream.NewByteStreamClient(conn)
 	assetClient = asset.NewFetchClient(conn)
+	healthClient = grpc_health_v1.NewHealthClient(conn)
 
 	os.Exit(m.Run())
 }
@@ -975,7 +978,7 @@ func TestGrpcByteStreamZstdWrite(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if bswResp.CommittedSize != int64(len(compressedBlob)) {
+	if bswResp.CommittedSize != int64(len(compressedBlob)) && bswResp.CommittedSize != -1 {
 		t.Fatalf("Error: expected to write: %d but committed: %d\n",
 			len(testBlob), bswResp.CommittedSize)
 	}
@@ -1132,6 +1135,108 @@ func TestGrpcByteStreamSkippedWrite(t *testing.T) {
 	}
 }
 
+func TestGrpcByteStreamQueryWriteStatus(t *testing.T) {
+
+	testBlob, testBlobHash := testutils.RandomDataAndHash(123)
+	testBlobDigest := pb.Digest{
+		Hash:      testBlobHash,
+		SizeBytes: int64(len(testBlob)),
+	}
+
+	instance := "unusedInstance"
+
+	resourceName := fmt.Sprintf(
+		"%s/uploads/%s/blobs/%s/%d",
+		instance,
+		uuid.New().String(),
+		testBlobDigest.Hash,
+		len(testBlob),
+	)
+
+	req := &bytestream.QueryWriteStatusRequest{ResourceName: resourceName}
+
+	resp, err := bsClient.QueryWriteStatus(context.Background(), req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Confirm that the blob does not exist, and we get "0 bytes committed, incomplete".
+
+	if resp == nil {
+		t.Fatal("Expected non-nil *bytestream.QueryWriteStatusResponse")
+	}
+	if resp.CommittedSize != 0 {
+		t.Fatalf("Expected CommittedSize == 0, got: %d", resp.CommittedSize)
+	}
+	if resp.Complete {
+		t.Fatal("Expected incomplete response")
+	}
+
+	// Write the blob.
+	bswc, err := bsClient.Write(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = bswc.Send(&bytestream.WriteRequest{
+		ResourceName: resourceName,
+		FinishWrite:  true,
+		Data:         testBlob,
+		WriteOffset:  0,
+	})
+	if err != nil && err == io.EOF {
+		t.Fatal(err)
+	}
+
+	_, err = bswc.CloseAndRecv()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Confirm that the blob now exists, and we get "len(testBlob) bytes committed, complete".
+
+	resp, err = bsClient.QueryWriteStatus(context.Background(), req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if resp == nil {
+		t.Fatal("Expected non-nil *bytestream.QueryWriteStatusResponse")
+	}
+	if resp.CommittedSize != int64(len(testBlob)) {
+		t.Fatalf("Expected CommittedSize == %d, got: %d", len(testBlob), resp.CommittedSize)
+	}
+	if !resp.Complete {
+		t.Fatal("Expected complete response")
+	}
+
+	// Check the empty blob special case, which should always exist.
+
+	resourceName = fmt.Sprintf(
+		"%s/uploads/%s/blobs/%s/0",
+		instance,
+		uuid.New().String(),
+		emptySha256,
+	)
+
+	req = &bytestream.QueryWriteStatusRequest{ResourceName: resourceName}
+
+	resp, err = bsClient.QueryWriteStatus(context.Background(), req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if resp == nil {
+		t.Fatal("Expected non-nil *bytestream.QueryWriteStatusResponse")
+	}
+	if resp.CommittedSize != 0 {
+		t.Fatalf("Expected CommittedSize == 0, got: %d", resp.CommittedSize)
+	}
+	if !resp.Complete {
+		t.Fatal("Expected complete response")
+	}
+}
+
 func TestGrpcCasBasics(t *testing.T) {
 
 	testBlob, testBlobHash := testutils.RandomDataAndHash(256)
@@ -2257,3 +2362,19 @@ func TestCompressedBatchReadsAndWrites(t *testing.T) {
 		t.Fatal("Response data did not match")
 	}
 }
+
+func TestHealthCheck(t *testing.T) {
+	req := grpc_health_v1.HealthCheckRequest{Service: grpcHealthServiceName}
+	resp, err := healthClient.Check(ctx, &req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if resp == nil {
+		t.Fatal("Expected non-nil *HealthCheckResponse")
+	}
+
+	if resp.Status != grpc_health_v1.HealthCheckResponse_SERVING {
+		t.Fatalf("Expected health check to return SERVING status, got: %s", resp.Status.String())
+	}
+}