Skip to content

Commit

Permalink
Fix CEL cost handling of zero length replacement strings
Browse files Browse the repository at this point in the history
  • Loading branch information
jpbetz committed Aug 22, 2023
1 parent 70a2844 commit 87cf890
Show file tree
Hide file tree
Showing 3 changed files with 159 additions and 10 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -1683,17 +1683,19 @@ func TestCostEstimation(t *testing.T) {
name: "extended library replace",
schemaGenerator: func(max *int64) *schema.Structural {
strType := withMaxLength(primitiveType("string", ""), max)
beforeLen := int64(2)
afterLen := int64(4)
objType := objectType(map[string]schema.Structural{
"str": strType,
"before": strType,
"after": strType,
"before": withMaxLength(primitiveType("string", ""), &beforeLen),
"after": withMaxLength(primitiveType("string", ""), &afterLen),
})
objType = withRule(objType, "self.str.replace(self.before, self.after) == 'does not matter'")
return &objType
},
expectedCalcCost: 629152, // cost is based on the result size of the replace() call
setMaxElements: 10,
expectedSetCost: 14,
expectedCalcCost: 629154, // cost is based on the result size of the replace() call
setMaxElements: 4,
expectedSetCost: 12,
},
{
name: "extended library split",
Expand Down
44 changes: 39 additions & 5 deletions staging/src/k8s.io/apiserver/pkg/cel/library/cost.go
Original file line number Diff line number Diff line change
Expand Up @@ -126,13 +126,47 @@ func (l *CostEstimator) EstimateCallCost(function, overloadId string, target *ch
sz := l.sizeEstimate(*target)
toReplaceSz := l.sizeEstimate(args[0])
replaceWithSz := l.sizeEstimate(args[1])
// smallest possible result: smallest input size composed of the largest possible substrings being replaced by smallest possible replacement
minSz := uint64(math.Ceil(float64(sz.Min)/float64(toReplaceSz.Max))) * replaceWithSz.Min
// largest possible result: largest input size composed of the smallest possible substrings being replaced by largest possible replacement
maxSz := uint64(math.Ceil(float64(sz.Max)/float64(toReplaceSz.Min))) * replaceWithSz.Max

var replaceCount, retainedSz checker.SizeEstimate
// find the longest replacement:
if toReplaceSz.Min == 0 {
// if the string being replaced is empty, replace surrounds all characters in the input string with the replacement.
if replaceCount.Max < math.MaxUint64 {
replaceCount.Max = sz.Max + 1
}
// Include the length of the longest possible original string length.
retainedSz.Max = sz.Max
} else if replaceWithSz.Max <= toReplaceSz.Min {
// If the replacement does not make the result longer, use the original string length.
replaceCount.Max = 0
retainedSz.Max = sz.Max
} else {
// Replace the smallest possible substrings with the largest possible replacement
// as many times as possible.
replaceCount.Max = uint64(math.Ceil(float64(sz.Max) / float64(toReplaceSz.Min)))
}

// find the shortest replacement:
if toReplaceSz.Max == 0 {
// if the string being replaced is empty, replace surrounds all characters in the input string with the replacement.
if replaceCount.Min < math.MaxUint64 {
replaceCount.Min = sz.Min + 1
}
// Include the length of the shortest possible original string length.
retainedSz.Min = sz.Min
} else if toReplaceSz.Max <= replaceWithSz.Min {
// If the replacement does not make the result shorter, use the original string length.
replaceCount.Min = 0
retainedSz.Min = sz.Min
} else {
// Replace the largest possible substrings being with the smallest possible replacement
// as many times as possible.
replaceCount.Min = uint64(math.Ceil(float64(sz.Min) / float64(toReplaceSz.Max)))
}
size := replaceCount.Multiply(replaceWithSz).Add(retainedSz)

// cost is the traversal plus the construction of the result
return &checker.CallEstimate{CostEstimate: sz.MultiplyByCostFactor(2 * common.StringTraversalCostFactor), ResultSize: &checker.SizeEstimate{Min: minSz, Max: maxSz}}
return &checker.CallEstimate{CostEstimate: sz.MultiplyByCostFactor(2 * common.StringTraversalCostFactor), ResultSize: &size}
}
case "split":
if target != nil {
Expand Down
113 changes: 113 additions & 0 deletions staging/src/k8s.io/apiserver/pkg/cel/library/cost_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,18 @@ func TestStringLibrary(t *testing.T) {
expectEsimatedCost: checker.CostEstimate{Min: 3, Max: 3},
expectRuntimeCost: 3,
},
{
name: "replace between all chars",
expr: "'abc 123 def 123'.replace('', 'x')",
expectEsimatedCost: checker.CostEstimate{Min: 3, Max: 3},
expectRuntimeCost: 3,
},
{
name: "replace with empty",
expr: "'abc 123 def 123'.replace('123', '')",
expectEsimatedCost: checker.CostEstimate{Min: 3, Max: 3},
expectRuntimeCost: 3,
},
{
name: "replace with limit",
expr: "'abc 123 def 123'.replace('123', '456', 1)",
Expand Down Expand Up @@ -437,6 +449,107 @@ func testCost(t *testing.T, expr string, expectEsimatedCost checker.CostEstimate
}
}

func TestSize(t *testing.T) {
exactSize := func(size int) checker.SizeEstimate {
return checker.SizeEstimate{Min: uint64(size), Max: uint64(size)}
}
exactSizes := func(sizes ...int) []checker.SizeEstimate {
results := make([]checker.SizeEstimate, len(sizes))
for i, size := range sizes {
results[i] = exactSize(size)
}
return results
}
cases := []struct {
name string
function string
overload string
targetSize checker.SizeEstimate
argSizes []checker.SizeEstimate
expectSize checker.SizeEstimate
}{
{
name: "replace empty with char",
function: "replace",
targetSize: exactSize(3), // e.g. abc
argSizes: exactSizes(0, 1), // e.g. replace "" with "_"
expectSize: exactSize(7), // e.g. _a_b_c_
},
{
name: "maybe replace char with empty",
function: "replace",
targetSize: exactSize(3),
argSizes: exactSizes(1, 0),
expectSize: checker.SizeEstimate{Min: 0, Max: 3},
},
{
name: "maybe replace repeated",
function: "replace",
targetSize: exactSize(4),
argSizes: exactSizes(2, 4),
expectSize: checker.SizeEstimate{Min: 4, Max: 8},
},
{
name: "maybe replace empty",
function: "replace",
targetSize: exactSize(4),
argSizes: []checker.SizeEstimate{{Min: 0, Max: 1}, {Min: 0, Max: 2}},
expectSize: checker.SizeEstimate{Min: 0, Max: 14}, // len(__a__a__a__a__) == 14
},
{
name: "replace non-empty size range, maybe larger",
function: "replace",
targetSize: exactSize(4),
argSizes: []checker.SizeEstimate{{Min: 1, Max: 1}, {Min: 1, Max: 2}},
expectSize: checker.SizeEstimate{Min: 4, Max: 8},
},
{
name: "replace non-empty size range, maybe smaller",
function: "replace",
targetSize: exactSize(4),
argSizes: []checker.SizeEstimate{{Min: 1, Max: 2}, {Min: 1, Max: 1}},
expectSize: checker.SizeEstimate{Min: 2, Max: 4},
},
}
est := &CostEstimator{SizeEstimator: &testCostEstimator{}}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
var targetNode checker.AstNode = testSizeNode{size: tc.targetSize}
argNodes := make([]checker.AstNode, len(tc.argSizes))
for i, arg := range tc.argSizes {
argNodes[i] = testSizeNode{size: arg}
}
result := est.EstimateCallCost(tc.function, tc.overload, &targetNode, argNodes)
if result.ResultSize == nil {
t.Fatalf("Expected ResultSize but got none")
}
if *result.ResultSize != tc.expectSize {
t.Fatalf("Expected %+v but got %+v", tc.expectSize, *result.ResultSize)
}
})
}
}

type testSizeNode struct {
size checker.SizeEstimate
}

func (t testSizeNode) Path() []string {
return nil // not needed
}

func (t testSizeNode) Type() *expr.Type {
return nil // not needed
}

func (t testSizeNode) Expr() *expr.Expr {
return nil // not needed
}

func (t testSizeNode) ComputedSize() *checker.SizeEstimate {
return &t.size
}

type testCostEstimator struct {
}

Expand Down

0 comments on commit 87cf890

Please sign in to comment.