ESQL: Limit how many bytes concat() can process (elastic#100360)

This adds a limit on the bytes length that a concatenated string can get to. The limit is set to 1MB (per entry). When the limit is hit, an exception is returned to the caller (similar to an accounting circuit breaking) and execution halted. Related: elastic#100288.
bpintea · Oct 9, 2023 · 48d946a · 48d946a
1 parent a1d2460
commit 48d946a
Show file tree

Hide file tree

Showing 4 changed files with 76 additions and 2 deletions.
diff --git a/docs/changelog/100360.yaml b/docs/changelog/100360.yaml
@@ -0,0 +1,5 @@
+pr: 100360
+summary: "ESQL: Limit how many bytes `concat()` can process"
+area: ES|QL
+type: bug
+issues: []
diff --git a/...-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/HeapAttackIT.java b/...-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/HeapAttackIT.java
@@ -152,10 +152,16 @@ public void testSmallConcat() throws IOException {
         assertMap(map, matchesMap().entry("columns", columns).entry("values", values));
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/99826")
     public void testHugeConcat() throws IOException {
         initSingleDocIndex();
-        assertCircuitBreaks(() -> concat(10));
+        ResponseException e = expectThrows(ResponseException.class, () -> concat(10));
+        Map<?, ?> map = XContentHelper.convertToMap(JsonXContent.jsonXContent, EntityUtils.toString(e.getResponse().getEntity()), false);
+        logger.info("expected request rejected {}", map);
+        assertMap(
+            map,
+            matchesMap().entry("status", 400)
+                .entry("error", matchesMap().extraOk().entry("reason", "concatenating more than [1048576] bytes is not supported"))
+        );
     }
 
     private Response concat(int evals) throws IOException {

diff --git a/.../src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/Concat.java b/.../src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/Concat.java
@@ -13,6 +13,8 @@
 import org.elasticsearch.compute.operator.BreakingBytesRefBuilder;
 import org.elasticsearch.compute.operator.EvalOperator;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.xpack.esql.EsqlClientException;
 import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper;
 import org.elasticsearch.xpack.ql.expression.Expression;
 import org.elasticsearch.xpack.ql.expression.Expressions;
@@ -27,13 +29,17 @@
 import java.util.function.Function;
 import java.util.stream.Stream;
 
+import static org.elasticsearch.common.unit.ByteSizeUnit.MB;
 import static org.elasticsearch.xpack.ql.expression.TypeResolutions.ParamOrdinal.DEFAULT;
 import static org.elasticsearch.xpack.ql.expression.TypeResolutions.isString;
 
 /**
  * Join strings.
  */
 public class Concat extends ScalarFunction implements EvaluatorMapper {
+
+    static final long MAX_CONCAT_LENGTH = MB.toBytes(1);
+
     public Concat(Source source, Expression first, List<? extends Expression> rest) {
         super(source, Stream.concat(Stream.of(first), rest.stream()).toList());
     }
@@ -83,13 +89,30 @@ public ExpressionEvaluator.Factory toEvaluator(Function<Expression, ExpressionEv
 
     @Evaluator
     static BytesRef process(@Fixed(includeInToString = false) BreakingBytesRefBuilder scratch, BytesRef[] values) {
+        scratch.grow(checkedTotalLength(values));
         scratch.clear();
         for (int i = 0; i < values.length; i++) {
             scratch.append(values[i]);
         }
         return scratch.bytesRefView();
     }
 
+    private static int checkedTotalLength(BytesRef[] values) {
+        int length = 0;
+        for (var v : values) {
+            length += v.length;
+        }
+        if (length > MAX_CONCAT_LENGTH) {
+            throw new EsqlClientException("concatenating more than [" + MAX_CONCAT_LENGTH + "] bytes is not supported") {
+                @Override
+                public RestStatus status() {
+                    return RestStatus.BAD_REQUEST; // return a 400 response
+                }
+            };
+        }
+        return length;
+    }
+
     @Override
     public Expression replaceChildren(List<Expression> newChildren) {
         return new Concat(source(), newChildren.get(0), newChildren.subList(1, newChildren.size()));

diff --git a/...test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/ConcatTests.java b/...test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/ConcatTests.java
@@ -13,22 +13,26 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.operator.EvalOperator;
+import org.elasticsearch.xpack.esql.EsqlClientException;
 import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase;
 import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
 import org.elasticsearch.xpack.esql.type.EsqlDataTypes;
 import org.elasticsearch.xpack.ql.expression.Expression;
+import org.elasticsearch.xpack.ql.expression.Literal;
 import org.elasticsearch.xpack.ql.tree.Source;
 import org.elasticsearch.xpack.ql.type.DataType;
 import org.elasticsearch.xpack.ql.type.DataTypes;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.function.Supplier;
 import java.util.stream.IntStream;
 
 import static org.elasticsearch.compute.data.BlockUtils.toJavaObject;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
 
 public class ConcatTests extends AbstractFunctionTestCase {
     public ConcatTests(@Name("TestCase") Supplier<TestCaseSupplier.TestCase> testCaseSupplier) {
@@ -118,11 +122,47 @@ public void testSomeConstant() {
             assertThat(expression.typeResolved().message(), equalTo(testCase.getExpectedTypeError()));
             return;
         }
+
+        int totalLength = testDataLength();
+        if (totalLength >= Concat.MAX_CONCAT_LENGTH || rarely()) {
+            boolean hasNulls = mix.stream().anyMatch(x -> x instanceof Literal l && l.value() == null)
+                || fieldValues.stream().anyMatch(Objects::isNull);
+            if (hasNulls == false) {
+                testOversized(totalLength, mix, fieldValues);
+                return;
+            }
+        }
+
         try (
             EvalOperator.ExpressionEvaluator eval = evaluator(expression).get(driverContext());
             Block.Ref ref = eval.eval(row(fieldValues))
         ) {
             assertThat(toJavaObject(ref.block(), 0), testCase.getMatcher());
         }
     }
+
+    private void testOversized(int totalLen, List<Expression> mix, List<Object> fieldValues) {
+        for (int len; totalLen < Concat.MAX_CONCAT_LENGTH; totalLen += len) {
+            len = randomIntBetween(1, (int) Concat.MAX_CONCAT_LENGTH);
+            mix.add(new Literal(Source.EMPTY, new BytesRef(randomAlphaOfLength(len)), DataTypes.KEYWORD));
+        }
+        Expression expression = build(testCase.getSource(), mix);
+        Exception e = expectThrows(EsqlClientException.class, () -> {
+            try (
+                EvalOperator.ExpressionEvaluator eval = evaluator(expression).get(driverContext());
+                Block.Ref ref = eval.eval(row(fieldValues));
+            ) {}
+        });
+        assertThat(e.getMessage(), is("concatenating more than [1048576] bytes is not supported"));
+    }
+
+    private int testDataLength() {
+        int totalLength = 0;
+        for (var data : testCase.getData()) {
+            if (data.data() instanceof BytesRef bytesRef) {
+                totalLength += bytesRef.length;
+            }
+        }
+        return totalLength;
+    }
 }