Skip to content

Commit 50ad13b

Browse files
mtparetmtparet
committed
refactor(mcp): use Envoy Gateway's ClaimToHeader type
Replace the custom ClaimToHeader type definition with egv1a1.ClaimToHeader to reduce code duplication between AI Gateway and Envoy Gateway. fix codegen to append license header Signed-off-by: Matthieu Paret <mp@blackfuel.ai>
1 parent 8294327 commit 50ad13b

9 files changed

Lines changed: 91 additions & 66 deletions

File tree

‎Makefile‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -122,20 +122,20 @@ codegen: ## Generate typed client, listers, and informers for the API.
122122
--clientset-name="versioned" \
123123
--input-base="" \
124124
--input="github.com/envoyproxy/ai-gateway/api/v1alpha1" \
125-
--go-header-file=/dev/null \
125+
--go-header-file=hack/boilerplate.go.txt \
126126
--output-dir="./api/v1alpha1/client/clientset" \
127127
--output-pkg="github.com/envoyproxy/ai-gateway/api/v1alpha1/client/clientset" \
128128
--plural-exceptions="BackendSecurityPolicy:BackendSecurityPolicies"
129129
@echo "codegen => generating listers..."
130130
@$(GO_TOOL) lister-gen \
131-
--go-header-file=/dev/null \
131+
--go-header-file=hack/boilerplate.go.txt \
132132
--output-dir="./api/v1alpha1/client/listers" \
133133
--output-pkg="github.com/envoyproxy/ai-gateway/api/v1alpha1/client/listers" \
134134
--plural-exceptions="BackendSecurityPolicy:BackendSecurityPolicies" \
135135
"github.com/envoyproxy/ai-gateway/api/v1alpha1"
136136
@echo "codegen => generating informers..."
137137
@$(GO_TOOL) informer-gen \
138-
--go-header-file=/dev/null \
138+
--go-header-file=hack/boilerplate.go.txt \
139139
--versioned-clientset-package="github.com/envoyproxy/ai-gateway/api/v1alpha1/client/clientset/versioned" \
140140
--listers-package="github.com/envoyproxy/ai-gateway/api/v1alpha1/client/listers" \
141141
--output-dir="./api/v1alpha1/client/informers" \

‎api/v1alpha1/mcp_route.go‎

Lines changed: 1 addition & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -259,28 +259,7 @@ type MCPRouteOAuth struct {
259259
// +kubebuilder:validation:Optional
260260
// +kubebuilder:validation:MaxItems=16
261261
// +optional
262-
ClaimToHeaders []ClaimToHeader `json:"claimToHeaders,omitempty"`
263-
}
264-
265-
// ClaimToHeader defines a mapping from a JWT claim to an HTTP header.
266-
type ClaimToHeader struct {
267-
// Claim is the JWT claim name to extract. Supports nested claims using dot notation
268-
// (e.g., "realm_access.roles", "resource_access.my-client.roles").
269-
// For array claims, the values will be JSON-encoded in the header.
270-
//
271-
// +kubebuilder:validation:Required
272-
// +kubebuilder:validation:MinLength=1
273-
// +kubebuilder:validation:MaxLength=256
274-
Claim string `json:"claim"`
275-
276-
// Header is the HTTP header name to forward the claim value as.
277-
// The header name should follow HTTP header naming conventions.
278-
//
279-
// +kubebuilder:validation:Required
280-
// +kubebuilder:validation:MinLength=1
281-
// +kubebuilder:validation:MaxLength=256
282-
// +kubebuilder:validation:Pattern=`^[A-Za-z0-9]([A-Za-z0-9-]*[A-Za-z0-9])?$`
283-
Header string `json:"header"`
262+
ClaimToHeaders []egv1a1.ClaimToHeader `json:"claimToHeaders,omitempty"`
284263
}
285264

286265
// MCPRouteAuthorization defines the authorization configuration for a MCPRoute.

‎api/v1alpha1/zz_generated.deepcopy.go‎

