diff --git a/ext/lists.go b/ext/lists.go index 7a7a33c6..cbf7ce4e 100644 --- a/ext/lists.go +++ b/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) } diff --git a/ext/lists_test.go b/ext/lists_test.go index 42a2db6b..bc5b5a7e 100644 --- a/ext/lists_test.go +++ b/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") + } +}