Add CustomHeaderMatcher to pass additional headers (flyteorg#5563)

Signed-off-by: Andrew Dye <[email protected]>
Signed-off-by: Vladyslav Libov <[email protected]>

flyteadmin/auth/handlers.go

@@ -5,11 +5,13 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"net/textproto"
 	"net/url"
 	"strings"
 	"time"
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/util/metautils"
+	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
 	"golang.org/x/oauth2"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -22,6 +24,7 @@ import (
 	"github.com/flyteorg/flyte/flyteadmin/pkg/common"
 	"github.com/flyteorg/flyte/flyteadmin/plugins"
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service"
+	"github.com/flyteorg/flyte/flytestdlib/contextutils"
 	"github.com/flyteorg/flyte/flytestdlib/errors"
 	"github.com/flyteorg/flyte/flytestdlib/logger"
 )
@@ -32,6 +35,8 @@ const (
 	FromHTTPVal          = "true"
 )
 
+var XRequestID = textproto.CanonicalMIMEHeaderKey(contextutils.RequestIDKey.String())
+
 type PreRedirectHookError struct {
 	Message string
 	Code    int
@@ -533,3 +538,18 @@ func GetUserInfoForwardResponseHandler() UserInfoForwardResponseHandler {
 		return nil
 	}
 }
+
+func GetCustomHeaderMatcher(pluginRegistry *plugins.Registry) runtime.HeaderMatcherFunc {
+	if fn := plugins.Get[runtime.HeaderMatcherFunc](pluginRegistry, plugins.PluginIDCustomerHeaderMatcher); fn != nil {
+		return fn
+	}
+	return func(key string) (string, bool) {
+		canonicalKey := textproto.CanonicalMIMEHeaderKey(key)
+		switch canonicalKey {
+		case XRequestID:
+			return canonicalKey, true
+		default:
+			return runtime.DefaultHeaderMatcher(key)
+		}
+	}
+}

flyteadmin/pkg/server/service.go

@@ -240,6 +240,9 @@ func newHTTPServer(ctx context.Context, pluginRegistry *plugins.Registry, cfg *c
 	// This option sets subject in the user info response
 	gwmuxOptions = append(gwmuxOptions, runtime.WithForwardResponseOption(auth.GetUserInfoForwardResponseHandler()))
 
+	// Use custom header matcher to allow additional headers to be passed through
+	gwmuxOptions = append(gwmuxOptions, runtime.WithIncomingHeaderMatcher(auth.GetCustomHeaderMatcher(pluginRegistry)))
+
 	if cfg.Security.UseAuth {
 		// Add HTTP handlers for OIDC endpoints
 		auth.RegisterHandlers(ctx, mux, authCtx, pluginRegistry)

flyteadmin/plugins/registry.go

@@ -9,12 +9,13 @@ import (
 type PluginID = string
 
 const (
-	PluginIDWorkflowExecutor       PluginID = "WorkflowExecutor"
+	PluginIDAdditionalGRPCService  PluginID = "AdditionalGRPCService"
+	PluginIDCustomerHeaderMatcher  PluginID = "CustomerHeaderMatcher"
 	PluginIDDataProxy              PluginID = "DataProxy"
-	PluginIDUnaryServiceMiddleware PluginID = "UnaryServiceMiddleware"
-	PluginIDPreRedirectHook        PluginID = "PreRedirectHook"
 	PluginIDLogoutHook             PluginID = "LogoutHook"
-	PluginIDAdditionalGRPCService  PluginID = "AdditionalGRPCService"
+	PluginIDPreRedirectHook        PluginID = "PreRedirectHook"
+	PluginIDUnaryServiceMiddleware PluginID = "UnaryServiceMiddleware"
+	PluginIDWorkflowExecutor       PluginID = "WorkflowExecutor"
 )
 
 type AtomicRegistry struct {