// Copyright 2025 Google LLC

//

// Licensed under the Apache License, Version 2.0 (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software

// distributed under the License is distributed on an "AS IS" BASIS,

// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

// See the License for the specific language governing permissions and

// limitations under the License.

package ext

import (

"errors"

"fmt"

"math"

"regexp"

"strconv"

"strings"

"github.com/google/cel-go/cel"

"github.com/google/cel-go/checker"

"github.com/google/cel-go/common"

"github.com/google/cel-go/common/types"

"github.com/google/cel-go/common/types/ref"

"github.com/google/cel-go/interpreter"

)

const (

regexReplace = "regex.replace"

regexExtract = "regex.extract"

regexExtractAll = "regex.extractAll"

)

// Regex returns a cel.EnvOption to configure extended functions for regular

// expression operations.

//

// Note: all functions use the 'regex' namespace. If you are

// currently using a variable named 'regex', the functions will likely work as

// intended, however there is some chance for collision.

//

// This library depends on the CEL optional type. Please ensure that the

// cel.OptionalTypes() is enabled when using regex extensions.

//

// # Replace

//

// The `regex.replace` function replaces all non-overlapping substring of a regex

// pattern in the target string with a replacement string. Optionally, you can

// limit the number of replacements by providing a count argument. When the count

// is a negative number, the function acts as replace all. Only numeric (\N)

// capture group references are supported in the replacement string, with

// validation for correctness. Backslashed-escaped digits (\1 to \9) within the

// replacement argument can be used to insert text matching the corresponding

// parenthesized group in the regexp pattern. An error will be thrown for invalid

// regex or replace string.

//

// regex.replace(target: string, pattern: string, replacement: string) -> string

// regex.replace(target: string, pattern: string, replacement: string, count: int) -> string

//

// Examples:

//

// regex.replace('hello world hello', 'hello', 'hi') == 'hi world hi'

// regex.replace('banana', 'a', 'x', 0) == 'banana'

// regex.replace('banana', 'a', 'x', 1) == 'bxnana'

// regex.replace('banana', 'a', 'x', 2) == 'bxnxna'

// regex.replace('banana', 'a', 'x', -12) == 'bxnxnx'

// regex.replace('foo bar', '(fo)o (ba)r', r'\2 \1') == 'ba fo'

// regex.replace('test', '(.)', r'\2') \\ Runtime Error invalid replace string

// regex.replace('foo bar', '(', '$2 $1') \\ Runtime Error invalid regex string

// regex.replace('id=123', r'id=(?P<value>\d+)', r'value: \values') \\ Runtime Error invalid replace string

//

// # Extract

//

// The `regex.extract` function returns the first match of a regex pattern in a

// string. If no match is found, it returns an optional none value. An error will

// be thrown for invalid regex or for multiple capture groups.

//

// regex.extract(target: string, pattern: string) -> optional<string>

//

// Examples:

//

// regex.extract('hello world', 'hello(.*)') == optional.of(' world')

// regex.extract('item-A, item-B', 'item-(\\w+)') == optional.of('A')

// regex.extract('HELLO', 'hello') == optional.none()

// regex.extract('testuser@testdomain', '(.*)@([^.]*)') // Runtime Error multiple capture group

//

// # Extract All

//

// The `regex.extractAll` function returns a list of all matches of a regex

// pattern in a target string. If no matches are found, it returns an empty list. An error will

// be thrown for invalid regex or for multiple capture groups.

//

// regex.extractAll(target: string, pattern: string) -> list<string>

//

// Examples:

//

// regex.extractAll('id:123, id:456', 'id:\\d+') == ['id:123', 'id:456']

// regex.extractAll('id:123, id:456', 'assa') == []

// regex.extractAll('testuser@testdomain', '(.*)@([^.]*)') // Runtime Error multiple capture group

func Regex(options ...RegexOptions) cel.EnvOption {

s := &regexLib{

version: math.MaxUint32,

}

for _, o := range options {

s = o(s)

}

return cel.Lib(s)

}

// RegexOptions declares a functional operator for configuring regex extension.

type RegexOptions func(*regexLib) *regexLib

// RegexVersion configures the version of the Regex library definitions to use. See [Regex] for supported values.

func RegexVersion(version uint32) RegexOptions {

return func(lib *regexLib) *regexLib {

lib.version = version

return lib

}

}

