Comprehension limit validator (#769)

cel/validator.go

@@ -92,6 +92,18 @@ func ValidateHomogeneousAggregateLiterals() ASTValidator {
 	return homogeneousAggregateLiteralValidator{}
 }
 
+// ValidateComprehensionNestingLimit ensures that comprehension nesting does not exceed the specified limit.
+//
+// This validator can be useful for preventing arbitrarily nested comprehensions which can take high polynomial
+// time to complete.
+//
+// Note, this limit does not apply to comprehensions with an empty iteration range, as these comprehensions have
+// no actual looping cost. The cel.bind() utilizes the comprehension structure to perform local variable
+// assignments and supplies an empty iteration range, so they won't count against the nesting limit either.
+func ValidateComprehensionNestingLimit(limit int) ASTValidator {
+	return nestingLimitValidator{limit: limit}
+}
+
 type argChecker func(env *Env, call, arg ast.NavigableExpr) error
 
 func newFormatValidator(funcName string, argNum int, check argChecker) formatValidator {
@@ -248,6 +260,51 @@ func (homogeneousAggregateLiteralValidator) typeMismatch(errs errorReporter, id
 
 func (homogeneousAggregateLiteralValidator) is_validator() {}
 
+type nestingLimitValidator struct {
+	limit int
+}
+
+func (v nestingLimitValidator) Name() string {
+	return "cel.lib.std.validate.comprehension_nesting_limit"
+}
+
+func (v nestingLimitValidator) Validate(e *Env, a *Ast, iss *Issues) {
+	errs := errorReporter{iss: iss, info: a.info}
+	root := ast.NavigateCheckedAST(astToCheckedAST(a))
+	comprehensions := ast.MatchDescendants(root, ast.KindMatcher(ast.ComprehensionKind))
+	if len(comprehensions) <= v.limit {
+		return
+	}
+	for _, comp := range comprehensions {
+		count := 0
+		e := comp
+		hasParent := true
+		for hasParent {
+			// When the expression is not a comprehension, continue to the next ancestor.
+			if e.Kind() != ast.ComprehensionKind {
+				e, hasParent = e.Parent()
+				continue
+			}
+			// When the comprehension has an empty range, continue to the next ancestor
+			// as this comprehension does not have any associated cost.
+			iterRange := e.AsComprehension().IterRange()
+			if iterRange.Kind() == ast.ListKind && iterRange.AsList().Size() == 0 {
+				e, hasParent = e.Parent()
+				continue
+			}
+			// Otherwise check the nesting limit.
+			count++
+			if count > v.limit {
+				errs.reportErrorAtID(comp.ID(), "comprehension exceeds nesting limit")
+				break
+			}
+			e, hasParent = e.Parent()
+		}
+	}
+}
+
+func (nestingLimitValidator) is_validator() {}
+
 type errorReporter struct {
 	iss  *Issues
 	info *exprpb.SourceInfo

cel/validator_test.go

@@ -289,6 +289,61 @@ func TestValidateHomogeneousAggregateLiterals(t *testing.T) {
 	}
 }
 
+func TestValidateComprehensionNestingLimit(t *testing.T) {
+	env, err := NewEnv(
+		ASTValidators(ValidateComprehensionNestingLimit(2)),
+	)
+	if err != nil {
+		t.Fatalf("NewEnv() failed: %v", err)
+	}
+	tests := []struct {
+		expr string
+		iss  string
+	}{
+		{
+			expr: `[1, 2, 3].exists(i, i < 1)`,
+		},
+		{
+			expr: `[1, 2, 3].exists(i, [4, 5, 6].filter(j, j % i != 0).size() > 0)`,
+		},
+		{
+			// three comprehensions, but not three levels deep
+			expr: `[1, 2, 3].exists(i, [4, 5, 6].filter(j, j % i != 0).size() > 0) && [1, 2, 3].exists(i, i < 1)`,
+		},
+		{
+			// the empty iteration range in [].all(k, k) does not impact the actual runtime complexity,
+			// so it does not trip the comprehension limit.
+			expr: `[1, 2, 3].exists(i, [4, 5, 6].filter(j, [].all(k, k) && j % i != 0).size() > 0)`,
+		},
+		{
+			// three comprehensions, three levels deep
+			expr: `[1, 2, 3].map(i, [4, 5, 6].map(j, [7, 8, 9].map(k, i * j * k)))`,
+			iss: `
+			ERROR: <input>:1:48: comprehension exceeds nesting limit
+             | [1, 2, 3].map(i, [4, 5, 6].map(j, [7, 8, 9].map(k, i * j * k)))
+             | ...............................................^`,
+		},
+	}
+	for _, tst := range tests {
+		tc := tst
+		t.Run(tc.expr, func(t *testing.T) {
+			_, iss := env.Compile(tc.expr)
+			if tc.iss != "" {
+				if iss.Err() == nil {
+					t.Fatalf("e.Compile(%v) returned ast, expected error: %v", tc.expr, tc.iss)
+				}
+				if !test.Compare(iss.Err().Error(), tc.iss) {
+					t.Fatalf("e.Compile(%v) returned %v, expected error: %v", tc.expr, iss.Err(), tc.iss)
+				}
+				return
+			}
+			if iss.Err() != nil {
+				t.Fatalf("e.Compile(%v) failed: %v", tc.expr, iss.Err())
+			}
+		})
+	}
+}
+
 func TestExtendedValidations(t *testing.T) {
 	env, err := NewEnv(
 		Variable("x", types.StringType),