From d9210d04c6782c155b537299b9f7923506fe0f95 Mon Sep 17 00:00:00 2001
From: Joe Elliott <number101010@gmail.com>
Date: Fri, 18 Aug 2023 09:00:22 -0400
Subject: [PATCH] Revert "Remove tenant_header_key from config (#2414)" (#2795)

This reverts commit affcf7c3db546bb8657e269d1c71798e07564094.
---
 CHANGELOG.md                    |  1 +
 cmd/tempo-query/tempo/config.go | 14 +++++++++++---
 cmd/tempo-query/tempo/plugin.go | 25 +++++++++++++------------
 3 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3de8d3909db..62ebf830e60 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@
 * [BUGFIX] Only search ingester blocks that fall within the request time range. [#2783](https://github.com/grafana/tempo/pull/2783) (@joe-elliott)
 * [BUGFIX] Fix incorrect metrics for index failures [#2781](https://github.com/grafana/tempo/pull/2781) (@zalegrala)
 * [BUGFIX] Fix panic in the metrics-generator when using multiple tenants with default overrides [#2786](https://github.com/grafana/tempo/pull/2786) (@kvrhdn)
+* [BUGFIX] Restore `tenant_header_key` removed in #2414. [#2786](https://github.com/grafana/tempo/pull/2795) (@joe-elliott)
 
 ## v2.2.0 / 2023-07-31
 
diff --git a/cmd/tempo-query/tempo/config.go b/cmd/tempo-query/tempo/config.go
index 535ed246d69..accf5a15586 100644
--- a/cmd/tempo-query/tempo/config.go
+++ b/cmd/tempo-query/tempo/config.go
@@ -2,14 +2,16 @@ package tempo
 
 import (
 	"github.com/grafana/dskit/crypto/tls"
+	"github.com/jaegertracing/jaeger/plugin/storage/grpc/shared"
 	"github.com/spf13/viper"
 )
 
 // Config holds the configuration for redbull.
 type Config struct {
-	Backend    string           `yaml:"backend"`
-	TLSEnabled bool             `yaml:"tls_enabled" category:"advanced"`
-	TLS        tls.ClientConfig `yaml:",inline"`
+	Backend         string           `yaml:"backend"`
+	TLSEnabled      bool             `yaml:"tls_enabled" category:"advanced"`
+	TLS             tls.ClientConfig `yaml:",inline"`
+	TenantHeaderKey string           `yaml:"tenant_header_key"`
 }
 
 // InitFromViper initializes the options struct with values from Viper
@@ -23,4 +25,10 @@ func (c *Config) InitFromViper(v *viper.Viper) {
 	c.TLS.InsecureSkipVerify = v.GetBool("tls_insecure_skip_verify")
 	c.TLS.CipherSuites = v.GetString("tls_cipher_suites")
 	c.TLS.MinVersion = v.GetString("tls_min_version")
+
+	tenantHeader := v.GetString("tenant_header_key")
+	if tenantHeader == "" {
+		tenantHeader = shared.BearerTokenKey
+	}
+	c.TenantHeaderKey = tenantHeader
 }
diff --git a/cmd/tempo-query/tempo/plugin.go b/cmd/tempo-query/tempo/plugin.go
index 9d295c6b8bb..6a2855433b7 100644
--- a/cmd/tempo-query/tempo/plugin.go
+++ b/cmd/tempo-query/tempo/plugin.go
@@ -24,7 +24,6 @@ import (
 	"google.golang.org/grpc/metadata"
 
 	jaeger "github.com/jaegertracing/jaeger/model"
-	"github.com/jaegertracing/jaeger/plugin/storage/grpc/shared"
 	jaeger_spanstore "github.com/jaegertracing/jaeger/storage/spanstore"
 
 	ot_jaeger "github.com/open-telemetry/opentelemetry-collector-contrib/pkg/translator/jaeger"
@@ -54,10 +53,11 @@ var tlsVersions = map[string]uint16{
 }
 
 type Backend struct {
-	tempoBackend string
-	tlsEnabled   bool
-	tls          tlsCfg.ClientConfig
-	httpClient   *http.Client
+	tempoBackend    string
+	tlsEnabled      bool
+	tls             tlsCfg.ClientConfig
+	httpClient      *http.Client
+	tenantHeaderKey string
 }
 
 func New(cfg *Config) (*Backend, error) {
@@ -66,10 +66,11 @@ func New(cfg *Config) (*Backend, error) {
 		return nil, err
 	}
 	return &Backend{
-		tempoBackend: cfg.Backend,
-		tlsEnabled:   cfg.TLSEnabled,
-		tls:          cfg.TLS,
-		httpClient:   httpClient,
+		tempoBackend:    cfg.Backend,
+		tlsEnabled:      cfg.TLSEnabled,
+		tls:             cfg.TLS,
+		httpClient:      httpClient,
+		tenantHeaderKey: cfg.TenantHeaderKey,
 	}, nil
 }
 
@@ -434,7 +435,7 @@ func (b *Backend) newGetRequest(ctx context.Context, url string, span opentracin
 
 	// currently Jaeger Query will only propagate bearer token to the grpc backend and no other headers
 	// so we are going to extract the tenant id from the header, if it exists and use it
-	tenantID, found := extractBearerToken(ctx)
+	tenantID, found := extractBearerToken(ctx, b.tenantHeaderKey)
 	if found {
 		req.Header.Set(user.OrgIDHeaderName, tenantID)
 	}
@@ -442,9 +443,9 @@ func (b *Backend) newGetRequest(ctx context.Context, url string, span opentracin
 	return req, nil
 }
 
-func extractBearerToken(ctx context.Context) (string, bool) {
+func extractBearerToken(ctx context.Context, tenantHeader string) (string, bool) {
 	if md, ok := metadata.FromIncomingContext(ctx); ok {
-		values := md.Get(shared.BearerTokenKey)
+		values := md.Get(tenantHeader)
 		if len(values) > 0 {
 			return values[0], true
 		}