flags: Allow configuring arbitrary gRPC headers (#3113)

Author: brancz

flags/flags.go

@@ -334,11 +334,12 @@ type FlagsRemoteStore struct {
 	RPCLoggingEnable   bool          `default:"false" help:"[deprecated] Enable gRPC logging."`
 	RPCUnaryTimeout    time.Duration `default:"5m"    help:"[deprecated] Maximum timeout window for unary gRPC requests including retries."`
 
-	GRPCMaxCallRecvMsgSize   int           `default:"33554432" help:"The maximum message size the client can receive."`
-	GRPCMaxCallSendMsgSize   int           `default:"33554432" help:"The maximum message size the client can send."`
-	GRPCStartupBackoffTime   time.Duration `default:"1m" help:"The time between failed gRPC requests during startup phase."`
-	GRPCConnectionTimeout    time.Duration `default:"3s" help:"The timeout duration for gRPC connection establishment."`
-	GRPCMaxConnectionRetries uint32        `default:"5" help:"The maximum number of retries to establish a gRPC connection."`
+	GRPCMaxCallRecvMsgSize   int               `default:"33554432" help:"The maximum message size the client can receive."`
+	GRPCMaxCallSendMsgSize   int               `default:"33554432" help:"The maximum message size the client can send."`
+	GRPCStartupBackoffTime   time.Duration     `default:"1m" help:"The time between failed gRPC requests during startup phase."`
+	GRPCConnectionTimeout    time.Duration     `default:"3s" help:"The timeout duration for gRPC connection establishment."`
+	GRPCMaxConnectionRetries uint32            `default:"5" help:"The maximum number of retries to establish a gRPC connection."`
+	GRPCHeaders              map[string]string `help:"Additional gRPC headers to send with each request (key=value pairs)."`
 }
 
 // FlagsDebuginfo contains flags to configure debuginfo.

flags/grpc.go

@@ -22,6 +22,7 @@ import (
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/encoding"
+	"google.golang.org/grpc/metadata"
 )
 
 // WaitGrpcEndpoint waits until the gRPC connection is established.
@@ -123,10 +124,10 @@ func (f FlagsRemoteStore) setupGrpcConnection(parent context.Context, metrics *g
 	}
 	propagators := propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{})
 
-	opts = append(opts,
-		grpc.WithChainUnaryInterceptor(
-			timeout.UnaryClientInterceptor(f.RPCUnaryTimeout), // 5m by default.
-			retry.UnaryClientInterceptor(
+	// Build interceptor chain
+	unaryInterceptors := []grpc.UnaryClientInterceptor{
+		timeout.UnaryClientInterceptor(f.RPCUnaryTimeout), // 5m by default.
+		retry.UnaryClientInterceptor(
 				// Back-off with Jitter: scalar: 1s, jitterFraction: 0,1, 10 runs
 				// i: 1		t:969.91774ms		total:969.91774ms
 				// i: 2		t:1.914221005s		total:2.884138745s
@@ -146,17 +147,28 @@ func (f FlagsRemoteStore) setupGrpcConnection(parent context.Context, metrics *g
 				// `WithPerRetryTimeout` allows you to shorten the deadline of each retry call, allowing you to fit multiple retries in the single parent deadline.
 				retry.WithPerRetryTimeout(2*time.Minute),
 			),
-			metrics.UnaryClientInterceptor(
-				grpc_prometheus.WithExemplarFromContext(exemplarFromContext),
-			),
-			logging.UnaryClientInterceptor(interceptorLogger(), logging.WithFieldsFromContext(logTraceID)),
+		metrics.UnaryClientInterceptor(
+			grpc_prometheus.WithExemplarFromContext(exemplarFromContext),
 		),
-		grpc.WithChainStreamInterceptor(
-			metrics.StreamClientInterceptor(
-				grpc_prometheus.WithExemplarFromContext(exemplarFromContext),
-			),
-			logging.StreamClientInterceptor(interceptorLogger(), logging.WithFieldsFromContext(logTraceID)),
+		logging.UnaryClientInterceptor(interceptorLogger(), logging.WithFieldsFromContext(logTraceID)),
+	}
+
+	streamInterceptors := []grpc.StreamClientInterceptor{
+		metrics.StreamClientInterceptor(
+			grpc_prometheus.WithExemplarFromContext(exemplarFromContext),
 		),
+		logging.StreamClientInterceptor(interceptorLogger(), logging.WithFieldsFromContext(logTraceID)),
+	}
+
+	// Add custom headers interceptor if headers are configured
+	if len(f.GRPCHeaders) > 0 {
+		unaryInterceptors = append([]grpc.UnaryClientInterceptor{customHeadersUnaryInterceptor(f.GRPCHeaders)}, unaryInterceptors...)
+		streamInterceptors = append([]grpc.StreamClientInterceptor{customHeadersStreamInterceptor(f.GRPCHeaders)}, streamInterceptors...)
+	}
+
+	opts = append(opts,
+		grpc.WithChainUnaryInterceptor(unaryInterceptors...),
+		grpc.WithChainStreamInterceptor(streamInterceptors...),
 		grpc.WithStatsHandler(tracing.NewClientHandler(
 			tracing.WithTracerProvider(tp),
 			tracing.WithPropagators(propagators),
@@ -191,6 +203,28 @@ func (t *perRequestBearerToken) RequireTransportSecurity() bool {
 	return !t.insecure
 }
 
+// customHeadersUnaryInterceptor adds custom headers to all unary RPC calls.
+func customHeadersUnaryInterceptor(headers map[string]string) grpc.UnaryClientInterceptor {
+	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+		// Add headers to outgoing context
+		for key, value := range headers {
+			ctx = metadata.AppendToOutgoingContext(ctx, key, value)
+		}
+		return invoker(ctx, method, req, reply, cc, opts...)
+	}
+}
+
+// customHeadersStreamInterceptor adds custom headers to all streaming RPC calls.
+func customHeadersStreamInterceptor(headers map[string]string) grpc.StreamClientInterceptor {
+	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+		// Add headers to outgoing context
+		for key, value := range headers {
+			ctx = metadata.AppendToOutgoingContext(ctx, key, value)
+		}
+		return streamer(ctx, desc, cc, method, opts...)
+	}
+}
+
 // interceptorLogger adapts go-kit logger to interceptor logger.
 func interceptorLogger() logging.Logger {
 	return logging.LoggerFunc(func(_ context.Context, lvl logging.Level, msg string, fields ...any) {