Lines changed: 1 addition & 21 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎internal/controller/mcp_route_security_policy.go‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -148,12 +148,7 @@ func (c *MCPRouteController) ensureSecurityPolicy(ctx context.Context, mcpRoute
148148
// Add ClaimToHeaders to extract JWT claims and set them as HTTP headers.
149149
// Envoy's JWT filter will extract these claims and add them to the request headers,
150150
// which can then be forwarded to backend MCP servers.
151-
for _, ctoh := range oauth.ClaimToHeaders {
152-
jwtProvider.ClaimToHeaders = append(jwtProvider.ClaimToHeaders, egv1a1.ClaimToHeader{
153-
Claim: ctoh.Claim,
154-
Header: ctoh.Header,
155-
})
156-
}
151+
jwtProvider.ClaimToHeaders = append(jwtProvider.ClaimToHeaders, oauth.ClaimToHeaders...)
157152

158153
securityPolicySpec.JWT = &egv1a1.JWT{
159154
Providers: []egv1a1.JWTProvider{jwtProvider},

‎internal/controller/mcp_route_security_policy_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -586,7 +586,7 @@ func TestMCPRouteController_syncMCPRouteSecurityPolicy_ClaimToHeaders(t *testing
586586
ProtectedResourceMetadata: aigv1a1.ProtectedResourceMetadata{
587587
Resource: "https://api.example.com/mcp",
588588
},
589-
ClaimToHeaders: []aigv1a1.ClaimToHeader{
589+
ClaimToHeaders: []egv1a1.ClaimToHeader{
590590
{Claim: "sub", Header: "X-User-Id"},
591591
{Claim: "email", Header: "X-User-Email"},
592592
{Claim: "realm_access.roles", Header: "X-User-Roles"},

‎manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_mcproutes.yaml‎

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3240,24 +3240,18 @@ spec:
32403240
- Only the specified claims are extracted; the full JWT is not forwarded to backends
32413241
- Consider using a header prefix (e.g., "X-Jwt-Claim-") to avoid conflicts with other headers
32423242
items:
3243-
description: ClaimToHeader defines a mapping from a JWT
3244-
claim to an HTTP header.
3243+
description: ClaimToHeader defines a configuration to convert
3244+
JWT claims into HTTP headers
32453245
properties:
32463246
claim:
32473247
description: |-
3248-
Claim is the JWT claim name to extract. Supports nested claims using dot notation
3249-
(e.g., "realm_access.roles", "resource_access.my-client.roles").
3250-
For array claims, the values will be JSON-encoded in the header.
3251-
maxLength: 256
3252-
minLength: 1
3248+
Claim is the JWT Claim that should be saved into the header : it can be a nested claim of type
3249+
(eg. "claim.nested.key", "sub"). The nested claim name must use dot "."
3250+
to separate the JSON name path.
32533251
type: string
32543252
header:
3255-
description: |-
3256-
Header is the HTTP header name to forward the claim value as.
3257-
The header name should follow HTTP header naming conventions.
3258-
maxLength: 256
3259-
minLength: 1
3260-
pattern: ^[A-Za-z0-9]([A-Za-z0-9-]*[A-Za-z0-9])?$
3253+
description: Header defines the name of the HTTP request
3254+
header that the JWT Claim will be saved into.
32613255
type: string
32623256
required:
32633257
- claim

‎tests/e2e/mcp_route_oauth_test.go‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,52 @@ func TestMCPRouteOAuth(t *testing.T) {
9696
t.Logf("Successfully connected with valid token and retrieved %d tools", len(tools.Tools))
9797
})
9898

99+
t.Run("claim headers forwarded to backend", func(t *testing.T) {
100+
// This test validates that JWT claims configured via claimToHeaders are extracted
101+
// and forwarded to the backend MCP server. The backend has TEST_EXPECTED_CLAIM_HEADERS
102+
// set, so its middleware will reject any request missing the expected claim headers.
103+
validToken := "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWUsImlhdCI6MTUxNjIzOTAyMn0.NHVaYe26MbtOYhSKkoKYdFVomg4i8ZJd8_-RU8VNbftc4TSMb4bXP3l3YlNWACwyXPGffz5aXHc6lty1Y2t4SWRqGteragsVdZufDn5BlnJl9pdR_kdVFUsra2rWKEofkZeIC4yWytE58sMIihvo9H1ScmmVwBcQP6XETqYd0aSHp1gOa9RdUPDvoXQ5oqygTqVtxaDr6wUFKrKItgBMzWIdNZ6y7O9E0DhEPTbE9rfBo6KTFsHAZnMg4k68CDp2woYIaXbmYTWcvbzIuHO7_37GT79XdIwkm95QJ7hYC9RiwrV7mesbY4PAahERJawntho0my942XheVLmGwLMBkQ" //nolint:gosec // Test JWT token
104+
105+
authHTTPClient := &http.Client{
106+
Timeout: 10 * time.Second,
107+
Transport: &mcpAuthTransport{
108+
token: validToken,
109+
base: http.DefaultTransport,
110+
},
111+
}
112+
113+
client := mcp.NewClient(&mcp.Implementation{Name: "demo-http-client", Version: "0.1.0"}, nil)
114+
ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second)
115+
t.Cleanup(cancel)
116+
117+
var sess *mcp.ClientSession
118+
require.Eventually(t, func() bool {
119+
var err error
120+
sess, err = client.Connect(
121+
ctx,
122+
&mcp.StreamableClientTransport{
123+
Endpoint: fmt.Sprintf("%s/mcp", fwd.Address()),
124+
HTTPClient: authHTTPClient,
125+
}, nil)
126+
if err != nil {
127+
t.Logf("failed to connect to MCP server: %v", err)
128+
return false
129+
}
130+
return true
131+
}, 30*time.Second, 100*time.Millisecond, "failed to connect to MCP server")
132+
t.Cleanup(func() { _ = sess.Close() })
133+
134+
// Call the echo tool — if claim headers are missing or wrong, the backend
135+
// middleware will reject the request and this call will fail.
136+
result, err := sess.CallTool(ctx, &mcp.CallToolParams{
137+
Name: "echo",
138+
Arguments: map[string]interface{}{"message": "claim-header-test"},
139+
})
140+
require.NoError(t, err)
141+
require.False(t, result.IsError, "echo tool should succeed when claim headers are forwarded")
142+
t.Logf("echo tool result: %+v", result)
143+
})
144+
99145
t.Run("without token using MCP client", func(t *testing.T) {
100146
// Create an MCP client and connect to the server over Streamable HTTP.
101147
client := mcp.NewClient(&mcp.Implementation{Name: "demo-http-client", Version: "0.1.0"}, nil)

‎tests/e2e/testdata/mcp_route_oauth.yaml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,11 @@ spec:
5050
group: ""
5151
kind: ConfigMap
5252
name: jwks-configmap
53+
claimToHeaders:
54+
- claim: sub
55+
header: x-jwt-claim-sub
56+
- claim: name
57+
header: x-jwt-claim-name
5358
protectedResourceMetadata:
5459
resource: "https://foo.bar.com/mcp"
5560
resourceName: "example-resource"
@@ -109,6 +114,9 @@ spec:
109114
imagePullPolicy: IfNotPresent
110115
ports:
111116
- containerPort: 1063
117+
env:
118+
- name: TEST_EXPECTED_CLAIM_HEADERS
119+
value: "x-jwt-claim-sub=1234567890,x-jwt-claim-name=John Doe"
112120
---
113121
apiVersion: v1
114122
kind: Service

‎tests/internal/testmcp/server.go‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,29 @@ func NewServer(opts *Options) (*http.Server, *mcp.Server) {
103103
})
104104
}
105105

106+
// Setup claim header validation when environment variable TEST_EXPECTED_CLAIM_HEADERS is set.
107+
// Format: "header1=value1,header2=value2"
108+
expectedClaimHeaders := os.Getenv("TEST_EXPECTED_CLAIM_HEADERS")
109+
if expectedClaimHeaders != "" {
110+
expected := map[string]string{}
111+
for _, pair := range strings.Split(expectedClaimHeaders, ",") {
112+
k, v, ok := strings.Cut(pair, "=")
113+
if ok {
114+
expected[k] = v
115+
}
116+
}
117+
s.AddReceivingMiddleware(func(handler mcp.MethodHandler) mcp.MethodHandler {
118+
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
119+
for header, value := range expected {
120+
if got := req.GetExtra().Header.Get(header); got != value {
121+
return nil, fmt.Errorf("expected header %q=%q, got %q", header, value, got)
122+
}
123+
}
124+
return handler(ctx, method, req)
125+
}
126+
})
127+
}
128+
106129
s.AddPrompt(CodeReviewPrompt, codReviewPromptHandler)
107130
s.AddResource(DummyResource, DummyResourceHandler())
108131
s.AddResourceTemplate(DummyResourceTemplate, DummyResourceHandler())

0 commit comments

Comments
 (0)