SQL: Limit how much space some string functions can use (#107333) (#107632)

* SQL: Limit how much space some string functions can use (#107333)

This will check and fail if certain functions would generate a result
exceeding a certain fixed byte size.

This prevents an operation/query to fail the entire VM.

(cherry picked from commit f1bcb33)

* Fix use of forbidden APIs

* Style

Author: bpintea

docs/changelog/107333.yaml

@@ -0,0 +1,18 @@
+pr: 107333
+summary: Limit how much space some string functions can use
+area: SQL
+type: breaking
+issues: []
+breaking:
+  title: Limit how much space some string functions can use
+  area: REST API
+  details: "Before this change, some of the string functions could return a result\
+    \ of any arbitrary length, which could force the VM to allocate large chunks of\
+    \ memory or even make it exit. Any user with access to the SQL API can invoke\
+    \ these functions. This change introduces a limitation  of how much memory the\
+    \ result returned by a function call can consume. The functions affected by this\
+    \ change are: CONCAT, INSERT, REPEAT,  REPLACE and SPACE."
+  impact: "The affected functions used to return a result of any length. After this\
+    \ change, a result can no longer exceed 1MB in length. Note that this is a bytes\
+    \ length, the character count may be lower."
+  notable: false

docs/reference/sql/functions/string.asciidoc

@@ -109,6 +109,8 @@ CONCAT(
 
 *Description*: Returns a character string that is the result of concatenating `string_exp1` to `string_exp2`.
 
+The resulting string cannot exceed a byte length of 1 MB.
+
 [source, sql]
 --------------------------------------------------
 include-tagged::{sql-specs}/docs/docs.csv-spec[stringConcat]
@@ -137,6 +139,8 @@ INSERT(
 
 *Description*: Returns a string where `length` characters have been deleted from `source`, beginning at `start`, and where `replacement` has been inserted into `source`, beginning at `start`.
 
+The resulting string cannot exceed a byte length of 1 MB.
+
 [source, sql]
 --------------------------------------------------
 include-tagged::{sql-specs}/docs/docs.csv-spec[stringInsert]
@@ -330,6 +334,8 @@ REPEAT(
 
 *Description*: Returns a character string composed of `string_exp` repeated `count` times.
 
+The resulting string cannot exceed a byte length of 1 MB.
+
 [source, sql]
 --------------------------------------------------
 include-tagged::{sql-specs}/docs/docs.csv-spec[stringRepeat]
@@ -356,6 +362,8 @@ REPLACE(
 
 *Description*: Search `source` for occurrences of `pattern`, and replace with `replacement`.
 
+The resulting string cannot exceed a byte length of 1 MB.
+
 [source, sql]
 --------------------------------------------------
 include-tagged::{sql-specs}/docs/docs.csv-spec[stringReplace]
@@ -423,6 +431,8 @@ SPACE(count) <1>
 
 *Description*: Returns a character string consisting of `count` spaces.
 
+The resulting string cannot exceed a byte length of 1 MB.
+
 [source, sql]
 --------------------------------------------------
 include-tagged::{sql-specs}/docs/docs.csv-spec[stringSpace]

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/string/BinaryStringNumericProcessor.java

@@ -16,6 +16,8 @@
 import java.io.IOException;
 import java.util.function.BiFunction;
 
+import static org.elasticsearch.xpack.sql.expression.function.scalar.string.StringProcessor.checkResultLength;
+
 /**
  * Processor class covering string manipulating functions that have the first parameter as string,
  * second parameter as numeric and a string result.
@@ -42,7 +44,7 @@ public enum BinaryStringNumericOperation implements BiFunction<String, Number, S
             if (i <= 0) {
                 return null;
             }
-
+            checkResultLength(s.length() * c.longValue()); // mul is safe: c's checked by doProcess() to be within Integer's range
             StringBuilder sb = new StringBuilder(s.length() * i);
             for (int j = 0; j < i; j++) {
                 sb.append(s);

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/string/ConcatFunctionProcessor.java

@@ -16,6 +16,8 @@
 import java.io.IOException;
 import java.util.Objects;
 
+import static org.elasticsearch.xpack.sql.expression.function.scalar.string.StringProcessor.checkResultLength;
+
 public class ConcatFunctionProcessor extends BinaryProcessor {
 
     public static final String NAME = "scon";
@@ -62,7 +64,10 @@ public static Object process(Object source1, Object source2) {
             throw new SqlIllegalArgumentException("A string/char is required; received [{}]", source2);
         }
 
-        return source1.toString().concat(source2.toString());
+        String str1 = source1.toString();
+        String str2 = source2.toString();
+        checkResultLength(str1.length() + str2.length());
+        return str1.concat(str2);
     }
 
     @Override

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/string/InsertFunctionProcessor.java

@@ -15,6 +15,8 @@
 import java.io.IOException;
 import java.util.Objects;
 
+import static org.elasticsearch.xpack.sql.expression.function.scalar.string.StringProcessor.checkResultLength;
+
 public class InsertFunctionProcessor implements Processor {
 
     private final Processor input, start, length, replacement;
@@ -71,7 +73,9 @@ public static Object doProcess(Object input, Object start, Object length, Object
         StringBuilder sb = new StringBuilder(input.toString());
         String replString = (replacement.toString());
 
-        return sb.replace(realStart, realStart + ((Number) length).intValue(), replString).toString();
+        int cutLength = ((Number) length).intValue();
+        checkResultLength(sb.length() - cutLength + replString.length());
+        return sb.replace(realStart, realStart + cutLength, replString).toString();
     }
 
     @Override

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/string/ReplaceFunctionProcessor.java

@@ -58,11 +58,20 @@ public static Object doProcess(Object input, Object pattern, Object replacement)
             throw new SqlIllegalArgumentException("A string/char is required; received [{}]", replacement);
         }
 
-        return Strings.replace(
-            input instanceof Character ? input.toString() : (String) input,
-            pattern instanceof Character ? pattern.toString() : (String) pattern,
-            replacement instanceof Character ? replacement.toString() : (String) replacement
-        );
+        String inputStr = input instanceof Character ? input.toString() : (String) input;
+        String patternStr = pattern instanceof Character ? pattern.toString() : (String) pattern;
+        String replacementStr = replacement instanceof Character ? replacement.toString() : (String) replacement;
+        checkResultLength(inputStr, patternStr, replacementStr);
+        return Strings.replace(inputStr, patternStr, replacementStr);
+    }
+
+    private static void checkResultLength(String input, String pattern, String replacement) {
+        int patternLen = pattern.length();
+        long matches = 0;
+        for (int i = input.indexOf(pattern); i >= 0; i = input.indexOf(pattern, i + patternLen)) {
+            matches++;
+        }
+        StringProcessor.checkResultLength(input.length() + matches * (replacement.length() - patternLen));
     }
 
     @Override

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/string/StringProcessor.java

@@ -17,8 +17,12 @@
 import java.util.Locale;
 import java.util.function.Function;
 
+import static org.elasticsearch.common.unit.ByteSizeUnit.MB;
+
 public class StringProcessor implements Processor {
 
+    static final long MAX_RESULT_LENGTH = MB.toBytes(1);
+
     private interface StringFunction<R> {
         default R apply(Object o) {
             if ((o instanceof String || o instanceof Character) == false) {
@@ -60,6 +64,7 @@ public enum StringOperation {
             if (i < 0) {
                 return null;
             }
+            checkResultLength(n.longValue());
             char[] spaces = new char[i];
             char whitespace = ' ';
             Arrays.fill(spaces, whitespace);
@@ -125,6 +130,17 @@ StringOperation processor() {
         return processor;
     }
 
+    static void checkResultLength(long needed) {
+        if (needed > MAX_RESULT_LENGTH) {
+            throw new SqlIllegalArgumentException(
+                "Required result length [{}] exceeds implementation limit [{}] bytes",
+                needed,
+                MAX_RESULT_LENGTH
+            );
+        }
+
+    }
+
     @Override
     public boolean equals(Object obj) {
         if (obj == null || obj.getClass() != getClass()) {

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/string/BinaryStringNumericProcessorTests.java

@@ -17,6 +17,8 @@
 
 import static org.elasticsearch.xpack.ql.expression.function.scalar.FunctionTestUtils.l;
 import static org.elasticsearch.xpack.ql.tree.Source.EMPTY;
+import static org.elasticsearch.xpack.sql.expression.function.scalar.string.StringFunctionProcessorTests.maxResultLengthTest;
+import static org.elasticsearch.xpack.sql.expression.function.scalar.string.StringProcessor.MAX_RESULT_LENGTH;
 
 public class BinaryStringNumericProcessorTests extends AbstractWireSerializingTestCase<BinaryStringNumericProcessor> {
 
@@ -150,6 +152,10 @@ public void testRepeatFunctionWithEdgeCases() {
         assertNull(new Repeat(EMPTY, l("foo"), l(-1)).makePipe().asProcessor().process(null));
         assertNull(new Repeat(EMPTY, l("foo"), l(0)).makePipe().asProcessor().process(null));
         assertNull(new Repeat(EMPTY, l('f'), l(Integer.MIN_VALUE)).makePipe().asProcessor().process(null));
+        assertEquals(
+            MAX_RESULT_LENGTH,
+            new Repeat(EMPTY, l('f'), l(MAX_RESULT_LENGTH)).makePipe().asProcessor().process(null).toString().length()
+        );
     }
 
     public void testRepeatFunctionInputsValidation() {
@@ -182,5 +188,14 @@ public void testRepeatFunctionInputsValidation() {
             () -> new Repeat(EMPTY, l("foo"), l(1.0)).makePipe().asProcessor().process(null)
         );
         assertEquals("A fixed point number is required for [count]; received [java.lang.Double]", siae.getMessage());
+
+        maxResultLengthTest(
+            MAX_RESULT_LENGTH + 1,
+            () -> new Repeat(EMPTY, l("f"), l(MAX_RESULT_LENGTH + 1)).makePipe().asProcessor().process(null)
+        );
+
+        String str = "foo";
+        long count = (MAX_RESULT_LENGTH / str.length()) + 1;
+        maxResultLengthTest(count * str.length(), () -> new Repeat(EMPTY, l(str), l(count)).makePipe().asProcessor().process(null));
     }
 }

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/string/ConcatProcessorTests.java

@@ -16,6 +16,8 @@
 
 import static org.elasticsearch.xpack.ql.expression.function.scalar.FunctionTestUtils.l;
 import static org.elasticsearch.xpack.ql.tree.Source.EMPTY;
+import static org.elasticsearch.xpack.sql.expression.function.scalar.string.StringFunctionProcessorTests.maxResultLengthTest;
+import static org.elasticsearch.xpack.sql.expression.function.scalar.string.StringProcessor.MAX_RESULT_LENGTH;
 
 public class ConcatProcessorTests extends AbstractWireSerializingTestCase<ConcatFunctionProcessor> {
 
@@ -60,4 +62,19 @@ public void testConcatFunctionInputsValidation() {
         );
         assertEquals("A string/char is required; received [3]", siae.getMessage());
     }
+
+    public void testMaxResultLength() {
+        String str = repeat('a', (int) MAX_RESULT_LENGTH - 1);
+        assertEquals(MAX_RESULT_LENGTH, new Concat(EMPTY, l(str), l("b")).makePipe().asProcessor().process(null).toString().length());
+
+        maxResultLengthTest(MAX_RESULT_LENGTH + 1, () -> new Concat(EMPTY, l(str), l("bb")).makePipe().asProcessor().process(null));
+    }
+
+    public static String repeat(char ch, int count) {
+        StringBuilder sb = new StringBuilder(count);
+        for (int i = 0; i < count; i++) {
+            sb.append(ch);
+        }
+        return sb.toString();
+    }
 }

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/string/InsertProcessorTests.java

@@ -16,6 +16,9 @@
 
 import static org.elasticsearch.xpack.ql.expression.function.scalar.FunctionTestUtils.l;
 import static org.elasticsearch.xpack.ql.tree.Source.EMPTY;
+import static org.elasticsearch.xpack.sql.expression.function.scalar.string.ConcatProcessorTests.repeat;
+import static org.elasticsearch.xpack.sql.expression.function.scalar.string.StringFunctionProcessorTests.maxResultLengthTest;
+import static org.elasticsearch.xpack.sql.expression.function.scalar.string.StringProcessor.MAX_RESULT_LENGTH;
 
 public class InsertProcessorTests extends AbstractWireSerializingTestCase<InsertFunctionProcessor> {
 
@@ -117,5 +120,21 @@ public void testInsertInputsValidation() {
             () -> new Insert(EMPTY, l("foobar"), l(1), l((long) Integer.MAX_VALUE + 1), l("bar")).makePipe().asProcessor().process(null)
         );
         assertEquals("[length] out of the allowed range [0, 2147483647], received [2147483648]", siae.getMessage());
+
+        String str = repeat('a', (int) MAX_RESULT_LENGTH);
+        String replaceWith = "bar";
+        assertEquals(
+            MAX_RESULT_LENGTH,
+            new Insert(EMPTY, l(str), l(1), l(replaceWith.length()), l(replaceWith)).makePipe()
+                .asProcessor()
+                .process(null)
+                .toString()
+                .length()
+        );
+
+        maxResultLengthTest(
+            MAX_RESULT_LENGTH + 1,
+            () -> new Insert(EMPTY, l(str), l(1), l(replaceWith.length() - 1), l(replaceWith)).makePipe().asProcessor().process(null)
+        );
     }
 }