Add generic gRPC stream forwarding (#306)

To be able to pass Bazel's build event stream (BES) to the same DNS
name, without having to add an extra L7 router in front of the
bb-storage frontend, add a configuration to forward specific gRPC
methods to other backends. No authorization is possible on the passed
through messages because Buildbarn has no knowledge about the semantics
of the forwarded messages.

The gRPC reflection service has also been extended to forward requests
that cannot be resolved locally.

Author: moroten

MODULE.bazel

@@ -57,6 +57,7 @@ use_repo(
     "com_github_gorilla_mux",
     "com_github_grpc_ecosystem_go_grpc_middleware",
     "com_github_grpc_ecosystem_go_grpc_prometheus",
+    "com_github_jhump_protoreflect_v2",
     "com_github_jmespath_go_jmespath",
     "com_github_klauspost_compress",
     "com_github_prometheus_client_golang",

go.mod

@@ -28,6 +28,7 @@ require (
 	github.com/gorilla/mux v1.8.1
 	github.com/grpc-ecosystem/go-grpc-middleware v1.4.0
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
+	github.com/jhump/protoreflect/v2 v2.0.0-beta.2
 	github.com/jmespath/go-jmespath v0.4.0
 	github.com/klauspost/compress v1.18.1
 	github.com/prometheus/client_golang v1.23.2

go.sum

@@ -198,6 +198,8 @@ github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 h1:Ovs26xHkKqVztRpIrF/92Bcuy
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgfV/d3M/q6VIi02HzZEHgUlZvzk=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.3 h1:NmZ1PKzSTQbuGHw9DGPFomqkkLWMC+vZCkfs+FHv1Vg=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.3/go.mod h1:zQrxl1YP88HQlA6i9c63DSVPFklWpGX4OWAc9bFuaH4=
+github.com/jhump/protoreflect/v2 v2.0.0-beta.2 h1:qZU+rEZUOYTz1Bnhi3xbwn+VxdXkLVeEpAeZzVXLY88=
+github.com/jhump/protoreflect/v2 v2.0.0-beta.2/go.mod h1:4tnOYkB/mq7QTyS3YKtVtNrJv4Psqout8HA1U+hZtgM=
 github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
 github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
 github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8=

internal/mock/BUILD.bazel

@@ -250,6 +250,7 @@ gomock(
         "ClientConnInterface",
         "ClientStream",
         "ServerStream",
+        "ServerTransportStream",
         "StreamHandler",
         "Streamer",
         "UnaryHandler",

pkg/grpc/BUILD.bazel

@@ -13,6 +13,7 @@ go_library(
         "client_factory.go",
         "deduplicating_client_factory.go",
         "deny_authenticator.go",
+        "forwarding_stream_handler.go",
         "jmespath_extractor.go",
         "lazy_client_dialer.go",
         "metadata_adding_interceptor.go",
@@ -26,8 +27,10 @@ go_library(
         "peer_transport_credentials_linux.go",
         "proto_trace_attributes_extractor.go",
         "proxy_dialer.go",
+        "reflection_relay.go",
         "request_headers_authenticator.go",
         "request_metadata_tracing_interceptor.go",
+        "routing_stream_handler.go",
         "server.go",
         "tls_client_certificate_authenticator.go",
     ],
@@ -49,6 +52,8 @@ go_library(
         "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
         "@com_github_grpc_ecosystem_go_grpc_middleware//:go-grpc-middleware",
         "@com_github_grpc_ecosystem_go_grpc_prometheus//:go-grpc-prometheus",
+        "@com_github_jhump_protoreflect_v2//grpcreflect",
+        "@com_github_jhump_protoreflect_v2//protoresolve",
         "@io_opentelemetry_go_contrib_instrumentation_google_golang_org_grpc_otelgrpc//:otelgrpc",
         "@io_opentelemetry_go_otel//attribute",
         "@io_opentelemetry_go_otel_trace//:trace",
@@ -63,11 +68,14 @@ go_library(
         "@org_golang_google_grpc//metadata",
         "@org_golang_google_grpc//peer",
         "@org_golang_google_grpc//reflection",
+        "@org_golang_google_grpc//reflection/grpc_reflection_v1",
         "@org_golang_google_grpc//status",
         "@org_golang_google_grpc_security_advancedtls//:advancedtls",
         "@org_golang_google_protobuf//encoding/prototext",
         "@org_golang_google_protobuf//proto",
         "@org_golang_google_protobuf//reflect/protoreflect",
+        "@org_golang_google_protobuf//types/known/emptypb",
+        "@org_golang_x_sync//errgroup",
         "@org_golang_x_sync//semaphore",
     ] + select({
         "@rules_go//go/platform:android": [
@@ -98,6 +106,7 @@ go_test(
         "authenticating_interceptor_test.go",
         "deduplicating_client_factory_test.go",
         "deny_authenticator_test.go",
+        "forwarding_stream_handler_test.go",
         "jmespath_extractor_test.go",
         "lazy_client_dialer_test.go",
         "metadata_adding_interceptor_test.go",
@@ -107,6 +116,7 @@ go_test(
         "proto_trace_attributes_extractor_test.go",
         "request_headers_authenticator_test.go",
         "request_metadata_tracing_interceptor_test.go",
+        "routing_stream_handler_test.go",
     ] + select({
         "@rules_go//go/platform:android": [
             "peer_transport_credentials_test.go",

pkg/grpc/forwarding_stream_handler.go

@@ -0,0 +1,92 @@
+package grpc
+
+import (
+	"io"
+
+	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc"
+	"google.golang.org/protobuf/types/known/emptypb"
+)
+
+// NewForwardingStreamHandler creates a grpc.StreamHandler that forwards gRPC
+// calls to a grpc.ClientConnInterface backend.
+func NewForwardingStreamHandler(client grpc.ClientConnInterface) grpc.StreamHandler {
+	forwarder := &forwardingStreamHandler{
+		backend: client,
+	}
+	return forwarder.HandleStream
+}
+
+type forwardingStreamHandler struct {
+	backend grpc.ClientConnInterface
+}
+
+// HandleStream creates a new stream to the backend. Requests from
+// incomingStream are forwarded to the backend stream and responses from the
+// backend stream are sent back in the incomingStream.
+func (s *forwardingStreamHandler) HandleStream(srv any, incomingStream grpc.ServerStream) error {
+	// All gRPC invocations has a grpc.ServerTransportStream context.
+	method, _ := grpc.Method(incomingStream.Context())
+	desc := grpc.StreamDesc{
+		// According to grpc.StreamDesc documentation, StreamName and Handler
+		// are only used when registering handlers on a server.
+		StreamName: "",
+		Handler:    nil,
+		// Streaming behaviour is wanted, single message is treated the same on
+		// transport level, the application just closes the stream after the
+		// first message.
+		ServerStreams: true,
+		ClientStreams: true,
+	}
+	group, groupCtx := errgroup.WithContext(incomingStream.Context())
+	group.Go(func() error {
+		// groupCtx is guaranteed to be canceled before returning from this method, so outgoingStream will not leak resources.
+		outgoingStream, err := s.backend.NewStream(groupCtx, &desc, method)
+		if err != nil {
+			return err
+		}
+		// Avoid group.Go because incomingStream.RecvMsg might block returning
+		// an error from the outgoingStream and getting the context for
+		// incomingStream canceled.
+		go func() {
+			for {
+				msg := &emptypb.Empty{}
+				if err := incomingStream.RecvMsg(msg); err != nil {
+					if err == io.EOF {
+						// Let's continue to receive on outgoingStream, so don't
+						// cancel grouptCtx.
+						outgoingStream.CloseSend()
+						return
+					}
+					// Cancel groupCtx immediately.
+					group.Go(func() error { return err })
+					return
+				}
+				if err := outgoingStream.SendMsg(msg); err != nil {
+					if err == io.EOF {
+						// The error will be returned by outgoingStream.RecvMsg(),
+						// no need to cancel groupCtx now.
+						return
+					}
+					// Cancel groupCtx immediately.
+					group.Go(func() error { return err })
+					return
+				}
+			}
+		}()
+
+		for {
+			msg := &emptypb.Empty{}
+			if err := outgoingStream.RecvMsg(msg); err != nil {
+				if err == io.EOF {
+					return nil
+				}
+				return err
+			}
+			if err := incomingStream.SendMsg(msg); err != nil {
+				return err
+			}
+		}
+	})
+	return group.Wait()
+}