From a5133ccf40cbdb969944b2dc88d7a401eb7cf06b Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Wed, 27 Oct 2021 22:30:13 +1100
Subject: [PATCH 01/12] Add auth token propagation for metrics reader

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 .../metrics/prometheus/metricsstore/reader.go | 20 ++++++++++++-
 .../prometheus/metricsstore/reader_test.go    | 28 +++++++++++++++++--
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/plugin/metrics/prometheus/metricsstore/reader.go b/plugin/metrics/prometheus/metricsstore/reader.go
index 257d1d1616a..078bf2a0787 100644
--- a/plugin/metrics/prometheus/metricsstore/reader.go
+++ b/plugin/metrics/prometheus/metricsstore/reader.go
@@ -35,6 +35,7 @@ import (
 	"github.com/jaegertracing/jaeger/plugin/metrics/prometheus/metricsstore/dbmodel"
 	"github.com/jaegertracing/jaeger/proto-gen/api_v2/metrics"
 	"github.com/jaegertracing/jaeger/storage/metricsstore"
+	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
 const (
@@ -253,7 +254,7 @@ func getHTTPRoundTripper(c *config.Configuration, logger *zap.Logger) (rt http.R
 
 	// KeepAlive and TLSHandshake timeouts are kept to existing Prometheus client's
 	// DefaultRoundTripper to simplify user configuration and may be made configurable when required.
-	return &http.Transport{
+	httpTransport := &http.Transport{
 		Proxy: http.ProxyFromEnvironment,
 		DialContext: (&net.Dialer{
 			Timeout:   c.ConnectTimeout,
@@ -261,5 +262,22 @@ func getHTTPRoundTripper(c *config.Configuration, logger *zap.Logger) (rt http.R
 		}).DialContext,
 		TLSHandshakeTimeout: 10 * time.Second,
 		TLSClientConfig:     ctlsConfig,
+	}
+	return &tokenAuthTransport{
+		wrapped: httpTransport,
 	}, nil
 }
+
+// tokenAuthTransport wraps an instance of http.Transport for the purpose of
+// propagating an Authorization token from inbound to outbound HTTP requests.
+type tokenAuthTransport struct {
+	wrapped *http.Transport
+}
+
+// RoundTrip implements the http.RoundTripper interface, injecting the outbound
+// Authorization header with the token provided in the inbound request.
+func (tr *tokenAuthTransport) RoundTrip(r *http.Request) (*http.Response, error) {
+	headerToken, _ := spanstore.GetBearerToken(r.Context())
+	r.Header.Set("Authorization", "Bearer "+headerToken)
+	return tr.wrapped.RoundTrip(r)
+}
diff --git a/plugin/metrics/prometheus/metricsstore/reader_test.go b/plugin/metrics/prometheus/metricsstore/reader_test.go
index 4857c938cf5..55d578537fa 100644
--- a/plugin/metrics/prometheus/metricsstore/reader_test.go
+++ b/plugin/metrics/prometheus/metricsstore/reader_test.go
@@ -35,6 +35,7 @@ import (
 	"github.com/jaegertracing/jaeger/pkg/prometheus/config"
 	"github.com/jaegertracing/jaeger/proto-gen/api_v2/metrics"
 	"github.com/jaegertracing/jaeger/storage/metricsstore"
+	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
 type (
@@ -331,12 +332,33 @@ func TestGetRoundTripper(t *testing.T) {
 				},
 			}, logger)
 			require.NoError(t, err)
-			assert.IsType(t, &http.Transport{}, rt)
+			assert.IsType(t, &tokenAuthTransport{}, rt)
 			if tc.tlsEnabled {
-				assert.NotNil(t, rt.(*http.Transport).TLSClientConfig)
+				assert.NotNil(t, rt.(*tokenAuthTransport).wrapped.TLSClientConfig)
 			} else {
-				assert.Nil(t, rt.(*http.Transport).TLSClientConfig)
+				assert.Nil(t, rt.(*tokenAuthTransport).wrapped.TLSClientConfig)
 			}
+
+			server := httptest.NewServer(
+				http.HandlerFunc(
+					func(w http.ResponseWriter, r *http.Request) {
+						assert.Equal(t, "Bearer foo", r.Header.Get("Authorization"))
+					},
+				),
+			)
+			defer server.Close()
+
+			req, err := http.NewRequestWithContext(
+				spanstore.ContextWithBearerToken(context.Background(), "foo"),
+				http.MethodGet,
+				server.URL,
+				nil,
+			)
+			require.NoError(t, err)
+
+			resp, err := rt.RoundTrip(req)
+			require.NoError(t, err)
+			assert.Equal(t, http.StatusOK, resp.StatusCode)
 		})
 	}
 }

From fa128b3fec17529eefcc191aed1a02e0f9689535 Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Thu, 28 Oct 2021 22:06:27 +1100
Subject: [PATCH 02/12] Consolidate token propagation code

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 cmd/query/app/server.go                       |  3 +-
 cmd/query/main.go                             |  4 +-
 .../bearertoken/propagation.go                | 39 ++++++++++++++----
 .../bearertoken/propagation_test.go           | 36 ++++++++++------
 pkg/es/config/config.go                       |  4 +-
 .../metrics/prometheus/metricsstore/reader.go |  4 +-
 .../prometheus/metricsstore/reader_test.go    |  4 +-
 plugin/storage/es/options.go                  |  4 +-
 plugin/storage/grpc/shared/grpc_client.go     |  5 ++-
 .../storage/grpc/shared/grpc_client_test.go   |  5 ++-
 storage/spanstore/token_propagation.go        | 41 -------------------
 storage/spanstore/token_propagation_test.go   | 31 --------------
 12 files changed, 74 insertions(+), 106 deletions(-)
 rename cmd/query/app/token_propagation_handler.go => pkg/bearertoken/propagation.go (54%)
 rename cmd/query/app/token_propagation_hander_test.go => pkg/bearertoken/propagation_test.go (78%)
 delete mode 100644 storage/spanstore/token_propagation.go
 delete mode 100644 storage/spanstore/token_propagation_test.go

