Skip to content

Commit f0b7ff9

Browse files
zhaohuabingzhaohuabing
committed
authorization against tool arguments
Signed-off-by: Huabing Zhao <[email protected]>
1 parent a17a1ad commit f0b7ff9

File tree

11 files changed

+335
-53
lines changed

11 files changed

+335
-53
lines changed

‎api/v1alpha1/mcp_route.go‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -308,25 +308,13 @@ type ToolCall struct {
308308
ToolName string `json:"toolName"`
309309

310310
// Arguments defines the arguments that must be present in the tool call for this rule to match.
311+
// Keys must exist and their values must match the provided RE2-compatible regular expressions.
312+
// If the argument is a non-string type, it will be matched against its JSON representation.
311313
//
312314
// +optional
313-
// Arguments map[string]string `json:"arguments,omitempty"`
315+
Arguments map[string]string `json:"arguments,omitempty"`
314316
}
315317

316-
/*type ToolArgument struct {
317-
// Name is the name of the argument.
318-
Name string `json:"name"`
319-
320-
// Value is the value of the argument.
321-
Value ArgumentValues `json:"value"`
322-
}
323-
324-
type ArgumentValues struct {
325-
Include []string `json:"include,omitempty"`
326-
327-
IncludeRegex []string `json:"includeRegex,omitempty"`
328-
}*/
329-
330318
// JWKS defines how to obtain JSON Web Key Sets (JWKS) either from a remote HTTP/HTTPS endpoint or from a local source.
331319
// +kubebuilder:validation:XValidation:rule="has(self.remoteJWKS) || has(self.localJWKS)", message="either remoteJWKS or localJWKS must be specified."
332320
// +kubebuilder:validation:XValidation:rule="!(has(self.remoteJWKS) && has(self.localJWKS))", message="remoteJWKS and localJWKS cannot both be specified."

‎api/v1alpha1/zz_generated.deepcopy.go‎

Lines changed: 10 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎internal/controller/gateway.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -500,6 +500,7 @@ func mcpConfig(mcpRoutes []aigv1a1.MCPRoute) *filterapi.MCPConfig {
500500
tools[i] = filterapi.ToolCall{
501501
BackendName: tool.BackendName,
502502
ToolName: tool.ToolName,
503+
Arguments: tool.Arguments,
503504
}
504505
}
505506

‎internal/filterapi/mcpconfig.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,4 +118,8 @@ type ToolCall struct {
118118

119119
// ToolName is the name of the tool.
120120
ToolName string `json:"toolName"`
121+
122+
// Arguments defines required arguments (exact key names with regex patterns for values).
123+
// All patterns must match for the rule to apply.
124+
Arguments map[string]string `json:"arguments,omitempty"`
121125
}

‎internal/mcpproxy/authorization.go‎

Lines changed: 51 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,19 @@
66
package mcpproxy
77

88
import (
9+
"encoding/json"
910
"errors"
1011
"log/slog"
1112
"net/http"
13+
"regexp"
1214
"strings"
1315

1416
"github.com/golang-jwt/jwt/v5"
1517

1618
"github.com/envoyproxy/ai-gateway/internal/filterapi"
1719
)
1820