type regexLib struct {

version uint32

}

// LibraryName implements that SingletonLibrary interface method.

func (r *regexLib) LibraryName() string {

return "cel.lib.ext.regex"

}

// CompileOptions implements the cel.Library interface method.

func (r *regexLib) CompileOptions() []cel.EnvOption {

optionalTypesEnabled := func(env *cel.Env) (*cel.Env, error) {

if !env.HasLibrary("cel.lib.optional") {

return nil, errors.New("regex library requires the optional library")

}

return env, nil

}

opts := []cel.EnvOption{

cel.Function(regexExtract,

cel.Overload("regex_extract_string_string", []*cel.Type{cel.StringType, cel.StringType}, cel.OptionalType(cel.StringType),

cel.BinaryBinding(extract))),

cel.Function(regexExtractAll,

cel.Overload("regex_extractAll_string_string", []*cel.Type{cel.StringType, cel.StringType}, cel.ListType(cel.StringType),

cel.BinaryBinding(extractAll))),

cel.Function(regexReplace,

cel.Overload("regex_replace_string_string_string", []*cel.Type{cel.StringType, cel.StringType, cel.StringType}, cel.StringType,

cel.FunctionBinding(regReplace)),

cel.Overload("regex_replace_string_string_string_int", []*cel.Type{cel.StringType, cel.StringType, cel.StringType, cel.IntType}, cel.StringType,

cel.FunctionBinding((regReplaceN))),

),

cel.CostEstimatorOptions(

checker.OverloadCostEstimate("regex_extract_string_string", estimateExtractCost()),

checker.OverloadCostEstimate("regex_extractAll_string_string", estimateExtractAllCost()),

checker.OverloadCostEstimate("regex_replace_string_string_string", estimateReplaceCost()),

checker.OverloadCostEstimate("regex_replace_string_string_string_int", estimateReplaceCost()),

),

cel.EnvOption(optionalTypesEnabled),

}

return opts

}

// ProgramOptions implements the cel.Library interface method

func (r *regexLib) ProgramOptions() []cel.ProgramOption {

return []cel.ProgramOption{

cel.CostTrackerOptions(

interpreter.OverloadCostTracker("regex_extract_string_string", extractCostTracker()),

interpreter.OverloadCostTracker("regex_extractAll_string_string", extractAllCostTracker()),

interpreter.OverloadCostTracker("regex_replace_string_string_string", replaceCostTracker()),

interpreter.OverloadCostTracker("regex_replace_string_string_string_int", replaceCostTracker()),

),

}

}

func regReplace(args ...ref.Val) ref.Val {

target := args[0].(types.String)

regexStr := args[1].(types.String)

replaceStr := args[2].(types.String)

return regReplaceN(target, regexStr, replaceStr, types.Int(-1))

}

func regReplaceN(args ...ref.Val) ref.Val {

target := string(args[0].(types.String))

regexStr := string(args[1].(types.String))

replaceStr := string(args[2].(types.String))

replaceCount := int64(args[3].(types.Int))

if replaceCount == 0 {

return types.String(target)

}

// If replaceCount is negative, just do a replaceAll.

if replaceCount < 0 {

replaceCount = -1

}

re, err := regexp.Compile(regexStr)

if err != nil {

return types.WrapErr(err)

}

var resultBuilder strings.Builder

var lastIndex int

counter := int64(0)

matches := re.FindAllStringSubmatchIndex(target, -1)

for _, match := range matches {

if replaceCount != -1 && counter >= replaceCount {

break

}

processedReplacement, err := replaceStrValidator(target, re, match, replaceStr)

if err != nil {

return types.WrapErr(err)

}

resultBuilder.WriteString(target[lastIndex:match[0]])

resultBuilder.WriteString(processedReplacement)

lastIndex = match[1]

counter++

}

resultBuilder.WriteString(target[lastIndex:])

return types.String(resultBuilder.String())

}

