Add auth token propagation for metrics reader

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)

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))

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
+}

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) {
+	const token = "blah"
+	ctx := context.Background()
+	ctx = ContextWithBearerToken(ctx, token)
+	contextToken, ok := GetBearerToken(ctx)
+	assert.True(t, ok)
+	assert.Equal(t, contextToken, token)
+}

cmd/query/app/token_propagation_handler.go → pkg/bearertoken/http.go

@@ -12,18 +12,20 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package app
+package bearertoken
 
 import (
 	"net/http"
 	"strings"
 
 	"go.uber.org/zap"
-
-	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
-func bearerTokenPropagationHandler(logger *zap.Logger, h http.Handler) http.Handler {
+// 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 +42,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))
 		}
 	})
-
 }

cmd/query/app/token_propagation_hander_test.go → pkg/bearertoken/http_test.go

@@ -12,41 +12,44 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package app
+package bearertoken
 
 import (
 	"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_PropagationHandler(t *testing.T) {
+	httpClient := &http.Client{
+		Timeout: 2 * time.Second,
+	}
+
 	logger := zap.NewNop()
-	bearerToken := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJhZG1pbiIsIm5hbWUiOiJKb2huIERvZSIsImlhdCI"
+	const bearerToken = "blah"
 
 	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 +71,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 +84,4 @@ func Test_bearTokenPropagationHandler(t *testing.T) {
 			stop.Wait()
 		})
 	}
-
 }

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)
+}

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)
+			}
+		})
+	}
+}