diff --git a/cmd/query/app/server.go b/cmd/query/app/server.go
index 3c410294fba..d1f1b676ebe 100644
--- a/cmd/query/app/server.go
+++ b/cmd/query/app/server.go
@@ -32,6 +32,7 @@ import (
 
 	"github.com/jaegertracing/jaeger/cmd/query/app/apiv3"
 	"github.com/jaegertracing/jaeger/cmd/query/app/querysvc"
+	"github.com/jaegertracing/jaeger/pkg/bearertoken"
 	"github.com/jaegertracing/jaeger/pkg/healthcheck"
 	"github.com/jaegertracing/jaeger/pkg/netutils"
 	"github.com/jaegertracing/jaeger/pkg/recoveryhandler"
@@ -158,7 +159,7 @@ func createHTTPServer(querySvc *querysvc.QueryService, metricsQuerySvc querysvc.
 	var handler http.Handler = r
 	handler = additionalHeadersHandler(handler, queryOpts.AdditionalHeaders)
 	if queryOpts.BearerTokenPropagation {
-		handler = bearerTokenPropagationHandler(logger, handler)
+		handler = bearertoken.PropagationHandler(logger, handler)
 	}
 	handler = handlers.CompressHandler(handler)
 	recoveryHandler := recoveryhandler.NewRecoveryHandler(logger, true)
diff --git a/cmd/query/main.go b/cmd/query/main.go
index 4659ad84bab..5ec3e708e40 100644
--- a/cmd/query/main.go
+++ b/cmd/query/main.go
@@ -35,12 +35,12 @@ import (
 	"github.com/jaegertracing/jaeger/cmd/query/app"
 	"github.com/jaegertracing/jaeger/cmd/query/app/querysvc"
 	"github.com/jaegertracing/jaeger/cmd/status"
+	"github.com/jaegertracing/jaeger/pkg/bearertoken"
 	"github.com/jaegertracing/jaeger/pkg/config"
 	"github.com/jaegertracing/jaeger/pkg/version"
 	metricsPlugin "github.com/jaegertracing/jaeger/plugin/metrics"
 	"github.com/jaegertracing/jaeger/plugin/storage"
 	"github.com/jaegertracing/jaeger/ports"
-	"github.com/jaegertracing/jaeger/storage/spanstore"
 	storageMetrics "github.com/jaegertracing/jaeger/storage/spanstore/metrics"
 )
 
@@ -95,7 +95,7 @@ func main() {
 			opentracing.SetGlobalTracer(tracer)
 			queryOpts := new(app.QueryOptions).InitFromViper(v, logger)
 			// TODO: Need to figure out set enable/disable propagation on storage plugins.
-			v.Set(spanstore.StoragePropagationKey, queryOpts.BearerTokenPropagation)
+			v.Set(bearertoken.StoragePropagationKey, queryOpts.BearerTokenPropagation)
 			storageFactory.InitFromViper(v, logger)
 			if err := storageFactory.Initialize(baseFactory, logger); err != nil {
 				logger.Fatal("Failed to init storage factory", zap.Error(err))
diff --git a/cmd/query/app/token_propagation_handler.go b/pkg/bearertoken/propagation.go
similarity index 54%
rename from cmd/query/app/token_propagation_handler.go
rename to pkg/bearertoken/propagation.go
index b44a5087e8e..da0e1c81699 100644
--- a/cmd/query/app/token_propagation_handler.go
+++ b/pkg/bearertoken/propagation.go
@@ -12,18 +12,44 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package app
+package bearertoken
 
 import (
+	"context"
 	"net/http"
 	"strings"
 
 	"go.uber.org/zap"
-
-	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
-func bearerTokenPropagationHandler(logger *zap.Logger, h http.Handler) http.Handler {
+type contextKey string
+
+// Key is the string literal used internally in the implementation of this context.
+const Key = "bearer.token"
+const bearerToken = contextKey(Key)
+
+// StoragePropagationKey is a key for viper configuration to pass this option to storage plugins.
+const StoragePropagationKey = "storage.propagate.token"
+
+// ContextWithBearerToken set bearer token in context.
+func ContextWithBearerToken(ctx context.Context, token string) context.Context {
+	if token == "" {
+		return ctx
+	}
+	return context.WithValue(ctx, bearerToken, token)
+}
+
+// GetBearerToken from context, or empty string if there is no token.
+func GetBearerToken(ctx context.Context) (string, bool) {
+	val, ok := ctx.Value(bearerToken).(string)
+	return val, ok
+}
+
+// PropagationHandler returns a http.Handler containing the logic to extract
+// the Authorization token from the http.Request and inserts it into the http.Request
+// context for easier access to the request token via GetBearerToken for bearer token
+// propagation use cases.
+func PropagationHandler(logger *zap.Logger, h http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
 		authHeaderValue := r.Header.Get("Authorization")
@@ -40,15 +66,14 @@ func bearerTokenPropagationHandler(logger *zap.Logger, h http.Handler) http.Hand
 					token = headerValue[1]
 				}
 			} else if len(headerValue) == 1 {
-				// Tread all value as a token
+				// Treat the entire value as a token.
 				token = authHeaderValue
 			} else {
 				logger.Warn("Invalid authorization header value, skipping token propagation")
 			}
-			h.ServeHTTP(w, r.WithContext(spanstore.ContextWithBearerToken(ctx, token)))
+			h.ServeHTTP(w, r.WithContext(ContextWithBearerToken(ctx, token)))
 		} else {
 			h.ServeHTTP(w, r.WithContext(ctx))
 		}
 	})
-
 }
diff --git a/cmd/query/app/token_propagation_hander_test.go b/pkg/bearertoken/propagation_test.go
similarity index 78%
rename from cmd/query/app/token_propagation_hander_test.go
rename to pkg/bearertoken/propagation_test.go
index 9d238da4924..108ec87b6b6 100644
--- a/cmd/query/app/token_propagation_hander_test.go
+++ b/pkg/bearertoken/propagation_test.go
@@ -12,41 +12,54 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package app
+package bearertoken
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/zap"
-
-	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
-func Test_bearTokenPropagationHandler(t *testing.T) {
+func Test_GetBearerToken(t *testing.T) {
+	token := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJhZG1pbiIsIm5hbWUiOiJKb2huIERvZSIsImlhdCI"
+	ctx := context.Background()
+	ctx = ContextWithBearerToken(ctx, token)
+	contextToken, ok := GetBearerToken(ctx)
+	assert.True(t, ok)
+	assert.Equal(t, contextToken, token)
+}
+
+func Test_bearerTokenPropagationHandler(t *testing.T) {
+	httpClient := &http.Client{
+		Timeout: 2 * time.Second,
+	}
+
 	logger := zap.NewNop()
 	bearerToken := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJhZG1pbiIsIm5hbWUiOiJKb2huIERvZSIsImlhdCI"
 
 	validTokenHandler := func(stop *sync.WaitGroup) http.HandlerFunc {
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		return func(w http.ResponseWriter, r *http.Request) {
 			ctx := r.Context()
-			token, ok := spanstore.GetBearerToken(ctx)
+			token, ok := GetBearerToken(ctx)
 			assert.Equal(t, token, bearerToken)
 			assert.True(t, ok)
 			stop.Done()
-		})
+		}
 	}
 
 	emptyHandler := func(stop *sync.WaitGroup) http.HandlerFunc {
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		return func(w http.ResponseWriter, r *http.Request) {
 			ctx := r.Context()
-			token, _ := spanstore.GetBearerToken(ctx)
+			token, _ := GetBearerToken(ctx)
 			assert.Empty(t, token, bearerToken)
 			stop.Done()
-		})
+		}
 	}
 
 	testCases := []struct {
@@ -68,7 +81,7 @@ func Test_bearTokenPropagationHandler(t *testing.T) {
 		t.Run(testCase.name, func(t *testing.T) {
 			stop := sync.WaitGroup{}
 			stop.Add(1)
-			r := bearerTokenPropagationHandler(logger, testCase.handler(&stop))
+			r := PropagationHandler(logger, testCase.handler(&stop))
 			server := httptest.NewServer(r)
 			defer server.Close()
 			req, err := http.NewRequest("GET", server.URL, nil)
@@ -81,5 +94,4 @@ func Test_bearTokenPropagationHandler(t *testing.T) {
 			stop.Wait()
 		})
 	}
-
 }
diff --git a/pkg/es/config/config.go b/pkg/es/config/config.go
index f9c810e6c6e..a8e662c6c44 100644
--- a/pkg/es/config/config.go
+++ b/pkg/es/config/config.go
@@ -34,10 +34,10 @@ import (
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapgrpc"
 
+	"github.com/jaegertracing/jaeger/pkg/bearertoken"
 	"github.com/jaegertracing/jaeger/pkg/config/tlscfg"
 	"github.com/jaegertracing/jaeger/pkg/es"
 	eswrapper "github.com/jaegertracing/jaeger/pkg/es/wrapper"
-	"github.com/jaegertracing/jaeger/storage/spanstore"
 	storageMetrics "github.com/jaegertracing/jaeger/storage/spanstore/metrics"
 )
 
@@ -538,7 +538,7 @@ type tokenAuthTransport struct {
 func (tr *tokenAuthTransport) RoundTrip(r *http.Request) (*http.Response, error) {
 	token := tr.token
 	if tr.allowOverrideFromCtx {
-		headerToken, _ := spanstore.GetBearerToken(r.Context())
+		headerToken, _ := bearertoken.GetBearerToken(r.Context())
 		if headerToken != "" {
 			token = headerToken
 		}
diff --git a/plugin/metrics/prometheus/metricsstore/reader.go b/plugin/metrics/prometheus/metricsstore/reader.go
index 078bf2a0787..3415957769d 100644
--- a/plugin/metrics/prometheus/metricsstore/reader.go
+++ b/plugin/metrics/prometheus/metricsstore/reader.go
@@ -31,11 +31,11 @@ import (
 	promapi "github.com/prometheus/client_golang/api/prometheus/v1"
 	"go.uber.org/zap"
 
+	"github.com/jaegertracing/jaeger/pkg/bearertoken"
 	"github.com/jaegertracing/jaeger/pkg/prometheus/config"
 	"github.com/jaegertracing/jaeger/plugin/metrics/prometheus/metricsstore/dbmodel"
 	"github.com/jaegertracing/jaeger/proto-gen/api_v2/metrics"
 	"github.com/jaegertracing/jaeger/storage/metricsstore"
-	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
 const (
@@ -277,7 +277,7 @@ type tokenAuthTransport struct {
 // RoundTrip implements the http.RoundTripper interface, injecting the outbound
 // Authorization header with the token provided in the inbound request.
 func (tr *tokenAuthTransport) RoundTrip(r *http.Request) (*http.Response, error) {
-	headerToken, _ := spanstore.GetBearerToken(r.Context())
+	headerToken, _ := bearertoken.GetBearerToken(r.Context())
 	r.Header.Set("Authorization", "Bearer "+headerToken)
 	return tr.wrapped.RoundTrip(r)
 }
diff --git a/plugin/metrics/prometheus/metricsstore/reader_test.go b/plugin/metrics/prometheus/metricsstore/reader_test.go
index 55d578537fa..4b03bcb39c6 100644
--- a/plugin/metrics/prometheus/metricsstore/reader_test.go
+++ b/plugin/metrics/prometheus/metricsstore/reader_test.go
@@ -31,11 +31,11 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
 
+	"github.com/jaegertracing/jaeger/pkg/bearertoken"
 	"github.com/jaegertracing/jaeger/pkg/config/tlscfg"
 	"github.com/jaegertracing/jaeger/pkg/prometheus/config"
 	"github.com/jaegertracing/jaeger/proto-gen/api_v2/metrics"
 	"github.com/jaegertracing/jaeger/storage/metricsstore"
-	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
 type (
@@ -349,7 +349,7 @@ func TestGetRoundTripper(t *testing.T) {
 			defer server.Close()
 
 			req, err := http.NewRequestWithContext(
-				spanstore.ContextWithBearerToken(context.Background(), "foo"),
+				bearertoken.ContextWithBearerToken(context.Background(), "foo"),
 				http.MethodGet,
 				server.URL,
 				nil,
diff --git a/plugin/storage/es/options.go b/plugin/storage/es/options.go
index 93e3a78d418..58f2aff0149 100644
--- a/plugin/storage/es/options.go
+++ b/plugin/storage/es/options.go
@@ -22,9 +22,9 @@ import (
 
 	"github.com/spf13/viper"
 
+	"github.com/jaegertracing/jaeger/pkg/bearertoken"
 	"github.com/jaegertracing/jaeger/pkg/config/tlscfg"
 	"github.com/jaegertracing/jaeger/pkg/es/config"
-	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
 const (
@@ -328,7 +328,7 @@ func initFromViper(cfg *namespaceConfig, v *viper.Viper) {
 	cfg.UseILM = v.GetBool(cfg.namespace + suffixUseILM)
 
 	// TODO: Need to figure out a better way for do this.
-	cfg.AllowTokenFromContext = v.GetBool(spanstore.StoragePropagationKey)
+	cfg.AllowTokenFromContext = v.GetBool(bearertoken.StoragePropagationKey)
 	cfg.TLS = cfg.getTLSFlagsConfig().InitFromViper(v)
 
 	remoteReadClusters := stripWhiteSpace(v.GetString(cfg.namespace + suffixRemoteReadClusters))
diff --git a/plugin/storage/grpc/shared/grpc_client.go b/plugin/storage/grpc/shared/grpc_client.go
index 4d27c0eb269..0f3d49f1880 100644
--- a/plugin/storage/grpc/shared/grpc_client.go
+++ b/plugin/storage/grpc/shared/grpc_client.go
@@ -25,6 +25,7 @@ import (
 	"google.golang.org/grpc/status"
 
 	"github.com/jaegertracing/jaeger/model"
+	"github.com/jaegertracing/jaeger/pkg/bearertoken"
 	"github.com/jaegertracing/jaeger/proto-gen/storage_v1"
 	"github.com/jaegertracing/jaeger/storage/dependencystore"
 	"github.com/jaegertracing/jaeger/storage/spanstore"
@@ -67,13 +68,13 @@ func composeContextUpgradeFuncs(funcs ...ContextUpgradeFunc) ContextUpgradeFunc
 // in the request metadata, if the original context has bearer token attached.
 // Otherwise returns original context.
 func upgradeContextWithBearerToken(ctx context.Context) context.Context {
-	bearerToken, hasToken := spanstore.GetBearerToken(ctx)
+	bearerToken, hasToken := bearertoken.GetBearerToken(ctx)
 	if hasToken {
 		md, ok := metadata.FromOutgoingContext(ctx)
 		if !ok {
 			md = metadata.New(nil)
 		}
-		md.Set(spanstore.BearerTokenKey, bearerToken)
+		md.Set(bearertoken.Key, bearerToken)
 		return metadata.NewOutgoingContext(ctx, md)
 	}
 	return ctx
diff --git a/plugin/storage/grpc/shared/grpc_client_test.go b/plugin/storage/grpc/shared/grpc_client_test.go
index b77536542ed..518f95bc7aa 100644
--- a/plugin/storage/grpc/shared/grpc_client_test.go
+++ b/plugin/storage/grpc/shared/grpc_client_test.go
@@ -28,6 +28,7 @@ import (
 	"google.golang.org/grpc/status"
 
 	"github.com/jaegertracing/jaeger/model"
+	"github.com/jaegertracing/jaeger/pkg/bearertoken"
 	"github.com/jaegertracing/jaeger/proto-gen/storage_v1"
 	grpcMocks "github.com/jaegertracing/jaeger/proto-gen/storage_v1/mocks"
 	"github.com/jaegertracing/jaeger/storage/spanstore"
@@ -108,11 +109,11 @@ func withGRPCClient(fn func(r *grpcClientTest)) {
 
 func TestContextUpgradeWithToken(t *testing.T) {
 	testBearerToken := "test-bearer-token"
-	ctx := spanstore.ContextWithBearerToken(context.Background(), testBearerToken)
+	ctx := bearertoken.ContextWithBearerToken(context.Background(), testBearerToken)
 	upgradedToken := upgradeContextWithBearerToken(ctx)
 	md, ok := metadata.FromOutgoingContext(upgradedToken)
 	assert.Truef(t, ok, "Expected metadata in context")
-	bearerTokenFromMetadata := md.Get(spanstore.BearerTokenKey)
+	bearerTokenFromMetadata := md.Get(bearertoken.Key)
 	assert.Equal(t, []string{testBearerToken}, bearerTokenFromMetadata)
 }
 
diff --git a/storage/spanstore/token_propagation.go b/storage/spanstore/token_propagation.go
deleted file mode 100644
index ba023822b2e..00000000000
--- a/storage/spanstore/token_propagation.go
+++ /dev/null
@@ -1,41 +0,0 @@
-// Copyright (c) 2019 The Jaeger Authors.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package spanstore
-
-import "context"
-
-type contextKey string
-
-// BearerTokenKey is the string literal used internally in the implementation of this context.
-const BearerTokenKey = "bearer.token"
-const bearerToken = contextKey(BearerTokenKey)
-
-// StoragePropagationKey is a key for viper configuration to pass this option to storage plugins.
-const StoragePropagationKey = "storage.propagate.token"
-
-// ContextWithBearerToken set bearer token in context
-func ContextWithBearerToken(ctx context.Context, token string) context.Context {
-	if token == "" {
-		return ctx
-	}
-	return context.WithValue(ctx, bearerToken, token)
-
-}
-
-// GetBearerToken from context, or empty string if there is no token
-func GetBearerToken(ctx context.Context) (string, bool) {
-	val, ok := ctx.Value(bearerToken).(string)
-	return val, ok
-}
diff --git a/storage/spanstore/token_propagation_test.go b/storage/spanstore/token_propagation_test.go
deleted file mode 100644
index a5fe03b5381..00000000000
--- a/storage/spanstore/token_propagation_test.go
+++ /dev/null
@@ -1,31 +0,0 @@
-// Copyright (c) 2019 The Jaeger Authors.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package spanstore
-
-import (
-	"context"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-)
-
-func Test_GetBearerToken(t *testing.T) {
-	token := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJhZG1pbiIsIm5hbWUiOiJKb2huIERvZSIsImlhdCI"
-	ctx := context.Background()
-	ctx = ContextWithBearerToken(ctx, token)
-	contextToken, ok := GetBearerToken(ctx)
-	assert.True(t, ok)
-	assert.Equal(t, contextToken, token)
-}

From d1c96adaf45f148665ee525b6055bbd0922c49c5 Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Thu, 28 Oct 2021 22:11:43 +1100
Subject: [PATCH 03/12] Rename to use PropagationHandler for consistency

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 pkg/bearertoken/propagation_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pkg/bearertoken/propagation_test.go b/pkg/bearertoken/propagation_test.go
index 108ec87b6b6..8b31e5b1588 100644
--- a/pkg/bearertoken/propagation_test.go
+++ b/pkg/bearertoken/propagation_test.go
@@ -35,7 +35,7 @@ func Test_GetBearerToken(t *testing.T) {
 	assert.Equal(t, contextToken, token)
 }
 
-func Test_bearerTokenPropagationHandler(t *testing.T) {
+func Test_PropagationHandler(t *testing.T) {
 	httpClient := &http.Client{
 		Timeout: 2 * time.Second,
 	}

From f0774b21a240a7256fcd2393ee20352ac213030b Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Fri, 29 Oct 2021 22:45:18 +1100
Subject: [PATCH 04/12] Split and move common transports

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 pkg/bearertoken/context.go                    | 26 +++++
 pkg/bearertoken/context_test.go               | 17 ++++
 pkg/bearertoken/{propagation.go => http.go}   | 24 -----
 .../{propagation_test.go => http_test.go}     | 10 --
 pkg/bearertoken/transport.go                  | 63 ++++++++++++
 pkg/bearertoken/transport_test.go             | 98 +++++++++++++++++++
 pkg/es/config/config.go                       | 29 +-----
 .../metrics/prometheus/metricsstore/reader.go | 24 ++---
 .../prometheus/metricsstore/reader_test.go    |  6 --
 9 files changed, 215 insertions(+), 82 deletions(-)
 create mode 100644 pkg/bearertoken/context.go
 create mode 100644 pkg/bearertoken/context_test.go
 rename pkg/bearertoken/{propagation.go => http.go} (72%)
 rename pkg/bearertoken/{propagation_test.go => http_test.go} (90%)
 create mode 100644 pkg/bearertoken/transport.go
 create mode 100644 pkg/bearertoken/transport_test.go

diff --git a/pkg/bearertoken/context.go b/pkg/bearertoken/context.go
new file mode 100644
index 00000000000..1b332be4ad0
--- /dev/null
+++ b/pkg/bearertoken/context.go
@@ -0,0 +1,26 @@
+package bearertoken
+
+import "context"
+
+type contextKey string
+
+// Key is the string literal used internally in the implementation of this context.
+const Key = "bearer.token"
+const bearerToken = contextKey(Key)
+
+// StoragePropagationKey is a key for viper configuration to pass this option to storage plugins.
+const StoragePropagationKey = "storage.propagate.token"
+
+// ContextWithBearerToken set bearer token in context.
+func ContextWithBearerToken(ctx context.Context, token string) context.Context {
+	if token == "" {
+		return ctx
+	}
+	return context.WithValue(ctx, bearerToken, token)
+}
+
+// GetBearerToken from context, or empty string if there is no token.
+func GetBearerToken(ctx context.Context) (string, bool) {
+	val, ok := ctx.Value(bearerToken).(string)
+	return val, ok
+}
diff --git a/pkg/bearertoken/context_test.go b/pkg/bearertoken/context_test.go
new file mode 100644
index 00000000000..fa28a59050a
--- /dev/null
+++ b/pkg/bearertoken/context_test.go
@@ -0,0 +1,17 @@
+package bearertoken
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_GetBearerToken(t *testing.T) {
+	token := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJhZG1pbiIsIm5hbWUiOiJKb2huIERvZSIsImlhdCI"
+	ctx := context.Background()
+	ctx = ContextWithBearerToken(ctx, token)
+	contextToken, ok := GetBearerToken(ctx)
+	assert.True(t, ok)
+	assert.Equal(t, contextToken, token)
+}
diff --git a/pkg/bearertoken/propagation.go b/pkg/bearertoken/http.go
similarity index 72%
rename from pkg/bearertoken/propagation.go
rename to pkg/bearertoken/http.go
index da0e1c81699..b219effef79 100644
--- a/pkg/bearertoken/propagation.go
+++ b/pkg/bearertoken/http.go
@@ -15,36 +15,12 @@
 package bearertoken
 
 import (
-	"context"
 	"net/http"
 	"strings"
 
 	"go.uber.org/zap"
 )
 
-type contextKey string
-
-// Key is the string literal used internally in the implementation of this context.
-const Key = "bearer.token"
-const bearerToken = contextKey(Key)
-
-// StoragePropagationKey is a key for viper configuration to pass this option to storage plugins.
-const StoragePropagationKey = "storage.propagate.token"
-
-// ContextWithBearerToken set bearer token in context.
-func ContextWithBearerToken(ctx context.Context, token string) context.Context {
-	if token == "" {
-		return ctx
-	}
-	return context.WithValue(ctx, bearerToken, token)
-}
-
-// GetBearerToken from context, or empty string if there is no token.
-func GetBearerToken(ctx context.Context) (string, bool) {
-	val, ok := ctx.Value(bearerToken).(string)
-	return val, ok
-}
-
 // PropagationHandler returns a http.Handler containing the logic to extract
 // the Authorization token from the http.Request and inserts it into the http.Request
 // context for easier access to the request token via GetBearerToken for bearer token
diff --git a/pkg/bearertoken/propagation_test.go b/pkg/bearertoken/http_test.go
similarity index 90%
rename from pkg/bearertoken/propagation_test.go
rename to pkg/bearertoken/http_test.go
index 8b31e5b1588..6047cf961c1 100644
--- a/pkg/bearertoken/propagation_test.go
+++ b/pkg/bearertoken/http_test.go
@@ -15,7 +15,6 @@
 package bearertoken
 
 import (
-	"context"
 	"net/http"
 	"net/http/httptest"
 	"sync"
@@ -26,15 +25,6 @@ import (
 	"go.uber.org/zap"
 )
 
-func Test_GetBearerToken(t *testing.T) {
-	token := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJhZG1pbiIsIm5hbWUiOiJKb2huIERvZSIsImlhdCI"
-	ctx := context.Background()
-	ctx = ContextWithBearerToken(ctx, token)
-	contextToken, ok := GetBearerToken(ctx)
-	assert.True(t, ok)
-	assert.Equal(t, contextToken, token)
-}
-
 func Test_PropagationHandler(t *testing.T) {
 	httpClient := &http.Client{
 		Timeout: 2 * time.Second,
diff --git a/pkg/bearertoken/transport.go b/pkg/bearertoken/transport.go
new file mode 100644
index 00000000000..c3617193e7b
--- /dev/null
+++ b/pkg/bearertoken/transport.go
@@ -0,0 +1,63 @@
+package bearertoken
+
+import (
+	"errors"
+	"net/http"
+)
+
+// transport implements the http.RoundTripper interface,
+// itself wrapping an instance of http.RoundTripper.
+type transport struct {
+	defaultToken         string
+	allowOverrideFromCtx bool
+	wrapped              http.RoundTripper
+}
+
+// Option sets attributes of this transport.
+type Option func(*transport)
+
+// WithAllowOverrideFromCtx sets whether the defaultToken can be overridden
+// with the token within the request context.
+func WithAllowOverrideFromCtx(allow bool) Option {
+	return func(t *transport) {
+		t.allowOverrideFromCtx = allow
+	}
+}
+
+// WithToken sets the defaultToken that will be injected into the outbound HTTP
+// request's Authorization bearer token header.
+// If the WithAllowOverrideFromCtx(true) option is provided, the request context's
+// bearer token, will be used in preference to this token.
+func WithToken(token string) Option {
+	return func(t *transport) {
+		t.defaultToken = token
+	}
+}
+
+// NewTransport returns a new bearer token transport that wraps the given
+// http.RoundTripper, forwarding the authorization token from inbound to
+// outbound HTTP requests.
+func NewTransport(roundTripper http.RoundTripper, opts ...Option) http.RoundTripper {
+	t := &transport{wrapped: roundTripper}
+	for _, opt := range opts {
+		opt(t)
+	}
+	return t
+}
+
+// RoundTrip injects the outbound Authorization header with the
+// token provided in the inbound request.
+func (tr *transport) RoundTrip(r *http.Request) (*http.Response, error) {
+	if tr.wrapped == nil {
+		return nil, errors.New("no http.RoundTripper provided")
+	}
+	token := tr.defaultToken
+	if tr.allowOverrideFromCtx {
+		headerToken, _ := GetBearerToken(r.Context())
+		if headerToken != "" {
+			token = headerToken
+		}
+	}
+	r.Header.Set("Authorization", "Bearer "+token)
+	return tr.wrapped.RoundTrip(r)
+}
diff --git a/pkg/bearertoken/transport_test.go b/pkg/bearertoken/transport_test.go
new file mode 100644
index 00000000000..7126c9721b9
--- /dev/null
+++ b/pkg/bearertoken/transport_test.go
@@ -0,0 +1,98 @@
+package bearertoken
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type roundTripFunc func(r *http.Request) (*http.Response, error)
+
+func (s roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
+	return s(r)
+}
+
+func TestNewTransport(t *testing.T) {
+	for _, tc := range []struct {
+		name           string
+		roundTripper   http.RoundTripper
+		requestContext context.Context
+		options        []Option
+		wantError      bool
+	}{
+		{
+			name: "No options provided and request context set should have empty Bearer token",
+			roundTripper: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+				assert.Equal(t, "Bearer ", r.Header.Get("Authorization"))
+				return &http.Response{
+					StatusCode: http.StatusOK,
+				}, nil
+			}),
+			requestContext: ContextWithBearerToken(context.Background(), "tokenFromContext"),
+		},
+		{
+			name: "Allow override from context provided, and request context set should use request context token",
+			roundTripper: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+				assert.Equal(t, "Bearer tokenFromContext", r.Header.Get("Authorization"))
+				return &http.Response{
+					StatusCode: http.StatusOK,
+				}, nil
+			}),
+			requestContext: ContextWithBearerToken(context.Background(), "tokenFromContext"),
+			options: []Option{
+				WithAllowOverrideFromCtx(true),
+			},
+		},
+		{
+			name: "Allow override from context and token provided, and request context unset should use defaultToken",
+			roundTripper: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+				assert.Equal(t, "Bearer initToken", r.Header.Get("Authorization"))
+				return &http.Response{}, nil
+			}),
+			requestContext: context.Background(),
+			options: []Option{
+				WithAllowOverrideFromCtx(true),
+				WithToken("initToken"),
+			},
+		},
+		{
+			name: "Allow override from context and token provided, and request context set should use context token",
+			roundTripper: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+				assert.Equal(t, "Bearer tokenFromContext", r.Header.Get("Authorization"))
+				return &http.Response{}, nil
+			}),
+			requestContext: ContextWithBearerToken(context.Background(), "tokenFromContext"),
+			options: []Option{
+				WithAllowOverrideFromCtx(true),
+				WithToken("initToken"),
+			},
+		},
+		{
+			name:           "Nil roundTripper provided should return an error",
+			requestContext: context.Background(),
+			wantError:      true,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			server := httptest.NewServer(nil)
+			defer server.Close()
+			req, err := http.NewRequestWithContext(tc.requestContext, "GET", server.URL, nil)
+			require.NoError(t, err)
+
+			tr := NewTransport(tc.roundTripper, tc.options...)
+			resp, err := tr.RoundTrip(req)
+
+			if tc.wantError {
+				assert.Nil(t, resp)
+				assert.Error(t, err)
+			} else {
+				assert.NotNil(t, resp)
+				assert.NoError(t, err)
+			}
+		})
+	}
+}
diff --git a/pkg/es/config/config.go b/pkg/es/config/config.go
index a8e662c6c44..09ed57439bc 100644
--- a/pkg/es/config/config.go
+++ b/pkg/es/config/config.go
@@ -519,34 +519,15 @@ func GetHTTPRoundTripper(c *Configuration, logger *zap.Logger) (http.RoundTrippe
 		token = tokenFromFile
 	}
 	if token != "" || c.AllowTokenFromContext {
-		transport = &tokenAuthTransport{
-			token:                token,
-			allowOverrideFromCtx: c.AllowTokenFromContext,
-			wrapped:              httpTransport,
-		}
+		transport = bearertoken.NewTransport(
+			httpTransport,
+			bearertoken.WithAllowOverrideFromCtx(c.AllowTokenFromContext),
+			bearertoken.WithToken(token),
+		)
 	}
 	return transport, nil
 }
 
-// TokenAuthTransport
-type tokenAuthTransport struct {
-	token                string
-	allowOverrideFromCtx bool
-	wrapped              *http.Transport
-}
-
-func (tr *tokenAuthTransport) RoundTrip(r *http.Request) (*http.Response, error) {
-	token := tr.token
-	if tr.allowOverrideFromCtx {
-		headerToken, _ := bearertoken.GetBearerToken(r.Context())
-		if headerToken != "" {
-			token = headerToken
-		}
-	}
-	r.Header.Set("Authorization", "Bearer "+token)
-	return tr.wrapped.RoundTrip(r)
-}
-
 func loadToken(path string) (string, error) {
 	b, err := os.ReadFile(filepath.Clean(path))
 	if err != nil {
diff --git a/plugin/metrics/prometheus/metricsstore/reader.go b/plugin/metrics/prometheus/metricsstore/reader.go
index 3415957769d..4908ca0d337 100644
--- a/plugin/metrics/prometheus/metricsstore/reader.go
+++ b/plugin/metrics/prometheus/metricsstore/reader.go
@@ -24,6 +24,8 @@ import (
 	"time"
 	"unicode"
 
+	"github.com/jaegertracing/jaeger/pkg/bearertoken"
+
 	"github.com/opentracing/opentracing-go"
 	ottag "github.com/opentracing/opentracing-go/ext"
 	otlog "github.com/opentracing/opentracing-go/log"
@@ -31,7 +33,6 @@ import (
 	promapi "github.com/prometheus/client_golang/api/prometheus/v1"
 	"go.uber.org/zap"
 
-	"github.com/jaegertracing/jaeger/pkg/bearertoken"
 	"github.com/jaegertracing/jaeger/pkg/prometheus/config"
 	"github.com/jaegertracing/jaeger/plugin/metrics/prometheus/metricsstore/dbmodel"
 	"github.com/jaegertracing/jaeger/proto-gen/api_v2/metrics"
@@ -263,21 +264,8 @@ func getHTTPRoundTripper(c *config.Configuration, logger *zap.Logger) (rt http.R
 		TLSHandshakeTimeout: 10 * time.Second,
 		TLSClientConfig:     ctlsConfig,
 	}
-	return &tokenAuthTransport{
-		wrapped: httpTransport,
-	}, nil
-}
-
-// tokenAuthTransport wraps an instance of http.Transport for the purpose of
-// propagating an Authorization token from inbound to outbound HTTP requests.
-type tokenAuthTransport struct {
-	wrapped *http.Transport
-}
-
-// RoundTrip implements the http.RoundTripper interface, injecting the outbound
-// Authorization header with the token provided in the inbound request.
-func (tr *tokenAuthTransport) RoundTrip(r *http.Request) (*http.Response, error) {
-	headerToken, _ := bearertoken.GetBearerToken(r.Context())
-	r.Header.Set("Authorization", "Bearer "+headerToken)
-	return tr.wrapped.RoundTrip(r)
+	return bearertoken.NewTransport(
+		httpTransport,
+		bearertoken.WithAllowOverrideFromCtx(true),
+	), nil
 }
diff --git a/plugin/metrics/prometheus/metricsstore/reader_test.go b/plugin/metrics/prometheus/metricsstore/reader_test.go
index 4b03bcb39c6..3c25bb52516 100644
--- a/plugin/metrics/prometheus/metricsstore/reader_test.go
+++ b/plugin/metrics/prometheus/metricsstore/reader_test.go
@@ -332,12 +332,6 @@ func TestGetRoundTripper(t *testing.T) {
 				},
 			}, logger)
 			require.NoError(t, err)
-			assert.IsType(t, &tokenAuthTransport{}, rt)
-			if tc.tlsEnabled {
-				assert.NotNil(t, rt.(*tokenAuthTransport).wrapped.TLSClientConfig)
-			} else {
-				assert.Nil(t, rt.(*tokenAuthTransport).wrapped.TLSClientConfig)
-			}
 
 			server := httptest.NewServer(
 				http.HandlerFunc(

From 8e5f1b934aef14b44358aaa618aebb991e86dadc Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Fri, 29 Oct 2021 22:46:18 +1100
Subject: [PATCH 05/12] Give bearerToken value less meaning

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 pkg/bearertoken/http_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pkg/bearertoken/http_test.go b/pkg/bearertoken/http_test.go
index 6047cf961c1..88d2d4618b7 100644
--- a/pkg/bearertoken/http_test.go
+++ b/pkg/bearertoken/http_test.go
@@ -31,7 +31,7 @@ func Test_PropagationHandler(t *testing.T) {
 	}
 
 	logger := zap.NewNop()
-	bearerToken := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJhZG1pbiIsIm5hbWUiOiJKb2huIERvZSIsImlhdCI"
+	const bearerToken = "blah"
 
 	validTokenHandler := func(stop *sync.WaitGroup) http.HandlerFunc {
 		return func(w http.ResponseWriter, r *http.Request) {

From 73db36159830b06e890dd7dfc7f481c7d57222a0 Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Fri, 29 Oct 2021 22:52:02 +1100
Subject: [PATCH 06/12] Fix import grouping

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 plugin/metrics/prometheus/metricsstore/reader.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/plugin/metrics/prometheus/metricsstore/reader.go b/plugin/metrics/prometheus/metricsstore/reader.go
index 4908ca0d337..ac9b4eed5a1 100644
--- a/plugin/metrics/prometheus/metricsstore/reader.go
+++ b/plugin/metrics/prometheus/metricsstore/reader.go
@@ -24,8 +24,6 @@ import (
 	"time"
 	"unicode"
 
-	"github.com/jaegertracing/jaeger/pkg/bearertoken"
-
 	"github.com/opentracing/opentracing-go"
 	ottag "github.com/opentracing/opentracing-go/ext"
 	otlog "github.com/opentracing/opentracing-go/log"
@@ -33,6 +31,7 @@ import (
 	promapi "github.com/prometheus/client_golang/api/prometheus/v1"
 	"go.uber.org/zap"
 
+	"github.com/jaegertracing/jaeger/pkg/bearertoken"
 	"github.com/jaegertracing/jaeger/pkg/prometheus/config"
 	"github.com/jaegertracing/jaeger/plugin/metrics/prometheus/metricsstore/dbmodel"
 	"github.com/jaegertracing/jaeger/proto-gen/api_v2/metrics"

From 787b07bfe6ca7208d3bb0dc5743b7a110b291910 Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Fri, 29 Oct 2021 22:53:32 +1100
Subject: [PATCH 07/12] Give token value less meaning

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 pkg/bearertoken/context_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pkg/bearertoken/context_test.go b/pkg/bearertoken/context_test.go
index fa28a59050a..5f37d4280f9 100644
--- a/pkg/bearertoken/context_test.go
+++ b/pkg/bearertoken/context_test.go
@@ -8,7 +8,7 @@ import (
 )
 
 func Test_GetBearerToken(t *testing.T) {
-	token := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJhZG1pbiIsIm5hbWUiOiJKb2huIERvZSIsImlhdCI"
+	const token = "blah"
 	ctx := context.Background()
 	ctx = ContextWithBearerToken(ctx, token)
 	contextToken, ok := GetBearerToken(ctx)

From 9af8681a359fb6792c81585ba8b41a03d4673883 Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Fri, 29 Oct 2021 23:09:36 +1100
Subject: [PATCH 08/12] make fmt

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 pkg/bearertoken/context.go        | 14 ++++++++++++++
 pkg/bearertoken/context_test.go   | 14 ++++++++++++++
 pkg/bearertoken/transport.go      | 14 ++++++++++++++
 pkg/bearertoken/transport_test.go | 14 ++++++++++++++
 4 files changed, 56 insertions(+)

diff --git a/pkg/bearertoken/context.go b/pkg/bearertoken/context.go
index 1b332be4ad0..44a152a0cbe 100644
--- a/pkg/bearertoken/context.go
+++ b/pkg/bearertoken/context.go
@@ -1,3 +1,17 @@
+// Copyright (c) 2021 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package bearertoken
 
 import "context"
diff --git a/pkg/bearertoken/context_test.go b/pkg/bearertoken/context_test.go
index 5f37d4280f9..7a3f5184f20 100644
--- a/pkg/bearertoken/context_test.go
+++ b/pkg/bearertoken/context_test.go
@@ -1,3 +1,17 @@
+// Copyright (c) 2021 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package bearertoken
 
 import (
diff --git a/pkg/bearertoken/transport.go b/pkg/bearertoken/transport.go
index c3617193e7b..682c9340612 100644
--- a/pkg/bearertoken/transport.go
+++ b/pkg/bearertoken/transport.go
@@ -1,3 +1,17 @@
+// Copyright (c) 2021 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package bearertoken
 
 import (
diff --git a/pkg/bearertoken/transport_test.go b/pkg/bearertoken/transport_test.go
index 7126c9721b9..32f08848521 100644
--- a/pkg/bearertoken/transport_test.go
+++ b/pkg/bearertoken/transport_test.go
@@ -1,3 +1,17 @@
+// Copyright (c) 2021 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package bearertoken
 
 import (

From 56ac31bc17243d3d81ecfe6a495a4ba137d4fef8 Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Sat, 30 Oct 2021 07:05:25 +1100
Subject: [PATCH 09/12] Remove functional options

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 pkg/bearertoken/context.go                    | 16 ++---
 pkg/bearertoken/transport.go                  | 59 ++++++-------------
 pkg/bearertoken/transport_test.go             | 53 ++++++++---------
 .../metrics/prometheus/metricsstore/reader.go |  8 +--
 4 files changed, 55 insertions(+), 81 deletions(-)

diff --git a/pkg/bearertoken/context.go b/pkg/bearertoken/context.go
index 44a152a0cbe..ebbfdf7eb8c 100644
--- a/pkg/bearertoken/context.go
+++ b/pkg/bearertoken/context.go
@@ -16,25 +16,25 @@ package bearertoken
 
 import "context"
 
-type contextKey string
+type contextKeyType string
 
-// Key is the string literal used internally in the implementation of this context.
-const Key = "bearer.token"
-const bearerToken = contextKey(Key)
+const (
+	// StoragePropagationKey is a key for viper configuration to pass this option to storage plugins.
+	StoragePropagationKey = "storage.propagate.token"
 
-// StoragePropagationKey is a key for viper configuration to pass this option to storage plugins.
-const StoragePropagationKey = "storage.propagate.token"
+	contextKey = contextKeyType("bearer.token")
+)
 
 // ContextWithBearerToken set bearer token in context.
 func ContextWithBearerToken(ctx context.Context, token string) context.Context {
 	if token == "" {
 		return ctx
 	}
-	return context.WithValue(ctx, bearerToken, token)
+	return context.WithValue(ctx, contextKey, token)
 }
 
 // GetBearerToken from context, or empty string if there is no token.
 func GetBearerToken(ctx context.Context) (string, bool) {
-	val, ok := ctx.Value(bearerToken).(string)
+	val, ok := ctx.Value(contextKey).(string)
 	return val, ok
 }
diff --git a/pkg/bearertoken/transport.go b/pkg/bearertoken/transport.go
index 682c9340612..7e13e6300b5 100644
--- a/pkg/bearertoken/transport.go
+++ b/pkg/bearertoken/transport.go
@@ -19,59 +19,34 @@ import (
 	"net/http"
 )
 
-// transport implements the http.RoundTripper interface,
-// itself wrapping an instance of http.RoundTripper.
-type transport struct {
-	defaultToken         string
-	allowOverrideFromCtx bool
-	wrapped              http.RoundTripper
-}
+// RoundTripper wraps another http.RoundTripper and injects
+// an authentication header with bearer token into requests.
+type RoundTripper struct {
+	// Transport is the underlying http.RoundTripper being wrapped. Required.
+	Transport http.RoundTripper
 
-// Option sets attributes of this transport.
-type Option func(*transport)
+	// StaticToken is the pre-configured bearer token. Optional.
+	StaticToken string
 
-// WithAllowOverrideFromCtx sets whether the defaultToken can be overridden
-// with the token within the request context.
-func WithAllowOverrideFromCtx(allow bool) Option {
-	return func(t *transport) {
-		t.allowOverrideFromCtx = allow
-	}
-}
-
-// WithToken sets the defaultToken that will be injected into the outbound HTTP
-// request's Authorization bearer token header.
-// If the WithAllowOverrideFromCtx(true) option is provided, the request context's
-// bearer token, will be used in preference to this token.
-func WithToken(token string) Option {
-	return func(t *transport) {
-		t.defaultToken = token
-	}
-}
-
-// NewTransport returns a new bearer token transport that wraps the given
-// http.RoundTripper, forwarding the authorization token from inbound to
-// outbound HTTP requests.
-func NewTransport(roundTripper http.RoundTripper, opts ...Option) http.RoundTripper {
-	t := &transport{wrapped: roundTripper}
-	for _, opt := range opts {
-		opt(t)
-	}
-	return t
+	// OverrideFromCtx enables reading bearer token from Context.
+	OverrideFromCtx bool
 }
 
 // RoundTrip injects the outbound Authorization header with the
 // token provided in the inbound request.
-func (tr *transport) RoundTrip(r *http.Request) (*http.Response, error) {
-	if tr.wrapped == nil {
+func (tr RoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	if tr.Transport == nil {
 		return nil, errors.New("no http.RoundTripper provided")
 	}
-	token := tr.defaultToken
-	if tr.allowOverrideFromCtx {
+	token := tr.StaticToken
+	if tr.OverrideFromCtx {
 		headerToken, _ := GetBearerToken(r.Context())
 		if headerToken != "" {
 			token = headerToken
 		}
 	}
-	r.Header.Set("Authorization", "Bearer "+token)
-	return tr.wrapped.RoundTrip(r)
+	if token != "" {
+		r.Header.Set("Authorization", "Bearer "+token)
+	}
+	return tr.Transport.RoundTrip(r)
 }
diff --git a/pkg/bearertoken/transport_test.go b/pkg/bearertoken/transport_test.go
index 32f08848521..05992c5b97b 100644
--- a/pkg/bearertoken/transport_test.go
+++ b/pkg/bearertoken/transport_test.go
@@ -30,18 +30,19 @@ func (s roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
 	return s(r)
 }
 
-func TestNewTransport(t *testing.T) {
+func TestRoundTripper(t *testing.T) {
 	for _, tc := range []struct {
-		name           string
-		roundTripper   http.RoundTripper
-		requestContext context.Context
-		options        []Option
-		wantError      bool
+		name             string
+		staticToken      string
+		overrideFromCtx  bool
+		wrappedTransport http.RoundTripper
+		requestContext   context.Context
+		wantError        bool
 	}{
 		{
-			name: "No options provided and request context set should have empty Bearer token",
-			roundTripper: roundTripFunc(func(r *http.Request) (*http.Response, error) {
-				assert.Equal(t, "Bearer ", r.Header.Get("Authorization"))
+			name: "Default RoundTripper and request context set should have empty Bearer token",
+			wrappedTransport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+				assert.Empty(t, r.Header.Get("Authorization"))
 				return &http.Response{
 					StatusCode: http.StatusOK,
 				}, nil
@@ -49,41 +50,35 @@ func TestNewTransport(t *testing.T) {
 			requestContext: ContextWithBearerToken(context.Background(), "tokenFromContext"),
 		},
 		{
-			name: "Allow override from context provided, and request context set should use request context token",
-			roundTripper: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+			name:            "Override from context provided, and request context set should use request context token",
+			overrideFromCtx: true,
+			wrappedTransport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
 				assert.Equal(t, "Bearer tokenFromContext", r.Header.Get("Authorization"))
 				return &http.Response{
 					StatusCode: http.StatusOK,
 				}, nil
 			}),
 			requestContext: ContextWithBearerToken(context.Background(), "tokenFromContext"),
-			options: []Option{
-				WithAllowOverrideFromCtx(true),
-			},
 		},
 		{
-			name: "Allow override from context and token provided, and request context unset should use defaultToken",
-			roundTripper: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+			name:            "Allow override from context and token provided, and request context unset should use defaultToken",
+			overrideFromCtx: true,
+			staticToken:     "initToken",
+			wrappedTransport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
 				assert.Equal(t, "Bearer initToken", r.Header.Get("Authorization"))
 				return &http.Response{}, nil
 			}),
 			requestContext: context.Background(),
-			options: []Option{
-				WithAllowOverrideFromCtx(true),
-				WithToken("initToken"),
-			},
 		},
 		{
-			name: "Allow override from context and token provided, and request context set should use context token",
-			roundTripper: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+			name:            "Allow override from context and token provided, and request context set should use context token",
+			overrideFromCtx: true,
+			staticToken:     "initToken",
+			wrappedTransport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
 				assert.Equal(t, "Bearer tokenFromContext", r.Header.Get("Authorization"))
 				return &http.Response{}, nil
 			}),
 			requestContext: ContextWithBearerToken(context.Background(), "tokenFromContext"),
-			options: []Option{
-				WithAllowOverrideFromCtx(true),
-				WithToken("initToken"),
-			},
 		},
 		{
 			name:           "Nil roundTripper provided should return an error",
@@ -97,7 +92,11 @@ func TestNewTransport(t *testing.T) {
 			req, err := http.NewRequestWithContext(tc.requestContext, "GET", server.URL, nil)
 			require.NoError(t, err)
 
-			tr := NewTransport(tc.roundTripper, tc.options...)
+			tr := RoundTripper{
+				Transport:       tc.wrappedTransport,
+				OverrideFromCtx: tc.overrideFromCtx,
+				StaticToken:     tc.staticToken,
+			}
 			resp, err := tr.RoundTrip(req)
 
 			if tc.wantError {
diff --git a/plugin/metrics/prometheus/metricsstore/reader.go b/plugin/metrics/prometheus/metricsstore/reader.go
index ac9b4eed5a1..a5d05e131d3 100644
--- a/plugin/metrics/prometheus/metricsstore/reader.go
+++ b/plugin/metrics/prometheus/metricsstore/reader.go
@@ -263,8 +263,8 @@ func getHTTPRoundTripper(c *config.Configuration, logger *zap.Logger) (rt http.R
 		TLSHandshakeTimeout: 10 * time.Second,
 		TLSClientConfig:     ctlsConfig,
 	}
-	return bearertoken.NewTransport(
-		httpTransport,
-		bearertoken.WithAllowOverrideFromCtx(true),
-	), nil
+	return bearertoken.RoundTripper{
+		Transport:       httpTransport,
+		OverrideFromCtx: true,
+	}, nil
 }

From f5e3801dbc07e4fb8095a077678840c88900edc5 Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Sat, 30 Oct 2021 07:40:48 +1100
Subject: [PATCH 10/12] Fix build

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 pkg/bearertoken/context.go |  5 ++++-
 pkg/es/config/config.go    | 10 +++++-----
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/pkg/bearertoken/context.go b/pkg/bearertoken/context.go
index ebbfdf7eb8c..5320db4520d 100644
--- a/pkg/bearertoken/context.go
+++ b/pkg/bearertoken/context.go
@@ -19,10 +19,13 @@ import "context"
 type contextKeyType string
 
 const (
+	// Key is the key name for the bearer token context value.
+	Key = "bearer.token"
+
 	// StoragePropagationKey is a key for viper configuration to pass this option to storage plugins.
 	StoragePropagationKey = "storage.propagate.token"
 
-	contextKey = contextKeyType("bearer.token")
+	contextKey = contextKeyType(Key)
 )
 
 // ContextWithBearerToken set bearer token in context.
diff --git a/pkg/es/config/config.go b/pkg/es/config/config.go
index 09ed57439bc..7fe513aa776 100644
--- a/pkg/es/config/config.go
+++ b/pkg/es/config/config.go
@@ -519,11 +519,11 @@ func GetHTTPRoundTripper(c *Configuration, logger *zap.Logger) (http.RoundTrippe
 		token = tokenFromFile
 	}
 	if token != "" || c.AllowTokenFromContext {
-		transport = bearertoken.NewTransport(
-			httpTransport,
-			bearertoken.WithAllowOverrideFromCtx(c.AllowTokenFromContext),
-			bearertoken.WithToken(token),
-		)
+		transport = bearertoken.RoundTripper{
+			Transport:       httpTransport,
+			OverrideFromCtx: c.AllowTokenFromContext,
+			StaticToken:     token,
+		}
 	}
 	return transport, nil
 }

From 696cb5f9498bdea07c4a57490c58fef81941acea Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Tue, 2 Nov 2021 22:31:46 +1100
Subject: [PATCH 11/12] Address review comments

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 pkg/bearertoken/context.go                     | 13 ++++---------
 pkg/bearertoken/http.go                        |  5 ++---
 plugin/storage/grpc/README.md                  |  5 +++--
 plugin/storage/grpc/shared/grpc_client.go      |  5 ++++-
 plugin/storage/grpc/shared/grpc_client_test.go |  2 +-
 5 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/pkg/bearertoken/context.go b/pkg/bearertoken/context.go
index 5320db4520d..9e9a7fe8c3a 100644
--- a/pkg/bearertoken/context.go
+++ b/pkg/bearertoken/context.go
@@ -16,17 +16,12 @@ package bearertoken
 
 import "context"
 
-type contextKeyType string
+type contextKeyType int
 
-const (
-	// Key is the key name for the bearer token context value.
-	Key = "bearer.token"
+const contextKey = iota
 
-	// StoragePropagationKey is a key for viper configuration to pass this option to storage plugins.
-	StoragePropagationKey = "storage.propagate.token"
-
-	contextKey = contextKeyType(Key)
-)
+// StoragePropagationKey is a key for viper configuration to pass this option to storage plugins.
+const StoragePropagationKey = "storage.propagate.token"
 
 // ContextWithBearerToken set bearer token in context.
 func ContextWithBearerToken(ctx context.Context, token string) context.Context {
diff --git a/pkg/bearertoken/http.go b/pkg/bearertoken/http.go
index b219effef79..76cb0014bec 100644
--- a/pkg/bearertoken/http.go
+++ b/pkg/bearertoken/http.go
@@ -22,9 +22,8 @@ import (
 )
 
 // PropagationHandler returns a http.Handler containing the logic to extract
-// the Authorization token from the http.Request and inserts it into the http.Request
-// context for easier access to the request token via GetBearerToken for bearer token
-// propagation use cases.
+// the Bearer token from the Authorization header of the http.Request and insert it into request.Context
+// for propagation. The token can be accessed via GetBearerToken.
 func PropagationHandler(logger *zap.Logger, h http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
diff --git a/plugin/storage/grpc/README.md b/plugin/storage/grpc/README.md
index fd64d06c79a..0efca2d4dd0 100644
--- a/plugin/storage/grpc/README.md
+++ b/plugin/storage/grpc/README.md
@@ -180,15 +180,16 @@ When using `--query.bearer-token-propagation=true`, the bearer token will be pro
 import (
     // ... other imports
     "fmt"
-    "github.com/jaegertracing/jaeger/storage/spanstore"
     "google.golang.org/grpc/metadata"
+
+    "github.com/jaegertracing/jaeger/plugin/storage/grpc"
 )
 
 // ... spanReader type declared here
 
 func (r *spanReader) extractBearerToken(ctx context.Context) (string, bool) {
 	if md, ok := metadata.FromIncomingContext(ctx); ok {
-		values := md.Get(spanstore.BearerTokenKey)
+		values := md.Get(grpc.BearerTokenKey)
 		if len(values) > 0 {
 			return values[0], true
 		}
diff --git a/plugin/storage/grpc/shared/grpc_client.go b/plugin/storage/grpc/shared/grpc_client.go
index 0f3d49f1880..33d166de717 100644
--- a/plugin/storage/grpc/shared/grpc_client.go
+++ b/plugin/storage/grpc/shared/grpc_client.go
@@ -31,6 +31,9 @@ import (
 	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
+// BearerTokenKey is the key name for the bearer token context value.
+const BearerTokenKey = "bearer.token"
+
 var (
 	_ StoragePlugin        = (*grpcClient)(nil)
 	_ ArchiveStoragePlugin = (*grpcClient)(nil)
@@ -74,7 +77,7 @@ func upgradeContextWithBearerToken(ctx context.Context) context.Context {
 		if !ok {
 			md = metadata.New(nil)
 		}
-		md.Set(bearertoken.Key, bearerToken)
+		md.Set(BearerTokenKey, bearerToken)
 		return metadata.NewOutgoingContext(ctx, md)
 	}
 	return ctx
diff --git a/plugin/storage/grpc/shared/grpc_client_test.go b/plugin/storage/grpc/shared/grpc_client_test.go
index 518f95bc7aa..51e1252a5d9 100644
--- a/plugin/storage/grpc/shared/grpc_client_test.go
+++ b/plugin/storage/grpc/shared/grpc_client_test.go
@@ -113,7 +113,7 @@ func TestContextUpgradeWithToken(t *testing.T) {
 	upgradedToken := upgradeContextWithBearerToken(ctx)
 	md, ok := metadata.FromOutgoingContext(upgradedToken)
 	assert.Truef(t, ok, "Expected metadata in context")
-	bearerTokenFromMetadata := md.Get(bearertoken.Key)
+	bearerTokenFromMetadata := md.Get(BearerTokenKey)
 	assert.Equal(t, []string{testBearerToken}, bearerTokenFromMetadata)
 }
 

From 2742a11eb74355740ff644c474cc128d05deaee3 Mon Sep 17 00:00:00 2001
From: albertteoh <see.kwang.teoh@gmail.com>
Date: Wed, 3 Nov 2021 01:03:17 +1100
Subject: [PATCH 12/12] Fix staticcheck err

Signed-off-by: albertteoh <see.kwang.teoh@gmail.com>
---
 pkg/bearertoken/context.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pkg/bearertoken/context.go b/pkg/bearertoken/context.go
index 9e9a7fe8c3a..4f9221fe631 100644
--- a/pkg/bearertoken/context.go
+++ b/pkg/bearertoken/context.go
@@ -18,7 +18,7 @@ import "context"
 
 type contextKeyType int
 
-const contextKey = iota
+const contextKey = contextKeyType(iota)
 
 // StoragePropagationKey is a key for viper configuration to pass this option to storage plugins.
 const StoragePropagationKey = "storage.propagate.token"