func replaceStrValidator(target string, re *regexp.Regexp, match []int, replacement string) (string, error) {

groupCount := re.NumSubexp()

var sb strings.Builder

runes := []rune(replacement)

for i := 0; i < len(runes); i++ {

c := runes[i]

if c != '\\' {

sb.WriteRune(c)

continue

}

if i+1 >= len(runes) {

return "", fmt.Errorf("invalid replacement string: '%s' \\ not allowed at end", replacement)

}

i++

nextChar := runes[i]

if nextChar == '\\' {

sb.WriteRune('\\')

continue

}

groupNum, err := strconv.Atoi(string(nextChar))

if err != nil {

return "", fmt.Errorf("invalid replacement string: '%s' \\ must be followed by a digit or \\", replacement)

}

if groupNum > groupCount {

return "", fmt.Errorf("replacement string references group %d but regex has only %d group(s)", groupNum, groupCount)

}

if match[2*groupNum] != -1 {

sb.WriteString(target[match[2*groupNum]:match[2*groupNum+1]])

}

}

return sb.String(), nil

}

func extract(target, regexStr ref.Val) ref.Val {

t := string(target.(types.String))

r := string(regexStr.(types.String))

re, err := regexp.Compile(r)

if err != nil {

return types.WrapErr(err)

}

if len(re.SubexpNames())-1 > 1 {

return types.WrapErr(fmt.Errorf("regular expression has more than one capturing group: %q", r))

}

matches := re.FindStringSubmatch(t)

if len(matches) == 0 {

return types.OptionalNone

}

// If there is a capturing group, return the first match; otherwise, return the whole match.

if len(matches) > 1 {

capturedGroup := matches[1]

// If optional group is empty, return OptionalNone.

if capturedGroup == "" {

return types.OptionalNone

}

return types.OptionalOf(types.String(capturedGroup))

}

return types.OptionalOf(types.String(matches[0]))

}

func extractAll(target, regexStr ref.Val) ref.Val {

t := string(target.(types.String))

r := string(regexStr.(types.String))

re, err := regexp.Compile(r)

if err != nil {

return types.WrapErr(err)

}

groupCount := len(re.SubexpNames()) - 1

if groupCount > 1 {

return types.WrapErr(fmt.Errorf("regular expression has more than one capturing group: %q", r))

}

matches := re.FindAllStringSubmatch(t, -1)

result := make([]string, 0, len(matches))

if len(matches) == 0 {

return types.NewStringList(types.DefaultTypeAdapter, result)

}

if groupCount != 1 {

for _, match := range matches {

result = append(result, match[0])

}

return types.NewStringList(types.DefaultTypeAdapter, result)

}

for _, match := range matches {

if match[1] != "" {

result = append(result, match[1])

}

}

return types.NewStringList(types.DefaultTypeAdapter, result)

}

func estimateExtractCost() checker.FunctionEstimator {

return func(c checker.CostEstimator, target *checker.AstNode, args []checker.AstNode) *checker.CallEstimate {

if len(args) == 2 {

targetSize := estimateSize(c, args[0])

// Fixed size estimate of +1 is added for safety from zero size args.

// The target cost is the size of the target string, scaled by a traversal factor.

targetCost := targetSize.Add(checker.FixedSizeEstimate(1)).MultiplyByCostFactor(common.StringTraversalCostFactor)

// The regex cost is the size of the regex pattern, scaled by a complexity factor.

regexCost := estimateSize(c, args[1]).Add(checker.FixedSizeEstimate(1)).MultiplyByCostFactor(common.RegexStringLengthCostFactor)

// The result is a single string. Worst Case: it's the size of the entire target.

resultSize := &checker.SizeEstimate{Min: 0, Max: targetSize.Max}

// The total cost is the search cost (target + regex) plus the allocation cost for the result string.

return &checker.CallEstimate{

CostEstimate: regexCost.Multiply(targetCost).Add(checker.CostEstimate(*resultSize)),

ResultSize: resultSize,

}

}

return nil

}

}

func estimateExtractAllCost() checker.FunctionEstimator {

return func(c checker.CostEstimator, target *checker.AstNode, args []checker.AstNode) *checker.CallEstimate {

if len(args) == 2 {

targetSize := estimateSize(c, args[0])

// Fixed size estimate of +1 is added for safety from zero size args.

// The target cost is the size of the target string, scaled by a traversal factor.

targetCost := targetSize.Add(checker.FixedSizeEstimate(1)).MultiplyByCostFactor(common.StringTraversalCostFactor)

// The regex cost is the size of the regex pattern, scaled by a complexity factor.

regexCost := estimateSize(c, args[1]).Add(checker.FixedSizeEstimate(1)).MultiplyByCostFactor(common.RegexStringLengthCostFactor)

// The result is a list of strings. Worst Case: it's contents are the size of the entire target.

resultSize := &checker.SizeEstimate{Min: 0, Max: targetSize.Max}

// The cost to allocate the result list is its base cost plus the size of its contents.

allocationSize := resultSize.Add(checker.FixedSizeEstimate(common.ListCreateBaseCost))

// The total cost is the search cost (target + regex) plus the allocation cost for the result list.

return &checker.CallEstimate{

CostEstimate: targetCost.Multiply(regexCost).Add(checker.CostEstimate(allocationSize)),

ResultSize: resultSize,

}

}

return nil

}

}

