ESQL: Limit how many bytes concat() can process

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: []

x-pack/plugin/esql/qa/server/single-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 {

x-pack/plugin/esql/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()));

x-pack/plugin/esql/src/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;
+    }
 }