19-
func (m *MCPProxy) authorizeRequest(authorization *filterapi.MCPRouteAuthorization, headers http.Header, backendName, toolName string) bool {
21+
func (m *MCPProxy) authorizeRequest(authorization *filterapi.MCPRouteAuthorization, headers http.Header, backendName, toolName string, argments any) bool {
2022
defaultAction := authorization.DefaultAction == filterapi.AuthorizationActionAllow
2123

2224
// If there are no rules, return the default action.
@@ -44,9 +46,14 @@ func (m *MCPProxy) authorizeRequest(authorization *filterapi.MCPRouteAuthorizati
4446
scopeSet[scope] = struct{}{}
4547
}
4648

47-
target := filterapi.ToolCall{BackendName: backendName, ToolName: toolName}
4849
for _, rule := range authorization.Rules {
49-
if !toolTargetMatches(target, rule.Target.Tools) {
50+
var args map[string]any
51+
if argments != nil {
52+
if cast, ok := argments.(map[string]any); ok {
53+
args = cast
54+
}
55+
}
56+
if !toolMatches(args, filterapi.ToolCall{BackendName: backendName, ToolName: toolName}, rule.Target.Tools) {
5057
continue
5158
}
5259
if scopesSatisfied(scopeSet, rule.Source.JWTSource.Scopes) {
@@ -98,15 +105,54 @@ func extractScopes(claims jwt.MapClaims) []string {
98105
}
99106
}
100107

101-
func toolTargetMatches(target filterapi.ToolCall, tools []filterapi.ToolCall) bool {
108+
func toolMatches(args map[string]any, target filterapi.ToolCall, tools []filterapi.ToolCall) bool {
102109
if len(tools) == 0 {
103110
return true
104111
}
112+
105113
for _, t := range tools {
106-
if t.BackendName == target.BackendName && t.ToolName == target.ToolName {
114+
if t.BackendName != target.BackendName || t.ToolName != target.ToolName {
115+
continue
116+
}
117+
if len(t.Arguments) == 0 {
118+
return true
119+
}
120+
if args == nil {
121+
return false
122+
}
123+
allMatch := true
124+
for key, pattern := range t.Arguments {
125+
rawVal, ok := args[key]
126+
if !ok {
127+
allMatch = false
128+
break
129+
}
130+
re, err := regexp.Compile(pattern)
131+
if err != nil {
132+
allMatch = false
133+
break
134+
}
135+
var data []byte
136+
if s, ok := rawVal.(string); ok {
137+
data = []byte(s)
138+
} else {
139+
jsonVal, err := json.Marshal(rawVal)
140+
if err != nil {
141+
allMatch = false
142+
break
143+
}
144+
data = jsonVal
145+
}
146+
if !re.Match(data) {
147+
allMatch = false
148+
break
149+
}
150+
}
151+
if allMatch {
107152
return true
108153
}
109154
}
155+
// If no matching tool entry or no arguments matched, fail.
110156
return false
111157
}
112158

‎internal/mcpproxy/authorization_test.go‎

Lines changed: 152 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ func TestAuthorizeRequest(t *testing.T) {
3636
header string
3737
backendName string
3838
toolName string
39+
args map[string]any
3940
expectAllowed bool
4041
}{
4142
{
@@ -59,6 +60,156 @@ func TestAuthorizeRequest(t *testing.T) {
5960
toolName: "tool1",
6061
expectAllowed: true,
6162
},
63+
{
64+
name: "matching tool scope and arguments regex allowed",
65+
auth: &filterapi.MCPRouteAuthorization{
66+
DefaultAction: filterapi.AuthorizationActionDeny,
67+
Rules: []filterapi.MCPRouteAuthorizationRule{
68+
{
69+
Source: filterapi.MCPAuthorizationSource{
70+
JWTSource: filterapi.JWTSource{Scopes: []string{"read"}},
71+
},
72+
Target: filterapi.MCPAuthorizationTarget{
73+
Tools: []filterapi.ToolCall{{
74+
BackendName: "backend1",
75+
ToolName: "tool1",
76+
Arguments: map[string]string{
77+
"mode": "fast|slow",
78+
"user": "u-[0-9]+",
79+
"debug": "true",
80+
},
81+
}},
82+
},
83+
Action: filterapi.AuthorizationActionAllow,
84+
},
85+
},
86+
},
87+
header: "Bearer " + makeToken("read"),
88+
backendName: "backend1",
89+
toolName: "tool1",
90+
args: map[string]any{
91+
"mode": "fast",
92+
"user": "u-123",
93+
"debug": "true",
94+
},
95+
expectAllowed: true,
96+
},
97+
{
98+
name: "argument regex mismatch denied",
99+
auth: &filterapi.MCPRouteAuthorization{
100+
DefaultAction: filterapi.AuthorizationActionDeny,
101+
Rules: []filterapi.MCPRouteAuthorizationRule{
102+
{
103+
Source: filterapi.MCPAuthorizationSource{
104+
JWTSource: filterapi.JWTSource{Scopes: []string{"read"}},
105+
},
106+
Target: filterapi.MCPAuthorizationTarget{
107+
Tools: []filterapi.ToolCall{{
108+
BackendName: "backend1",
109+
ToolName: "tool1",
110+
Arguments: map[string]string{
111+
"mode": "fast|slow",
112+
},
113+
}},
114+
},
115+
Action: filterapi.AuthorizationActionAllow,
116+
},
117+
},
118+
},
119+
header: "Bearer " + makeToken("read"),
120+
backendName: "backend1",
121+
toolName: "tool1",
122+
args: map[string]any{
123+
"mode": "other",
124+
},
125+
expectAllowed: false,
126+
},
127+
{
128+
name: "missing argument denies when required",
129+
auth: &filterapi.MCPRouteAuthorization{
130+
DefaultAction: filterapi.AuthorizationActionDeny,
131+
Rules: []filterapi.MCPRouteAuthorizationRule{
132+
{
133+
Source: filterapi.MCPAuthorizationSource{
134+
JWTSource: filterapi.JWTSource{Scopes: []string{"read"}},
135+
},
136+
Target: filterapi.MCPAuthorizationTarget{
137+
Tools: []filterapi.ToolCall{{
138+
BackendName: "backend1",
139+
ToolName: "tool1",
140+
Arguments: map[string]string{
141+
"mode": "fast",
142+
},
143+
}},
144+
},
145+
Action: filterapi.AuthorizationActionAllow,
146+
},
147+
},
148+
},
149+
header: "Bearer " + makeToken("read"),
150+
backendName: "backend1",
151+
toolName: "tool1",
152+
args: map[string]any{},
153+
expectAllowed: false,
154+
},
155+
{
156+
name: "numeric argument matches via JSON string",
157+
auth: &filterapi.MCPRouteAuthorization{
158+
DefaultAction: filterapi.AuthorizationActionDeny,
159+
Rules: []filterapi.MCPRouteAuthorizationRule{
160+
{
161+
Source: filterapi.MCPAuthorizationSource{
162+
JWTSource: filterapi.JWTSource{Scopes: []string{"read"}},
163+
},
164+
Target: filterapi.MCPAuthorizationTarget{
165+
Tools: []filterapi.ToolCall{{
166+
BackendName: "backend1",
167+
ToolName: "tool1",
168+
Arguments: map[string]string{
169+
"count": "^4[0-9]$",
170+
},
171+
}},
172+
},
173+
Action: filterapi.AuthorizationActionAllow,
174+
},
175+
},
176+
},
177+
header: "Bearer " + makeToken("read"),
178+
backendName: "backend1",
179+
toolName: "tool1",
180+
args: map[string]any{"count": 42},
181+
expectAllowed: true,
182+
},
183+
{
184+
name: "object argument can be matched via JSON string",
185+
auth: &filterapi.MCPRouteAuthorization{
186+
DefaultAction: filterapi.AuthorizationActionDeny,
187+
Rules: []filterapi.MCPRouteAuthorizationRule{
188+
{
189+
Source: filterapi.MCPAuthorizationSource{
190+
JWTSource: filterapi.JWTSource{Scopes: []string{"read"}},
191+
},
192+
Target: filterapi.MCPAuthorizationTarget{
193+
Tools: []filterapi.ToolCall{{
194+
BackendName: "backend1",
195+
ToolName: "tool1",
196+
Arguments: map[string]string{
197+
"payload": `"kind":"test"`,
198+
},
199+
}},
200+
},
201+
Action: filterapi.AuthorizationActionAllow,
202+
},
203+
},
204+
},
205+
header: "Bearer " + makeToken("read"),
206+
backendName: "backend1",
207+
toolName: "tool1",
208+
args: map[string]any{
209+
"payload": map[string]any{"kind": "test", "value": 123},
210+
},
211+
expectAllowed: true,
212+
},
62213
{
63214
name: "matching tool but insufficient scopes not allowed",
64215
auth: &filterapi.MCPRouteAuthorization{
@@ -294,7 +445,7 @@ func TestAuthorizeRequest(t *testing.T) {
294445
if tt.header != "" {
295446
headers.Set("Authorization", tt.header)
296447
}
297-
allowed := proxy.authorizeRequest(tt.auth, headers, tt.backendName, tt.toolName)
448+
allowed := proxy.authorizeRequest(tt.auth, headers, tt.backendName, tt.toolName, tt.args)
298449
if allowed != tt.expectAllowed {
299450
t.Fatalf("expected %v, got %v", tt.expectAllowed, allowed)
300451
}

‎internal/mcpproxy/handlers.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -542,7 +542,7 @@ func (m *MCPProxy) handleToolCallRequest(ctx context.Context, s *session, w http
542542

543543
// Enforce authentication if required by the route.
544544
if route.authorization != nil {
545-
if !m.authorizeRequest(route.authorization, headers, backendName, toolName) {
545+
if !m.authorizeRequest(route.authorization, headers, backendName, toolName, p.Arguments) {
546546
onErrorResponse(w, http.StatusUnauthorized, "authorization failed")
547547
return fmt.Errorf("authorization failed")
548548
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -642,6 +642,14 @@ spec:
642642
rule applies to.
643643
items:
644644
properties:
645+
arguments:
646+
additionalProperties:
647+
type: string
648+
description: |-
649+
Arguments defines the arguments that must be present in the tool call for this rule to match.
650+
Keys must exist and their values must match the provided RE2-compatible regular expressions.
651+
If the argument is a non-string type, it will be matched against its JSON representation.
652+
type: object
645653
backendName:
646654
description: BackendName is the name of the
647655
backend this tool belongs to.

‎site/docs/api/api.mdx‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1985,6 +1985,11 @@ References:
19851985
type="string"
19861986
required="true"
19871987
description="ToolName is the name of the tool."
1988+
/><ApiField
1989+
name="arguments"
1990+
type="object (keys:string, values:string)"
1991+
required="false"
1992+
description="Arguments defines the arguments that must be present in the tool call for this rule to match.<br />Keys must exist and their values must match the provided RE2-compatible regular expressions.<br />If the argument is a non-string type, it will be matched against its JSON representation."
19881993
/>
19891994

19901995

0 commit comments

Comments
 (0)