Address review comments

Signed-off-by: Ed Snible <[email protected]>
jaegertracing · Jul 13, 2022 · 27b2bb8 · 27b2bb8
1 parent c8aa9d0
commit 27b2bb8
Show file tree

Hide file tree

Showing 7 changed files with 49 additions and 50 deletions.
diff --git a/cmd/query/app/handler_options.go b/cmd/query/app/handler_options.go
@@ -63,9 +63,9 @@ func (handlerOptions) QueryLookbackDuration(queryLookbackDuration time.Duration)
 }
 
 // Tracer creates a HandlerOption that initializes tenancy
-func (handlerOptions) Tenancy(options *tenancy.Options) HandlerOption {
+func (handlerOptions) Tenancy(tc *tenancy.TenancyConfig) HandlerOption {
 	return func(apiHandler *APIHandler) {
-		apiHandler.tenancyConfig = tenancy.NewTenancyConfig(options)
+		apiHandler.tenancyConfig = tc
 	}
 }
 

diff --git a/cmd/query/app/http_handler.go b/cmd/query/app/http_handler.go
@@ -147,7 +147,7 @@ func (aH *APIHandler) handleFunc(
 	var handler http.Handler
 	handler = http.HandlerFunc(f)
 	if aH.tenancyConfig.Enabled {
-		handler = tenancy.ExtractTenantHTTPHandler(aH.tenancyConfig, handler, zap.NewNop())
+		handler = tenancy.ExtractTenantHTTPHandler(aH.tenancyConfig, handler)
 	}
 	traceMiddleware := nethttp.Middleware(
 		aH.tracer,

diff --git a/cmd/query/app/http_handler_test.go b/cmd/query/app/http_handler_test.go
@@ -116,7 +116,7 @@ func initializeTestServerWithOptions(tenancyConfig tenancy.TenancyConfig, queryO
 	handler := NewAPIHandler(qs, options...)
 	handler.RegisterRoutes(r)
 	return &testServer{
-		server:           httptest.NewServer(tenancy.ExtractTenantHTTPHandler(&tenancyConfig, r, zap.NewNop())),
+		server:           httptest.NewServer(tenancy.ExtractTenantHTTPHandler(&tenancyConfig, r)),
 		spanReader:       readStorage,
 		dependencyReader: dependencyStorage,
 		handler:          handler,

diff --git a/cmd/query/app/server.go b/cmd/query/app/server.go
@@ -157,7 +157,7 @@ func createHTTPServer(querySvc *querysvc.QueryService, metricsQuerySvc querysvc.
 		HandlerOptions.Logger(logger),
 		HandlerOptions.Tracer(tracer),
 		HandlerOptions.MetricsQueryService(metricsQuerySvc),
-		HandlerOptions.Tenancy(&queryOpts.Tenancy),
+		HandlerOptions.Tenancy(tenancy.NewTenancyConfig(&queryOpts.Tenancy)),
 	}
 
 	apiHandler := NewAPIHandler(

diff --git a/pkg/tenancy/grpc.go b/pkg/tenancy/grpc.go
@@ -33,39 +33,55 @@ func (tss *tenantedServerStream) Context() context.Context {
 	return tss.context
 }
 
+func getValidTenant(ctx context.Context, tc *TenancyConfig) (string, error) {
+	// Handle case where tenant is already directly in the context
+	tenant := GetTenant(ctx)
+	if tenant != "" {
+		if !tc.Valid(tenant) {
+			return tenant, status.Errorf(codes.PermissionDenied, "unknown tenant")
+		}
+		return tenant, nil
+	}
+
+	// Handle case where tenant is in the context metadata
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return "", status.Errorf(codes.PermissionDenied, "missing tenant header")
+	}
+
+	var err error
+	tenant, err = tenantFromMetadata(md, tc.Header)
+	if err != nil {
+		return "", err
+	}
+	if !tc.Valid(tenant) {
+		return tenant, status.Errorf(codes.PermissionDenied, "unknown tenant")
+	}
+
+	return tenant, nil
+}
+
+func directlyAttachedTenant(ctx context.Context) bool {
+	return GetTenant(ctx) != ""
+}
+
 // NewGuardingStreamInterceptor blocks handling of streams whose tenancy header doesn't meet tenancy requirements.
 // It also ensures the tenant is directly in the context, rather than context metadata.
 func NewGuardingStreamInterceptor(tc *TenancyConfig) grpc.StreamServerInterceptor {
 	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
-		ctx := ss.Context()
-		// Handle case where tenant is directly in the context
-		tenant := GetTenant(ctx)
-		if tenant != "" {
-			if !tc.Valid(tenant) {
-				return status.Errorf(codes.PermissionDenied, "unknown tenant header")
-			}
-			return handler(srv, ss)
-		}
-
-		// Handle case where tenant is in the context metadata
-		md, ok := metadata.FromIncomingContext(ctx)
-		if !ok {
-			return status.Errorf(codes.PermissionDenied, "missing tenant header")
-		}
-
-		var err error
-		tenant, err = tenantFromMetadata(md, tc.Header)
+		tenant, err := getValidTenant(ss.Context(), tc)
 		if err != nil {
 			return err
 		}
-		if !tc.Valid(tenant) {
-			return status.Errorf(codes.PermissionDenied, "unknown tenant")
+
+		if directlyAttachedTenant(ss.Context()) {
+			return handler(srv, ss)
 		}
 
-		// Apply the tenant directly the context (in addition to metadata)
+		// "upgrade" the tenant to be part of the context, rather than just incoming metadata
 		return handler(srv, &tenantedServerStream{
 			ServerStream: ss,
-			context:      WithTenant(ctx, tenant),
+			context:      WithTenant(ss.Context(), tenant),
 		})
 	}
 }
@@ -85,28 +101,13 @@ func tenantFromMetadata(md metadata.MD, tenancyHeader string) (string, error) {
 // It also ensures the tenant is directly in the context, rather than context metadata.
 func NewGuardingUnaryInterceptor(tc *TenancyConfig) grpc.UnaryServerInterceptor {
 	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-		// Handle case where tenant is directly in the context
-		tenant := GetTenant(ctx)
-		if tenant != "" {
-			if !tc.Valid(tenant) {
-				return nil, status.Errorf(codes.PermissionDenied, "unknown tenant header")
-			}
-			return handler(ctx, req)
-		}
-
-		// Handle case where tenant is in the context metadata
-		md, ok := metadata.FromIncomingContext(ctx)
-		if !ok {
-			return nil, status.Errorf(codes.PermissionDenied, "missing tenant header")
-		}
-
-		var err error
-		tenant, err = tenantFromMetadata(md, tc.Header)
+		tenant, err := getValidTenant(ctx, tc)
 		if err != nil {
 			return nil, err
 		}
-		if !tc.Valid(tenant) {
-			return nil, status.Errorf(codes.PermissionDenied, "unknown tenant")
+
+		if directlyAttachedTenant(ctx) {
+			return handler(ctx, req)
 		}
 
 		return handler(WithTenant(ctx, tenant), req)

diff --git a/pkg/tenancy/http.go b/pkg/tenancy/http.go
@@ -18,14 +18,13 @@ import (
 	"context"
 	"net/http"
 
-	"go.uber.org/zap"
 	"google.golang.org/grpc/metadata"
 )
 
 // PropagationHandler returns a http.Handler containing the logic to extract
 // the tenancy header of the http.Request and insert the tenant into request.Context
 // for propagation. The token can be accessed via tenancy.GetTenant().
-func ExtractTenantHTTPHandler(tc *TenancyConfig, h http.Handler, logger *zap.Logger) http.Handler {
+func ExtractTenantHTTPHandler(tc *TenancyConfig, h http.Handler) http.Handler {
 	if !tc.Enabled {
 		return h
 	}

diff --git a/pkg/tenancy/http_test.go b/pkg/tenancy/http_test.go
@@ -23,7 +23,6 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"go.uber.org/zap"
 )
 
 type testHttpHandler struct {
@@ -70,7 +69,7 @@ func TestProgationHandler(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			handler := &testHttpHandler{}
-			propH := ExtractTenantHTTPHandler(test.tenancyConfig, handler, zap.NewNop())
+			propH := ExtractTenantHTTPHandler(test.tenancyConfig, handler)
 			req, err := http.NewRequest("GET", "/", strings.NewReader(""))
 			for k, vs := range test.requestHeaders {
 				for _, v := range vs {