Use GRPC interceptors for bearer token

cmd/query/app/server.go

@@ -124,14 +124,23 @@
 
 		grpcOpts = append(grpcOpts, grpc.Creds(creds))
 	}
+
+	unaryInterceptors := []grpc.UnaryServerInterceptor{
+		bearertoken.NewUnaryServerInterceptor(),
+	}
+	streamInterceptors := []grpc.StreamServerInterceptor{
+		bearertoken.NewStreamServerInterceptor(),
+	}
 	if tm.Enabled {
-		//nolint:contextcheck
-		grpcOpts = append(grpcOpts,
-			grpc.StreamInterceptor(tenancy.NewGuardingStreamInterceptor(tm)),
-			grpc.UnaryInterceptor(tenancy.NewGuardingUnaryInterceptor(tm)),
-		)
+		unaryInterceptors = append(unaryInterceptors, tenancy.NewGuardingUnaryInterceptor(tm))
+		streamInterceptors = append(streamInterceptors, tenancy.NewGuardingStreamInterceptor(tm))
 	}
 
+	grpcOpts = append(grpcOpts,
+		grpc.ChainUnaryInterceptor(unaryInterceptors...),
+		grpc.ChainStreamInterceptor(streamInterceptors...),
+	)
+
 	server := grpc.NewServer(grpcOpts...)
 	return server, nil
 }

cmd/remote-storage/app/server.go

