use CEL for arguments matching

Signed-off-by: Huabing Zhao <[email protected]>

Author: zhaohuabing

api/v1alpha1/mcp_route.go

@@ -305,12 +305,14 @@ type ToolCall struct {
 	// +kubebuilder:validation:Required
 	ToolName string `json:"toolName"`
 
-	// Arguments defines the arguments that must be present in the tool call for this rule to match.
-	// Keys must exist and their values must match the provided RE2-compatible regular expressions.
-	// If the argument is a non-string type, it will be matched against its JSON representation.
+	// Arguments is a CEL expression that must evaluate to true for the rule to match.
+	// The expression is evaluated with a single variable "args" bound to the tool call arguments as a dynamic object.
+	// Guard against missing fields with null checks (e.g., args["foo"] != null && args["foo"]["bar"] == "val").
 	//
+	// +kubebuilder:validation:Optional
+	// +kubebuilder:validation:MaxLength=4096
 	// +optional
-	Arguments map[string]string `json:"arguments,omitempty"`
+	Arguments *string `json:"arguments,omitempty"`
 }
 
 // JWKS defines how to obtain JSON Web Key Sets (JWKS) either from a remote HTTP/HTTPS endpoint or from a local source.

api/v1alpha1/zz_generated.deepcopy.go

@@ -108,7 +108,7 @@ type ToolCall struct {
 	// ToolName is the name of the tool.
 	ToolName string `json:"toolName"`
 
-	// Arguments defines required arguments (exact key names with regex patterns for values).
-	// All patterns must match for the rule to apply.
-	Arguments map[string]string `json:"arguments,omitempty"`
+	// Arguments is a CEL expression evaluated against the tool call arguments map.
+	// The expression must evaluate to true for the rule to apply.
+	Arguments *string `json:"arguments,omitempty"`
 }

internal/filterapi/mcpconfig.go

@@ -6,23 +6,88 @@
 package mcpproxy
 
 import (
-	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
 	"net/http"
-	"regexp"
 	"strings"
 
 	"github.com/golang-jwt/jwt/v5"
+	"github.com/google/cel-go/cel"
+	"github.com/google/cel-go/common/types"
 	"k8s.io/apimachinery/pkg/util/sets"
 
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
 )
 
+type compiledAuthorization struct {
+	ResourceMetadataURL string
+	Rules               []compiledAuthorizationRule
+}
+
+type compiledAuthorizationRule struct {
+	Source filterapi.MCPAuthorizationSource
+	Target []compiledToolCall
+}
+
+type compiledToolCall struct {
+	BackendName string
+	ToolName    string
+	Expression  string
+	program     cel.Program
+}
+
+// compileAuthorization compiles the MCPRouteAuthorization into a compiledAuthorization for efficient CEL evaluation.
+func compileAuthorization(auth *filterapi.MCPRouteAuthorization) (*compiledAuthorization, error) {
+	if auth == nil {
+		return nil, nil
+	}
+
+	env, err := cel.NewEnv(
+		cel.Variable("args", cel.DynType),
+		cel.OptionalTypes(),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create CEL environment: %w", err)
+	}
+
+	compiled := &compiledAuthorization{
+		ResourceMetadataURL: auth.ResourceMetadataURL,
+	}
+
+	for _, rule := range auth.Rules {
+		cr := compiledAuthorizationRule{
+			Source: rule.Source,
+		}
+		for _, tool := range rule.Target.Tools {
+			ct := compiledToolCall{
+				BackendName: tool.BackendName,
+				ToolName:    tool.ToolName,
+			}
+			if tool.Arguments != nil && strings.TrimSpace(*tool.Arguments) != "" {
+				expr := strings.TrimSpace(*tool.Arguments)
+				ast, issues := env.Compile(expr)
+				if issues != nil && issues.Err() != nil {
+					return nil, fmt.Errorf("failed to compile arguments CEL for tool %s/%s: %w", tool.BackendName, tool.ToolName, issues.Err())
+				}
+				program, err := env.Program(ast, cel.CostLimit(10000))
+				if err != nil {
+					return nil, fmt.Errorf("failed to build arguments CEL program for tool %s/%s: %w", tool.BackendName, tool.ToolName, err)
+				}
+				ct.Expression = expr
+				ct.program = program
+			}
+			cr.Target = append(cr.Target, ct)
+		}
+		compiled.Rules = append(compiled.Rules, cr)
+	}
+
+	return compiled, nil
+}
+
 // authorizeRequest authorizes the request based on the given MCPRouteAuthorization configuration.
 
