Add generic gRPC stream forwarding

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_lazybeaver_xorshift",

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/lazybeaver/xorshift v0.0.0-20170702203709-ce511d4823dd

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=

pkg/grpc/BUILD.bazel

@@ -26,9 +26,12 @@ 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_forwarder.go",
         "server.go",
+        "single_stream_forwarder.go",
         "tls_client_certificate_authenticator.go",
     ],
     importpath = "github.com/buildbarn/bb-storage/pkg/grpc",
@@ -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,15 @@ 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//reflection/grpc_reflection_v1alpha",
         "@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": [

pkg/grpc/reflection_relay.go

@@ -0,0 +1,83 @@
+package grpc
+
+import (
+	"context"
+	"maps"
+	"strings"
+
+	"github.com/buildbarn/bb-storage/pkg/program"
+	grpcpb "github.com/buildbarn/bb-storage/pkg/proto/configuration/grpc"
+	"github.com/buildbarn/bb-storage/pkg/util"
+	"github.com/jhump/protoreflect/v2/grpcreflect"
+	"github.com/jhump/protoreflect/v2/protoresolve"
+	v1reflectiongrpc "google.golang.org/grpc/reflection/grpc_reflection_v1"
+	v1alphareflectiongrpc "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/reflection"
+	"google.golang.org/grpc/status"
+)
+
+type combinedServiceInfoProvider struct {
+	server        reflection.ServiceInfoProvider
+	extraServices map[string]grpc.ServiceInfo
+}
+
+var _ reflection.ServiceInfoProvider = (*combinedServiceInfoProvider)(nil)
+
+// GetServiceInfo returns the currently available services, which might have
+// changed since the creation of this reflection server.
+func (p *combinedServiceInfoProvider) GetServiceInfo() map[string]grpc.ServiceInfo {
+	services := make(map[string]grpc.ServiceInfo)
+	maps.Copy(services, p.extraServices)
+	maps.Copy(services, p.server.GetServiceInfo())
+	return services
+}
+
+// registerReflection registers the google.golang.org/grpc/reflection/ service
+// on a grpc.Server and calls remote backends in case for relayed services. The
+// connections to the backend will run with the backendCtx.
+func registerReflection(backendCtx context.Context, s *grpc.Server, serverRelayConfiguration []*grpcpb.ServerRelayConfiguration, group program.Group, grpcClientFactory ClientFactory) error {
+	// Accumulate all the service names.
+	relayServices := make(map[string]grpc.ServiceInfo)
+	for _, relay := range serverRelayConfiguration {
+		for _, serviceMethod := range relay.Methods {
+			if !strings.HasPrefix(serviceMethod, "/") {
+				return status.Errorf(codes.InvalidArgument, "Malformed service method name %q, should start with '/'", serviceMethod)
+			}
+			pos := strings.LastIndex(serviceMethod, "/")
+			if pos == -1 || pos == 0 {
+				return status.Errorf(codes.InvalidArgument, "Malformed name %q, expected '/' between service and method", serviceMethod)
+			}
+			serviceName := serviceMethod[1:pos]
+			// According to ServiceInfoProvider docs for ServerOptions.Services,
+			// the reflection service is only interested in the service names.
+			relayServices[serviceName] = grpc.ServiceInfo{}
+		}
+	}
+
+	// Make a combined descriptor and extension resolver.
+	reflectionBackends := []protoresolve.Resolver{}
+	for relayIdx, relay := range serverRelayConfiguration {
+		grpcClient, err := grpcClientFactory.NewClientFromConfiguration(relay.Endpoint, group)
+		if err != nil {
+			return util.StatusWrapf(err, "Failed to create relay RPC client %d", relayIdx+1)
+		}
+		resolver := grpcreflect.NewClientAuto(backendCtx, grpcClient).AsResolver()
+		reflectionBackends = append(reflectionBackends, resolver)
+	}
+	combinedRemoteResolver := protoresolve.Combine(reflectionBackends...)
+
+	serverOptions := reflection.ServerOptions{
+		Services: &combinedServiceInfoProvider{
+			server:        s,
+			extraServices: relayServices,
+		},
+		DescriptorResolver: combinedRemoteResolver,
+		ExtensionResolver:  protoresolve.TypesFromDescriptorPool(combinedRemoteResolver),
+	}
+	v1reflectiongrpc.RegisterServerReflectionServer(s, reflection.NewServerV1(serverOptions))
+	v1alphareflectiongrpc.RegisterServerReflectionServer(s, reflection.NewServer(serverOptions))
+	return nil
+}

pkg/grpc/routing_stream_forwarder.go

@@ -0,0 +1,29 @@
+package grpc
+
+import (
+	"fmt"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+type RoutingStreamForwarder struct {
+	// A mapping from the full service and method name to forward, to the
+	// GrpcStreamHandler to be called.
+	RouteTable map[string]GrpcStreamHandler
+}
+
+func NewRoutingStreamForwarder() *RoutingStreamForwarder {
+	return &RoutingStreamForwarder{
+		RouteTable: make(map[string]GrpcStreamHandler),
+	}
+}
+
+func (s *RoutingStreamForwarder) HandleStream(method string, stream grpc.ServerStream) error {
+	if streamHandler, ok := s.RouteTable[method]; ok {
+		return streamHandler.HandleStream(method, stream)
+	}
+	errDesc := fmt.Sprintf("no route for method %v", method)
+	return status.Error(codes.Unimplemented, errDesc)
+}

pkg/grpc/server.go

@@ -9,7 +9,7 @@ import (
 	configuration "github.com/buildbarn/bb-storage/pkg/proto/configuration/grpc"
 	grpcpb "github.com/buildbarn/bb-storage/pkg/proto/configuration/grpc"
 	"github.com/buildbarn/bb-storage/pkg/util"
-	"github.com/grpc-ecosystem/go-grpc-prometheus"
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -18,7 +18,6 @@ import (
 	"google.golang.org/grpc/health"
 	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/keepalive"
-	"google.golang.org/grpc/reflection"
 	"google.golang.org/grpc/status"
 
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
@@ -147,6 +146,18 @@ func NewServersFromConfigurationAndServe(configurations []*configuration.ServerC
 			}))
 		}
 
+		if len(configuration.Relays) != 0 {
+			handler, err := newStreamRoutingFromConfiguration(configuration.Relays, grpcClientFactory, group)
+			if err != nil {
+				return util.StatusWrap(err, "Failed to create authenticator RPC client")
+			}
+			serverOptions = append(serverOptions, grpc.UnknownServiceHandler(func(srv any, stream grpc.ServerStream) error {
+				transportStream := grpc.ServerTransportStreamFromContext(stream.Context())
+				method := transportStream.Method()
+				return handler.HandleStream(method, stream)
+			}))
+		}
+
 		// Create server.
 		s := grpc.NewServer(serverOptions...)
 		stopFunc := s.Stop
@@ -162,7 +173,9 @@ func NewServersFromConfigurationAndServe(configurations []*configuration.ServerC
 
 		// Enable default services.
 		grpc_prometheus.Register(s)
-		reflection.Register(s)
+		if err := registerReflection(context.Background(), s, configuration.Relays, group, grpcClientFactory); err != nil {
+			return util.StatusWrap(err, "Failed to create reflection service")
+		}
 		h := health.NewServer()
 		grpc_health_v1.RegisterHealthServer(s, h)
 		// TODO: Construct an API for the caller to indicate
@@ -208,3 +221,22 @@ func NewServersFromConfigurationAndServe(configurations []*configuration.ServerC
 	}
 	return nil
 }
+
+func newStreamRoutingFromConfiguration(serverRelayConfiguration []*grpcpb.ServerRelayConfiguration, grpcClientFactory ClientFactory, group program.Group) (GrpcStreamHandler, error) {
+	handler := NewRoutingStreamForwarder()
+	for _, relay := range serverRelayConfiguration {
+		grpcClient, err := grpcClientFactory.NewClientFromConfiguration(relay.GetEndpoint(), group)
+		if err != nil {
+			return nil, util.StatusWrap(err, "Failed to create authenticator RPC client")
+		}
+		for _, method := range relay.GetMethods() {
+			if _, ok := handler.RouteTable[method]; ok {
+				return nil, status.Errorf(codes.InvalidArgument, "Duplicated relay for %v", method)
+			}
+			handler.RouteTable[method] = NewSingleStreamForwarder(
+				grpcClient,
+			)
+		}
+	}
+	return handler, nil
+}

pkg/grpc/single_stream_forwarder.go

@@ -0,0 +1,86 @@
+package grpc
+
+import (
+	"io"
+
+	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc"
+	"google.golang.org/protobuf/types/known/emptypb"
+)
+
+type GrpcStreamHandler interface {
+	// Process the gRPC stream with the given service and method name.
+	HandleStream(method string, stream grpc.ServerStream) error
+}
+
+func NewSingleStreamForwarder(client grpc.ClientConnInterface) GrpcStreamHandler {
+	return &singleStreamForwarder{
+		backend: client,
+	}
+}
+
+type singleStreamForwarder struct {
+	backend grpc.ClientConnInterface
+}
+
+func (s *singleStreamForwarder) HandleStream(method string, incomingStream grpc.ServerStream) error {
+	desc := grpc.StreamDesc{
+		// 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())
+	outgoingStream, err := s.backend.NewStream(groupCtx, &desc, method)
+	if err != nil {
+		return err
+	}
+	go func() {
+		for {
+			msg := &emptypb.Empty{}
+			if err := incomingStream.RecvMsg(msg); err != nil {
+				if err == io.EOF {
+					// Let's 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
+			}
+		}
+	}()
+	group.Go(func() error {
+		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
+			}
+		}
+	})
+	// group.Wait() may block a bit on incomingStream.SendMsg(), but that
+	// shouldn't be for too long.
+	return group.Wait()
+}