Cap regex replacement output growth

Author: mgomes

vibes/builtins.go

@@ -465,20 +465,60 @@ func builtinRegexReplaceInternal(args []Value, kwargs map[string]Value, block Va
 	if len(text) > maxRegexInputBytes {
 		return NewNil(), fmt.Errorf("%s text exceeds limit %d bytes", method, maxRegexInputBytes)
 	}
+	if len(replacement) > maxRegexInputBytes {
+		return NewNil(), fmt.Errorf("%s replacement exceeds limit %d bytes", method, maxRegexInputBytes)
+	}
 
 	re, err := regexp.Compile(pattern)
 	if err != nil {
 		return NewNil(), fmt.Errorf("%s invalid regex: %v", method, err)
 	}
 
 	if replaceAll {
-		return NewString(re.ReplaceAllString(text, replacement)), nil
+		replaced, err := regexReplaceAllWithLimit(re, text, replacement, method)
+		if err != nil {
+			return NewNil(), err
+		}
+		return NewString(replaced), nil
 	}
 
 	loc := re.FindStringSubmatchIndex(text)
 	if loc == nil {
 		return NewString(text), nil
 	}
 	replaced := string(re.ExpandString(nil, replacement, text, loc))
+	outputLen := len(text) - (loc[1] - loc[0]) + len(replaced)
+	if outputLen > maxRegexInputBytes {
+		return NewNil(), fmt.Errorf("%s output exceeds limit %d bytes", method, maxRegexInputBytes)
+	}
 	return NewString(text[:loc[0]] + replaced + text[loc[1]:]), nil
 }
+
+func regexReplaceAllWithLimit(re *regexp.Regexp, text string, replacement string, method string) (string, error) {
+	matches := re.FindAllStringSubmatchIndex(text, -1)
+	if len(matches) == 0 {
+		return text, nil
+	}
+
+	out := make([]byte, 0, len(text))
+	last := 0
+	for _, loc := range matches {
+		segmentLen := loc[0] - last
+		if len(out) > maxRegexInputBytes-segmentLen {
+			return "", fmt.Errorf("%s output exceeds limit %d bytes", method, maxRegexInputBytes)
+		}
+		out = append(out, text[last:loc[0]]...)
+		out = re.ExpandString(out, replacement, text, loc)
+		if len(out) > maxRegexInputBytes {
+			return "", fmt.Errorf("%s output exceeds limit %d bytes", method, maxRegexInputBytes)
+		}
+		last = loc[1]
+	}
+
+	tailLen := len(text) - last
+	if len(out) > maxRegexInputBytes-tailLen {
+		return "", fmt.Errorf("%s output exceeds limit %d bytes", method, maxRegexInputBytes)
+	}
+	out = append(out, text[last:]...)
+	return string(out), nil
+}

vibes/runtime_test.go

@@ -2028,6 +2028,10 @@ func TestJSONAndRegexSizeGuards(t *testing.T) {
     def regex_match_guard(pattern, text)
       Regex.match(pattern, text)
     end
+
+    def regex_replace_all_guard(text, pattern, replacement)
+      Regex.replace_all(text, pattern, replacement)
+    end
     `)
 	if err != nil {
 		t.Fatalf("compile error: %v", err)
@@ -2058,6 +2062,17 @@ func TestJSONAndRegexSizeGuards(t *testing.T) {
 	if err == nil || !strings.Contains(err.Error(), "Regex.match text exceeds limit") {
 		t.Fatalf("expected Regex.match text guard error, got %v", err)
 	}
+
+	hugeReplacement := strings.Repeat("x", maxRegexInputBytes/2)
+	_, err = script.Call(
+		context.Background(),
+		"regex_replace_all_guard",
+		[]Value{NewString("abc"), NewString(""), NewString(hugeReplacement)},
+		CallOptions{},
+	)
+	if err == nil || !strings.Contains(err.Error(), "Regex.replace_all output exceeds limit") {
+		t.Fatalf("expected Regex.replace_all output guard error, got %v", err)
+	}
 }
 
 func TestLocaleSensitiveOperationsDeterministic(t *testing.T) {