-func (m *MCPProxy) authorizeRequest(authorization *filterapi.MCPRouteAuthorization, headers http.Header, backendName, toolName string, argments any) (bool, []string) {
+func (m *MCPProxy) authorizeRequest(authorization *compiledAuthorization, headers http.Header, backendName, toolName string, arguments any) (bool, []string) {
 	if authorization == nil {
 		return true, nil
 	}
@@ -48,17 +113,11 @@ func (m *MCPProxy) authorizeRequest(authorization *filterapi.MCPRouteAuthorizati
 		return false, nil
 	}
 
-	scopeSet := sets.New[string](extractScopes(claims)...)
+	scopeSet := sets.New(extractScopes(claims)...)
 	var missingScopesForChallenge []string
 
 	for _, rule := range authorization.Rules {
-		var args map[string]any
-		if argments != nil {
-			if cast, ok := argments.(map[string]any); ok {
-				args = cast
-			}
-		}
-		if !m.toolMatches(filterapi.ToolCall{BackendName: backendName, ToolName: toolName}, rule.Target.Tools, args) {
+		if !m.toolMatches(backendName, toolName, rule.Target, arguments) {
 			continue
 		}
 
@@ -117,53 +176,36 @@ func extractScopes(claims jwt.MapClaims) []string {
 	}
 }
 
-func (m *MCPProxy) toolMatches(target filterapi.ToolCall, tools []filterapi.ToolCall, args map[string]any) bool {
+func (m *MCPProxy) toolMatches(backendName, toolName string, tools []compiledToolCall, args any) bool {
 	if len(tools) == 0 {
 		return true
 	}
 
 	for _, t := range tools {
-		if t.BackendName != target.BackendName || t.ToolName != target.ToolName {
+		if t.BackendName != backendName || t.ToolName != toolName {
 			continue
 		}
-		if len(t.Arguments) == 0 {
+		if t.program == nil {
 			return true
 		}
-		if args == nil {
-			return false
+
+		result, _, err := t.program.Eval(map[string]any{"args": args})
+		if err != nil {
+			m.l.Error("failed to evaluate arguments CEL", slog.String("backend", t.BackendName), slog.String("tool", t.ToolName), slog.String("error", err.Error()))
+			continue
 		}
-		allMatch := true
-		for key, pattern := range t.Arguments {
-			rawVal, ok := args[key]
-			if !ok {
-				allMatch = false
-				break
-			}
-			re, err := regexp.Compile(pattern)
-			if err != nil {
-				m.l.Error("invalid argument regex pattern", slog.String("pattern", pattern), slog.String("error", err.Error()))
-				allMatch = false
-				break
-			}
-			var data []byte
-			if s, ok := rawVal.(string); ok {
-				data = []byte(s)
-			} else {
-				jsonVal, err := json.Marshal(rawVal)
-				if err != nil {
-					m.l.Error("failed to marshal argument value to json", slog.String("key", key), slog.String("error", err.Error()))
-					allMatch = false
-					break
-				}
-				data = jsonVal
+
+		switch v := result.Value().(type) {
+		case bool:
+			if v {
+				return true
 			}
-			if !re.Match(data) {
-				allMatch = false
-				break
+		case types.Bool:
+			if bool(v) {
+				return true
 			}
-		}
-		if allMatch {
-			return true
+		default:
+			m.l.Error("arguments CEL did not return a boolean", slog.String("backend", t.BackendName), slog.String("tool", t.ToolName), slog.String("expression", t.Expression))
 		}
 	}
 	// If no matching tool entry or no arguments matched, fail.

internal/mcpproxy/authorization.go

@@ -17,6 +17,8 @@ import (
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
 )
 
+func strPtr(s string) *string { return &s }
+
 func TestAuthorizeRequest(t *testing.T) {
 	makeToken := func(scopes ...string) string {
 		claims := jwt.MapClaims{}
@@ -62,7 +64,7 @@ func TestAuthorizeRequest(t *testing.T) {
 			expectScopes:  nil,
 		},
 		{
-			name: "matching tool scope and arguments regex",
+			name: "matching tool scope and arguments CEL",
 			auth: &filterapi.MCPRouteAuthorization{
 				Rules: []filterapi.MCPRouteAuthorizationRule{
 					{
@@ -73,11 +75,7 @@ func TestAuthorizeRequest(t *testing.T) {
 							Tools: []filterapi.ToolCall{{
 								BackendName: "backend1",
 								ToolName:    "tool1",
-								Arguments: map[string]string{
-									"mode":  "fast|slow",
-									"user":  "u-[0-9]+",
-									"debug": "true",
-								},
+								Arguments:   strPtr(`args.mode in ["fast", "slow"] && args.user.matches("u-[0-9]+") && args.debug == true`),
 							}},
 						},
 					},
@@ -89,13 +87,13 @@ func TestAuthorizeRequest(t *testing.T) {
 			args: map[string]any{
 				"mode":  "fast",
 				"user":  "u-123",
-				"debug": "true",
+				"debug": true,
 			},
 			expectAllowed: true,
 			expectScopes:  nil,
 		},
 		{
-			name: "numeric argument matches via JSON string",
+			name: "numeric argument matches via CEL",
 			auth: &filterapi.MCPRouteAuthorization{
 				Rules: []filterapi.MCPRouteAuthorizationRule{
 					{
@@ -106,9 +104,7 @@ func TestAuthorizeRequest(t *testing.T) {
 							Tools: []filterapi.ToolCall{{
 								BackendName: "backend1",
 								ToolName:    "tool1",
-								Arguments: map[string]string{
-									"count": "^4[0-9]$",
-								},
+								Arguments:   strPtr(`int(args.count) >= 40 && int(args.count) < 50`),
 							}},
 						},
 					},
@@ -122,7 +118,7 @@ func TestAuthorizeRequest(t *testing.T) {
 			expectScopes:  nil,
 		},
 		{
-			name: "object argument can be matched via JSON string",
+			name: "object argument can be matched via CEL safe navigation",
 			auth: &filterapi.MCPRouteAuthorization{
 				Rules: []filterapi.MCPRouteAuthorizationRule{
 					{
@@ -133,9 +129,7 @@ func TestAuthorizeRequest(t *testing.T) {
 							Tools: []filterapi.ToolCall{{
 								BackendName: "backend1",
 								ToolName:    "tool1",
-								Arguments: map[string]string{
-									"payload": `"kind":"test"`,
-								},
+								Arguments:   strPtr(`args["payload"] != null && args["payload"]["kind"] == "test" && args["payload"]["value"] == 123`),
 							}},
 						},
 					},
@@ -145,12 +139,9 @@ func TestAuthorizeRequest(t *testing.T) {
 			backendName: "backend1",
 			toolName:    "tool1",
 			args: map[string]any{
-				"payload": struct {
-					Kind  string `json:"kind"`
-					Value int    `json:"value"`
-				}{
-					Kind:  "test",
-					Value: 123,
+				"payload": map[string]any{
+					"kind":  "test",
+					"value": 123,
 				},
 			},
 			expectAllowed: true,
@@ -177,7 +168,7 @@ func TestAuthorizeRequest(t *testing.T) {
 			expectScopes:  []string{"read", "write"},
 		},
 		{
-			name: "argument regex mismatch denied",
+			name: "arguments CEL mismatch denied",
 			auth: &filterapi.MCPRouteAuthorization{
 				Rules: []filterapi.MCPRouteAuthorizationRule{
 					{
@@ -188,9 +179,7 @@ func TestAuthorizeRequest(t *testing.T) {
 							Tools: []filterapi.ToolCall{{
 								BackendName: "backend1",
 								ToolName:    "tool1",
-								Arguments: map[string]string{
-									"mode": "fast|slow",
-								},
+								Arguments:   strPtr(`args.mode in ["fast", "slow"]`),
 							}},
 						},
 					},
@@ -217,9 +206,7 @@ func TestAuthorizeRequest(t *testing.T) {
 							Tools: []filterapi.ToolCall{{
 								BackendName: "backend1",
 								ToolName:    "tool1",
-								Arguments: map[string]string{
-									"mode": "fast",
-								},
+								Arguments:   strPtr(`args["mode"] == "fast"`),
 							}},
 						},
 					},
@@ -349,7 +336,11 @@ func TestAuthorizeRequest(t *testing.T) {
 			if tt.header != "" {
 				headers.Set("Authorization", tt.header)
 			}
-			allowed, requiredScopes := proxy.authorizeRequest(tt.auth, headers, tt.backendName, tt.toolName, tt.args)
+			compiled, err := compileAuthorization(tt.auth)
+			if err != nil {
+				t.Fatalf("failed to compile authorization: %v", err)
+			}
+			allowed, requiredScopes := proxy.authorizeRequest(compiled, headers, tt.backendName, tt.toolName, tt.args)
 			if allowed != tt.expectAllowed {
 				t.Fatalf("expected %v, got %v", tt.expectAllowed, allowed)
 			}
@@ -371,3 +362,25 @@ func TestBuildInsufficientScopeHeader(t *testing.T) {
 		}
 	})
 }
+
+func TestCompileAuthorizationInvalidExpression(t *testing.T) {
+	_, err := compileAuthorization(&filterapi.MCPRouteAuthorization{
+		Rules: []filterapi.MCPRouteAuthorizationRule{
+			{
+				Source: filterapi.MCPAuthorizationSource{
+					JWTSource: filterapi.JWTSource{Scopes: []string{"read"}},
+				},
+				Target: filterapi.MCPAuthorizationTarget{
+					Tools: []filterapi.ToolCall{{
+						BackendName: "backend1",
+						ToolName:    "tool1",
+						Arguments:   strPtr("args."),
+					}},
+				},
+			},
+		},
+	})
+	if err == nil {
+		t.Fatalf("expected compile error for invalid CEL expression")
+	}
+}