fix oidc and token exchange config for k8s (#2537)

yrobla · taskbot · web-flow · commit 45716962c6f0 · 2025-11-13T13:20:47.000+01:00
* allow to consume the resource url for crd

* fixes from review

* fix bugs

* fix tests

* add token type to exchange

* properly set token exchange

---------

Co-authored-by: taskbot &lt;taskbot@users.noreply.github.com&gt;
diff --git a/cmd/thv-operator/api/v1alpha1/mcpexternalauthconfig_types.go b/cmd/thv-operator/api/v1alpha1/mcpexternalauthconfig_types.go
@@ -52,6 +52,16 @@ type TokenExchangeConfig struct {
 	// +optional
 	Scopes []string `json:"scopes,omitempty"`
 
+	// SubjectTokenType is the type of the incoming subject token.
+	// Accepts short forms: "access_token" (default), "id_token", "jwt"
+	// Or full URNs: "urn:ietf:params:oauth:token-type:access_token",
+	//               "urn:ietf:params:oauth:token-type:id_token",
+	//               "urn:ietf:params:oauth:token-type:jwt"
+	// For Google Workload Identity Federation with OIDC providers (like Okta), use "id_token"
+	// +kubebuilder:validation:Pattern=`^(access_token|id_token|jwt|urn:ietf:params:oauth:token-type:(access_token|id_token|jwt))?$`
+	// +optional
+	SubjectTokenType string `json:"subjectTokenType,omitempty"`
+
 	// ExternalTokenHeaderName is the name of the custom header to use for the exchanged token.
 	// If set, the exchanged token will be added to this custom header (e.g., "X-Upstream-Token").
 	// If empty or not set, the exchanged token will replace the Authorization header (default behavior).
diff --git a/cmd/thv-operator/controllers/mcpremoteproxy_runconfig.go b/cmd/thv-operator/controllers/mcpremoteproxy_runconfig.go
@@ -200,13 +200,24 @@ func (r *MCPRemoteProxyReconciler) createRunConfigFromMCPRemoteProxy(
 	// Use the RunConfigBuilder for operator context
 	// Deployer is nil for remote proxies because they connect to external services
 	// and do not require container deployment (unlike MCPServer which deploys containers)
-	return runner.NewOperatorRunConfigBuilder(
+	runConfig, err := runner.NewOperatorRunConfigBuilder(
 		context.Background(),
 		nil,
 		nil,
 		nil,
 		options...,
 	)
+	if err != nil {
+		return nil, err
+	}
+
+	// Populate middleware configs from the configuration fields
+	// This ensures that middleware_configs is properly set for serialization
+	if err := runner.PopulateMiddlewareConfigs(runConfig); err != nil {
+		return nil, fmt.Errorf("failed to populate middleware configs: %w", err)
+	}
+
+	return runConfig, nil
 }
 
 // validateRunConfigForRemoteProxy validates a RunConfig for remote proxy deployments
diff --git a/cmd/thv-operator/controllers/mcpserver_externalauth_runconfig_test.go b/cmd/thv-operator/controllers/mcpserver_externalauth_runconfig_test.go
@@ -29,6 +29,7 @@ import (
 	ctrlutil "github.com/stacklok/toolhive/cmd/thv-operator/pkg/controllerutil"
 	"github.com/stacklok/toolhive/pkg/container/kubernetes"
 	"github.com/stacklok/toolhive/pkg/runner"
+	"github.com/stacklok/toolhive/pkg/transport/types"
 )
 
 // TestAddExternalAuthConfigOptions tests the addExternalAuthConfigOptions function
@@ -532,14 +533,23 @@ func TestCreateRunConfigFromMCPServer_WithExternalAuth(t *testing.T) {
 				assert.Equal(t, "external-auth-server", config.Name)
 				assert.Equal(t, "test:v1", config.Image)
 
-				// Verify middleware config was added
+				// Verify middleware configs are populated (auth, tokenexchange, mcp-parser, usagemetrics)
 				assert.NotNil(t, config.MiddlewareConfigs)
-				assert.Len(t, config.MiddlewareConfigs, 1)
-				assert.Equal(t, "tokenexchange", config.MiddlewareConfigs[0].Type)
+				assert.GreaterOrEqual(t, len(config.MiddlewareConfigs), 1, "Should have at least tokenexchange middleware")
+
+				// Find the tokenexchange middleware
+				var tokenExchangeMw *types.MiddlewareConfig
+				for i := range config.MiddlewareConfigs {
+					if config.MiddlewareConfigs[i].Type == "tokenexchange" {
+						tokenExchangeMw = &config.MiddlewareConfigs[i]
+						break
+					}
+				}
+				require.NotNil(t, tokenExchangeMw, "tokenexchange middleware should be present")
 
 				// Verify middleware parameters
 				var params map[string]interface{}
-				err := json.Unmarshal(config.MiddlewareConfigs[0].Parameters, &params)
+				err := json.Unmarshal(tokenExchangeMw.Parameters, &params)
 				require.NoError(t, err)
 
 				tokenExchangeConfig, ok := params["token_exchange_config"].(map[string]interface{})
diff --git a/cmd/thv-operator/controllers/mcpserver_runconfig.go b/cmd/thv-operator/controllers/mcpserver_runconfig.go
@@ -218,13 +218,24 @@ func (r *MCPServerReconciler) createRunConfigFromMCPServer(m *mcpv1alpha1.MCPSer
 	}
 
 	// Use the RunConfigBuilder for operator context with full builder pattern
-	return runner.NewOperatorRunConfigBuilder(
+	runConfig, err := runner.NewOperatorRunConfigBuilder(
 		context.Background(),
 		nil,
 		envVars,
 		nil,
 		options...,
 	)
+	if err != nil {
+		return nil, err
+	}
+
+	// Populate middleware configs from the configuration fields
+	// This ensures that middleware_configs is properly set for serialization
+	if err := runner.PopulateMiddlewareConfigs(runConfig); err != nil {
+		return nil, fmt.Errorf("failed to populate middleware configs: %w", err)
+	}
+
+	return runConfig, nil
 }
 
 // labelsForRunConfig returns labels for run config ConfigMap
diff --git a/cmd/thv-operator/pkg/controllerutil/tokenexchange.go b/cmd/thv-operator/pkg/controllerutil/tokenexchange.go
@@ -2,7 +2,6 @@ package controllerutil
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 
 	corev1 "k8s.io/api/core/v1"
@@ -12,7 +11,6 @@ import (
 	mcpv1alpha1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1alpha1"
 	"github.com/stacklok/toolhive/pkg/auth/tokenexchange"
 	"github.com/stacklok/toolhive/pkg/runner"
-	transporttypes "github.com/stacklok/toolhive/pkg/transport/types"
 )
 
 // GenerateOpenTelemetryEnvVars generates OpenTelemetry environment variables
@@ -93,7 +91,8 @@ func GenerateTokenExchangeEnvVars(
 }
 
 // AddExternalAuthConfigOptions adds external authentication configuration options to builder options
-// This creates middleware configuration for token exchange and is shared between MCPServer and MCPRemoteProxy
+// This creates token exchange configuration which will be automatically converted to middleware by
+// PopulateMiddlewareConfigs() when the runner starts. This ensures correct middleware ordering.
 func AddExternalAuthConfigOptions(
 	ctx context.Context,
 	c client.Client,
@@ -138,55 +137,34 @@ func AddExternalAuthConfigOptions(
 		}
 	}
 
-	// Use scopes array directly from spec
-	scopes := tokenExchangeSpec.Scopes
-
 	// Determine header strategy based on ExternalTokenHeaderName
 	headerStrategy := "replace" // Default strategy
 	if tokenExchangeSpec.ExternalTokenHeaderName != "" {
 		headerStrategy = "custom"
 	}
 
-	// Build token exchange middleware configuration
-	// Client secret is provided via TOOLHIVE_TOKEN_EXCHANGE_CLIENT_SECRET environment variable
-	// to avoid embedding plaintext secrets in the ConfigMap
-	tokenExchangeConfig := map[string]interface{}{
-		"token_url": tokenExchangeSpec.TokenURL,
-		"client_id": tokenExchangeSpec.ClientID,
-		"audience":  tokenExchangeSpec.Audience,
-	}
-
-	if len(scopes) > 0 {
-		tokenExchangeConfig["scopes"] = scopes
-	}
-
-	if headerStrategy != "" {
-		tokenExchangeConfig["header_strategy"] = headerStrategy
-	}
-
-	if tokenExchangeSpec.ExternalTokenHeaderName != "" {
-		tokenExchangeConfig["external_token_header_name"] = tokenExchangeSpec.ExternalTokenHeaderName
-	}
-
-	// Create middleware parameters
-	middlewareParams := map[string]interface{}{
-		"token_exchange_config": tokenExchangeConfig,
-	}
-
-	// Marshal parameters to JSON
-	paramsJSON, err := json.Marshal(middlewareParams)
+	// Normalize SubjectTokenType to full URN (accepts both short forms and full URNs)
+	normalizedTokenType, err := tokenexchange.NormalizeTokenType(tokenExchangeSpec.SubjectTokenType)
 	if err != nil {
-		return fmt.Errorf("failed to marshal token exchange middleware parameters: %w", err)
+		return fmt.Errorf("invalid subject token type: %w", err)
 	}
 
-	// Create middleware config
-	middlewareConfig := transporttypes.MiddlewareConfig{
-		Type:       tokenexchange.MiddlewareType,
-		Parameters: json.RawMessage(paramsJSON),
-	}
-
-	// Add to options using the WithMiddlewareConfig builder option
-	*options = append(*options, runner.WithMiddlewareConfig([]transporttypes.MiddlewareConfig{middlewareConfig}))
+	// Build token exchange configuration
+	// Client secret is provided via TOOLHIVE_TOKEN_EXCHANGE_CLIENT_SECRET environment variable
+	// to avoid embedding plaintext secrets in the ConfigMap
+	tokenExchangeConfig := &tokenexchange.Config{
+		TokenURL:                tokenExchangeSpec.TokenURL,
+		ClientID:                tokenExchangeSpec.ClientID,
+		Audience:                tokenExchangeSpec.Audience,
+		Scopes:                  tokenExchangeSpec.Scopes,
+		SubjectTokenType:        normalizedTokenType,
+		HeaderStrategy:          headerStrategy,
+		ExternalTokenHeaderName: tokenExchangeSpec.ExternalTokenHeaderName,
+	}
+
+	// Use WithTokenExchangeConfig to add configuration
+	// The middleware will be automatically created by PopulateMiddlewareConfigs() in the correct order
+	*options = append(*options, runner.WithTokenExchangeConfig(tokenExchangeConfig))
 
 	return nil
 }
diff --git a/deploy/charts/operator-crds/Chart.yaml b/deploy/charts/operator-crds/Chart.yaml
@@ -2,5 +2,5 @@ apiVersion: v2
 name: toolhive-operator-crds
 description: A Helm chart for installing the ToolHive Operator CRDs into Kubernetes.
 type: application
-version: 0.0.54
+version: 0.0.55
 appVersion: "0.0.1"
diff --git a/deploy/charts/operator-crds/README.md b/deploy/charts/operator-crds/README.md
@@ -1,7 +1,7 @@
 
 # ToolHive Operator CRDs Helm Chart
 
-![Version: 0.0.54](https://img.shields.io/badge/Version-0.0.54-informational?style=flat-square)
+![Version: 0.0.55](https://img.shields.io/badge/Version-0.0.55-informational?style=flat-square)
 ![Type: application](https://img.shields.io/badge/Type-application-informational?style=flat-square)
 
 A Helm chart for installing the ToolHive Operator CRDs into Kubernetes.
diff --git a/deploy/charts/operator-crds/crds/toolhive.stacklok.dev_mcpexternalauthconfigs.yaml b/deploy/charts/operator-crds/crds/toolhive.stacklok.dev_mcpexternalauthconfigs.yaml
@@ -97,6 +97,16 @@ spec:
                     items:
                       type: string
                     type: array
+                  subjectTokenType:
+                    description: |-
+                      SubjectTokenType is the type of the incoming subject token.
+                      Accepts short forms: "access_token" (default), "id_token", "jwt"
+                      Or full URNs: "urn:ietf:params:oauth:token-type:access_token",
+                                    "urn:ietf:params:oauth:token-type:id_token",
+                                    "urn:ietf:params:oauth:token-type:jwt"
+                      For Google Workload Identity Federation with OIDC providers (like Okta), use "id_token"
+                    pattern: ^(access_token|id_token|jwt|urn:ietf:params:oauth:token-type:(access_token|id_token|jwt))?$
+                    type: string
                   tokenUrl:
                     description: TokenURL is the OAuth 2.0 token endpoint URL for
                       token exchange
diff --git a/docs/operator/crd-api.md b/docs/operator/crd-api.md
diff --git a/docs/server/docs.go b/docs/server/docs.go
diff --git a/docs/server/swagger.json b/docs/server/swagger.json
diff --git a/docs/server/swagger.yaml b/docs/server/swagger.yaml
diff --git a/pkg/auth/tokenexchange/exchange.go b/pkg/auth/tokenexchange/exchange.go
@@ -243,19 +243,14 @@ func (c *ExchangeConfig) Validate() error {
 	// don't require client credentials and rely on the trust relationship
 	// configured in the identity provider (e.g., Workload Identity Federation)
 
-	// Validate SubjectTokenType if provided
+	// Validate and normalize SubjectTokenType if provided
 	if c.SubjectTokenType != "" {
-		validTypes := []string{tokenTypeAccessToken, tokenTypeIDToken, tokenTypeJWT}
-		isValid := false
-		for _, validType := range validTypes {
-			if c.SubjectTokenType == validType {
-				isValid = true
-				break
-			}
-		}
-		if !isValid {
-			return fmt.Errorf("invalid SubjectTokenType %q: must be one of %v", c.SubjectTokenType, validTypes)
+		normalized, err := NormalizeTokenType(c.SubjectTokenType)
+		if err != nil {
+			return fmt.Errorf("invalid SubjectTokenType: %w", err)
 		}
+		// Update the config with the normalized value
+		c.SubjectTokenType = normalized
 	}
 
 	// Validate URL format
diff --git a/pkg/auth/tokenexchange/middleware.go b/pkg/auth/tokenexchange/middleware.go
@@ -345,6 +345,9 @@ func createTokenExchangeMiddleware(
 				return
 			}
 
+			// Token exchange succeeded
+			logger.Debug("Token exchange successful")
+
 			// Inject the exchanged token into the request using the pre-selected strategy
 			if err := injectToken(r, exchangedToken.AccessToken); err != nil {
 				logger.Warnf("Failed to inject token: %v", err)
diff --git a/pkg/runner/config.go b/pkg/runner/config.go
@@ -11,6 +11,7 @@ import (
 	"github.com/stacklok/toolhive/pkg/auth"
 	authoauth "github.com/stacklok/toolhive/pkg/auth/oauth"
 	"github.com/stacklok/toolhive/pkg/auth/remote"
+	"github.com/stacklok/toolhive/pkg/auth/tokenexchange"
 	"github.com/stacklok/toolhive/pkg/authz"
 	"github.com/stacklok/toolhive/pkg/container"
 	rt "github.com/stacklok/toolhive/pkg/container/runtime"
@@ -100,6 +101,9 @@ type RunConfig struct {
 	// OIDCConfig contains OIDC configuration
 	OIDCConfig *auth.TokenValidatorConfig `json:"oidc_config,omitempty" yaml:"oidc_config,omitempty"`
 
+	// TokenExchangeConfig contains token exchange configuration for external authentication
+	TokenExchangeConfig *tokenexchange.Config `json:"token_exchange_config,omitempty" yaml:"token_exchange_config,omitempty"`
+
 	// AuthzConfig contains the authorization configuration
 	AuthzConfig *authz.Config `json:"authz_config,omitempty" yaml:"authz_config,omitempty"`
 
diff --git a/pkg/runner/config_builder.go b/pkg/runner/config_builder.go
@@ -357,6 +357,14 @@ func WithOIDCConfig(
 	}
 }
 
+// WithTokenExchangeConfig sets the token exchange configuration
+func WithTokenExchangeConfig(config *tokenexchange.Config) RunConfigBuilderOption {
+	return func(b *runConfigBuilder) error {
+		b.config.TokenExchangeConfig = config
+		return nil
+	}
+}
+
 // WithTelemetryConfig configures telemetry settings (legacy - custom attributes handled via middleware)
 func WithTelemetryConfig(
 	otelEndpoint string,
diff --git a/pkg/runner/middleware.go b/pkg/runner/middleware.go
diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go
