ext/lists: add max size check to genRange() to prevent OOM

ext/lists.go

@@ -153,15 +153,18 @@ var comparableTypes = []*cel.Type{
 //	].sortBy(e, e.score).map(e, e.name)
 //	== ["bar", "foo", "baz"]
 func Lists(options ...ListsOption) cel.EnvOption {
-	l := &listsLib{version: math.MaxUint32}
+	l := &listsLib{version: math.MaxUint32, maxRangeSize: defaultMaxRangeSize}
 	for _, o := range options {
 		l = o(l)
 	}
 	return cel.Lib(l)
 }
 
+const defaultMaxRangeSize = 10_000_000
+
 type listsLib struct {
-	version uint32
+	version      uint32
+	maxRangeSize int64
 }
 
 // LibraryName implements the SingletonLibrary interface method.
@@ -188,6 +191,16 @@ func ListsVersion(version uint32) ListsOption {
 	}
 }
 
+// ListsMaxRangeSize sets the maximum number of elements lists.range() will
+// allocate. If not set, the default is 10,000,000. Setting this to zero
+// disables the limit (not recommended).
+func ListsMaxRangeSize(size int64) ListsOption {
+	return func(lib *listsLib) *listsLib {
+		lib.maxRangeSize = size
+		return lib
+	}
+}
+
 // CompileOptions implements the Library interface method.
 func (lib listsLib) CompileOptions() []cel.EnvOption {
 	listType := cel.ListType(cel.TypeParamType("T"))
@@ -309,11 +322,12 @@ func (lib listsLib) CompileOptions() []cel.EnvOption {
 			)...,
 		))
 
+		maxRange := lib.maxRangeSize
 		opts = append(opts, cel.Function("lists.range",
 			cel.Overload("lists_range",
 				[]*cel.Type{cel.IntType}, cel.ListType(cel.IntType),
 				cel.UnaryBinding(func(n ref.Val) ref.Val {
-					result, err := genRange(n.(types.Int))
+					result, err := genRange(n.(types.Int), maxRange)
 					if err != nil {
 						return types.WrapErr(err)
 					}
@@ -403,8 +417,14 @@ func (lib *listsLib) ProgramOptions() []cel.ProgramOption {
 	return opts
 }
 
-func genRange(n types.Int) (ref.Val, error) {
-	var newList []ref.Val
+func genRange(n types.Int, maxSize int64) (ref.Val, error) {
+	if n < 0 {
+		return nil, fmt.Errorf("lists.range: size must be non-negative, got %d", n)
+	}
+	if maxSize > 0 && int64(n) > maxSize {
+		return nil, fmt.Errorf("lists.range: size %d exceeds maximum allowed (%d)", n, maxSize)
+	}
+	newList := make([]ref.Val, 0, n)
 	for i := types.Int(0); i < n; i++ {
 		newList = append(newList, i)
 	}

ext/lists_test.go

@@ -521,3 +521,35 @@ func testListsEnv(t *testing.T, opts ...cel.EnvOption) *cel.Env {
 	}
 	return env
 }
+
+func TestGenRangeMaxSize(t *testing.T) {
+	// Negative size should fail.
+	_, err := genRange(-1, defaultMaxRangeSize)
+	if err == nil {
+		t.Error("genRange(-1) should fail")
+	}
+
+	// Small size should work.
+	val, err := genRange(10, defaultMaxRangeSize)
+	if err != nil {
+		t.Fatalf("genRange(10) failed: %v", err)
+	}
+	if val == nil {
+		t.Fatal("genRange(10) returned nil")
+	}
+
+	// Over the limit should fail, not allocate.
+	_, err = genRange(defaultMaxRangeSize+1, defaultMaxRangeSize)
+	if err == nil {
+		t.Error("genRange(defaultMaxRangeSize+1) should fail")
+	}
+
+	// Zero limit disables the check.
+	val, err = genRange(100, 0)
+	if err != nil {
+		t.Fatalf("genRange(100, 0) with disabled limit failed: %v", err)
+	}
+	if val == nil {
+		t.Fatal("genRange(100, 0) returned nil")
+	}
+}