@@ -16,6 +16,7 @@ import (
 	"google.golang.org/grpc/reflection"
 
 	"github.com/jaegertracing/jaeger/cmd/query/app/querysvc"
+	"github.com/jaegertracing/jaeger/pkg/bearertoken"
 	"github.com/jaegertracing/jaeger/pkg/telemetery"
 	"github.com/jaegertracing/jaeger/pkg/tenancy"
 	"github.com/jaegertracing/jaeger/plugin/storage/grpc/shared"
@@ -96,13 +97,23 @@ func createGRPCServer(opts *Options, tm *tenancy.Manager, handler *shared.GRPCHa
 		creds := credentials.NewTLS(tlsCfg)
 		grpcOpts = append(grpcOpts, grpc.Creds(creds))
 	}
+
+	unaryInterceptors := []grpc.UnaryServerInterceptor{
+		bearertoken.NewUnaryServerInterceptor(),
+	}
+	streamInterceptors := []grpc.StreamServerInterceptor{
+		bearertoken.NewStreamServerInterceptor(),
+	}
 	if tm.Enabled {
-		grpcOpts = append(grpcOpts,
-			grpc.StreamInterceptor(tenancy.NewGuardingStreamInterceptor(tm)),
-			grpc.UnaryInterceptor(tenancy.NewGuardingUnaryInterceptor(tm)),
-		)
+		unaryInterceptors = append(unaryInterceptors, tenancy.NewGuardingUnaryInterceptor(tm))
+		streamInterceptors = append(streamInterceptors, tenancy.NewGuardingStreamInterceptor(tm))
 	}
 
+	grpcOpts = append(grpcOpts,
+		grpc.ChainUnaryInterceptor(unaryInterceptors...),
+		grpc.ChainStreamInterceptor(streamInterceptors...),
+	)
+
 	server := grpc.NewServer(grpcOpts...)
 	healthServer := health.NewServer()
 	reflection.Register(server)

pkg/bearertoken/grpc.go

@@ -0,0 +1,154 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package bearertoken
+
+import (
+	"context"
+	"fmt"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+const Key = "bearer.token"
+
+type tokenatedServerStream struct {
+	grpc.ServerStream
+	context context.Context
+}
+
+func (tss *tokenatedServerStream) Context() context.Context {
+	return tss.context
+}
+
+// getValidBearerToken attempts to retrieve the bearer token from the context.
+// It does not return an error if the token is missing.
+func getValidBearerToken(ctx context.Context, bearerHeader string) (string, error) {
+	bearerToken, ok := GetBearerToken(ctx)
+	if ok && bearerToken != "" {
+		return bearerToken, nil
+	}
+
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return "", nil
+	}
+
+	tokens := md.Get(bearerHeader)
+	if len(tokens) < 1 {
+		return "", nil
+	}
+	if len(tokens) > 1 {
+		return "", fmt.Errorf("malformed token: multiple tokens found")
+	}
+	return tokens[0], nil
+}
+
+// NewStreamServerInterceptor creates a new stream interceptor that injects the bearer token into the context if available.
+func NewStreamServerInterceptor() grpc.StreamServerInterceptor {
+	return streamServerInterceptor
+}
+
+func streamServerInterceptor(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+	bearerToken, err := getValidBearerToken(ss.Context(), Key)
+	if err != nil {
+		return err
+	}
+	// Upgrade the bearer token to be part of the context.
+
+	if token, _ := GetBearerToken(ss.Context()); token != "" {
+		return handler(srv, ss)
+	}
+
+	return handler(srv, &tokenatedServerStream{
+		ServerStream: ss,
+		context:      ContextWithBearerToken(ss.Context(), bearerToken),
+	})
+}
+
+// NewUnaryServerInterceptor creates a new unary interceptor that injects the bearer token into the context if available.
+func NewUnaryServerInterceptor() grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+		bearerToken, err := getValidBearerToken(ctx, Key)
+		if err != nil {
+			return nil, err
+		}
+
+		if token, _ := GetBearerToken(ctx); token != "" {
+			return handler(ctx, req)
+		}
+
+		return handler(ContextWithBearerToken(ctx, bearerToken), req)
+	}
+}
+
+// NewUnaryClientInterceptor injects the bearer token header into gRPC request metadata.
+func NewUnaryClientInterceptor() grpc.UnaryClientInterceptor {
+	return grpc.UnaryClientInterceptor(func(
+		ctx context.Context,
+		method string,
+		req, reply any,
+		cc *grpc.ClientConn,
+		invoker grpc.UnaryInvoker,
+		opts ...grpc.CallOption,
+	) error {
+		var token string
+		if md, ok := metadata.FromIncomingContext(ctx); ok {
+			tokens := md.Get(Key)
+			if len(tokens) > 1 {
+				return fmt.Errorf("malformed token: multiple tokens found")
+			}
+			if len(tokens) == 1 && tokens[0] != "" {
+				token = tokens[0]
+			}
+		}
+
+		if token == "" {
+			bearerToken, ok := GetBearerToken(ctx)
+			if ok && bearerToken != "" {
+				token = bearerToken
+			}
+		}
+
+		if token != "" {
+			ctx = metadata.AppendToOutgoingContext(ctx, Key, token)
+		}
+		return invoker(ctx, method, req, reply, cc, opts...)
+	})
+}
+
+// NewStreamClientInterceptor injects the bearer token header into gRPC request metadata.
+func NewStreamClientInterceptor() grpc.StreamClientInterceptor {
+	return grpc.StreamClientInterceptor(func(
+		ctx context.Context,
+		desc *grpc.StreamDesc,
+		cc *grpc.ClientConn,
+		method string,
+		streamer grpc.Streamer,
+		opts ...grpc.CallOption,
+	) (grpc.ClientStream, error) {
+		var token string
+		if md, ok := metadata.FromIncomingContext(ctx); ok {
+			tokens := md.Get(Key)
+			if len(tokens) > 1 {
+				return nil, fmt.Errorf("malformed token: multiple tokens found")
+			}
+			if len(tokens) == 1 && tokens[0] != "" {
+				token = tokens[0]
+			}
+		}
+
+		if token == "" {
+			bearerToken, ok := GetBearerToken(ctx)
+			if ok && bearerToken != "" {
+				token = bearerToken
+			}
+		}
+
+		if token != "" {
+			ctx = metadata.AppendToOutgoingContext(ctx, Key, token)
+		}
+		return streamer(ctx, desc, cc, method, opts...)
+	})
+}