From 7dd2229340f7f7ab6e708335dd80a587ff07af19 Mon Sep 17 00:00:00 2001
From: Andrew Dye <andrewwdye@gmail.com>
Date: Thu, 11 Jul 2024 10:59:36 -0700
Subject: [PATCH] Add CustomHeaderMatcher to pass additional headers

Signed-off-by: Andrew Dye <andrewwdye@gmail.com>
---
 flyteadmin/auth/handlers.go      | 20 ++++++++++++++++++++
 flyteadmin/pkg/server/service.go |  3 +++
 flyteadmin/plugins/registry.go   |  9 +++++----
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/flyteadmin/auth/handlers.go b/flyteadmin/auth/handlers.go
index a6220db6e3..b839cf26d0 100644
--- a/flyteadmin/auth/handlers.go
+++ b/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)
+		}
+	}
+}
diff --git a/flyteadmin/pkg/server/service.go b/flyteadmin/pkg/server/service.go
index ff80c343d3..296a4ee9da 100644
--- a/flyteadmin/pkg/server/service.go
+++ b/flyteadmin/pkg/server/service.go
@@ -219,6 +219,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)
diff --git a/flyteadmin/plugins/registry.go b/flyteadmin/plugins/registry.go
index 92644b1367..a89a8dfeae 100644
--- a/flyteadmin/plugins/registry.go
+++ b/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 {