Browse Source

ESQL: Limit how many bytes `concat()` can process (#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: #100288.
Bogdan Pintea 2 years ago
parent
commit
c879775528

+ 5 - 0
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: []

+ 8 - 2
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 class HeapAttackIT extends ESRestTestCase {
         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 {

+ 23 - 0
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.ann.Fixed;
 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,6 +29,7 @@ import java.util.List;
 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;
 
@@ -34,6 +37,9 @@ 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,6 +89,7 @@ public class Concat extends ScalarFunction implements EvaluatorMapper {
 
     @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]);
@@ -90,6 +97,22 @@ public class Concat extends ScalarFunction implements EvaluatorMapper {
         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()));

+ 40 - 0
x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/ConcatTests.java

@@ -13,22 +13,26 @@ import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 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,6 +122,17 @@ public class ConcatTests extends AbstractFunctionTestCase {
             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))
@@ -125,4 +140,29 @@ public class ConcatTests extends AbstractFunctionTestCase {
             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;
+    }
 }