Repository: buchgr/bazel-remote
Pull request #523 (2022-01-28 21:21:34)

Add gRPC health check service
This might be almost all that's needed for buchgr/bazel-remote#522 ?


Patch:
diff --git a/README.md b/README.md
index 24f9cc0..069446d 100644
--- a/README.md
+++ b/README.md
@@ -462,13 +462,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
+### 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
+  ```
 
-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.
 
 ### Build your own
 
diff --git a/main.go b/main.go
index b04b122..72ccf0a 100644
--- a/main.go
+++ b/main.go
@@ -181,6 +181,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 +219,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 +233,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/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/server/grpc_bytestream.go b/server/grpc_bytestream.go
index 5d73bd3..dd2490b 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()
 
diff --git a/server/grpc_test.go b/server/grpc_test.go
index eaf89fc..0ca78ad 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())
 }
@@ -2257,3 +2260,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())
+	}
+}