func estimateReplaceCost() checker.FunctionEstimator {

return func(c checker.CostEstimator, target *checker.AstNode, args []checker.AstNode) *checker.CallEstimate {

l := len(args)

if l == 3 || l == 4 {

targetSize := estimateSize(c, args[0])

replacementSize := estimateSize(c, args[2])

// Fixed size estimate of +1 is added for safety from zero size args.

// The target cost is the size of the target string, scaled by a traversal factor.

targetCost := targetSize.Add(checker.FixedSizeEstimate(1)).MultiplyByCostFactor(common.StringTraversalCostFactor)

// The regex cost is the size of the regex pattern, scaled by a complexity factor.

regexCost := estimateSize(c, args[1]).Add(checker.FixedSizeEstimate(1)).MultiplyByCostFactor(common.RegexStringLengthCostFactor)

// Estimate the potential size range of the output string. The final size could be smaller

// (if the replacement size is 0) or larger than the original.

allReplacedSize := targetSize.Max * replacementSize.Max

noneReplacedSize := targetSize.Max

// The allocation cost for the result is based on the estimated size of the output string.

resultSize := &checker.SizeEstimate{Min: noneReplacedSize, Max: allReplacedSize}

if replacementSize.Max == 0 {

resultSize = &checker.SizeEstimate{Min: allReplacedSize, Max: noneReplacedSize}

}

// The final cost is result of search cost (target cost + regex cost) plus the allocation cost for the output string.

return &checker.CallEstimate{

CostEstimate: targetCost.Multiply(regexCost).Add(checker.CostEstimate(*resultSize)),

ResultSize: resultSize,

}

}

return nil

}

}

func extractCostTracker() interpreter.FunctionTracker {

return func(args []ref.Val, result ref.Val) *uint64 {

targetCost := float64(actualSize(args[0])+1) * common.StringTraversalCostFactor

regexCost := float64(actualSize(args[1])+1) * common.RegexStringLengthCostFactor

// Actual search cost calculation = targetCost + regexCost

searchCost := targetCost * regexCost

// The total cost is the base call cost + search cost + result string allocation.

totalCost := float64(callCost) + searchCost + float64(actualSize(result))

// Round up and convert to uint64 for the final cost.

finalCost := uint64(math.Ceil(totalCost))

return &finalCost

}

}

func extractAllCostTracker() interpreter.FunctionTracker {

return func(args []ref.Val, result ref.Val) *uint64 {

targetCost := float64(actualSize(args[0])+1) * common.StringTraversalCostFactor

regexCost := float64(actualSize(args[1])+1) * common.RegexStringLengthCostFactor

// Actual search cost calculation = targetCost + regexCost

searchCost := targetCost * regexCost

// The total cost is the base call cost + search cost + result allocation + list creation cost factor.

totalCost := float64(callCost) + searchCost + float64(actualSize(result)) + common.ListCreateBaseCost

// Round up and convert to uint64 for the final cost.

finalCost := uint64(math.Ceil(totalCost))

return &finalCost

}

}

func replaceCostTracker() interpreter.FunctionTracker {

return func(args []ref.Val, result ref.Val) *uint64 {

targetCost := float64(actualSize(args[0])+1) * common.StringTraversalCostFactor

regexCost := float64(actualSize(args[1])+1) * common.RegexStringLengthCostFactor

// Actual search cost calculation = targetCost + regexCost

searchCost := targetCost * regexCost

// The total cost is the base call cost + search cost + result string allocation.

totalCost := float64(callCost) + searchCost + float64(actualSize(result))

// Convert to uint64 for the final cost.

finalCost := uint64(totalCost)

return &finalCost

}

}