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.