Skip to content

Commit

Permalink
Full implementation of tree depth Circuit Breaker
Browse files Browse the repository at this point in the history
  • Loading branch information
@matriv
matriv committed Sep 21, 2018
1 parent 37fb6bf commit 7e09985
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 15 deletions.
27 changes: 18 additions & 9 deletions x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/parser/SqlParser.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
*/
package org.elasticsearch.xpack.sql.parser;

import com.carrotsearch.hppc.ObjectIntHashMap;
import org.antlr.v4.runtime.BaseErrorListener;
import org.antlr.v4.runtime.CharStream;
import org.antlr.v4.runtime.CommonToken;
Expand Down Expand Up @@ -99,7 +100,7 @@ private <T> T invokeParser(String sql,
CommonTokenStream tokenStream = new CommonTokenStream(tokenSource);
SqlBaseParser parser = new SqlBaseParser(tokenStream);

parser.addParseListener(new CircuitBreakerProcessor());
parser.addParseListener(new CircuitBreakerListener());
parser.addParseListener(new PostProcessor(Arrays.asList(parser.getRuleNames())));

parser.removeErrorListeners();
Expand Down Expand Up @@ -212,18 +213,26 @@ public void exitNonReserved(SqlBaseParser.NonReservedContext context) {
/**
* Used to catch large expressions that can lead to stack overflows
*/
private class CircuitBreakerProcessor extends SqlBaseBaseListener {
private class CircuitBreakerListener extends SqlBaseBaseListener {

private static final short MAX_BOOLEAN_ELEMENTS = 1000;
private short countElementsInBooleanExpressions = 0;
private static final short MAX_DEPTH = 100;

// Keep current depth for every rule visited
ObjectIntHashMap<String> depthCounts = new ObjectIntHashMap<>(100);

@Override
public void enterLogicalBinary(SqlBaseParser.LogicalBinaryContext ctx) {
if (++countElementsInBooleanExpressions == MAX_BOOLEAN_ELEMENTS) {
throw new ParsingException("boolean expression is too large to parse, (exceeds {} elements)",
MAX_BOOLEAN_ELEMENTS);
public void enterEveryRule(ParserRuleContext ctx) {
int currentDepth = depthCounts.putOrAdd(ctx.getClass().getSimpleName(), 1, 1);
if (currentDepth > MAX_DEPTH) {
throw new ParsingException("expression is too large to parse, (tree's depth exceeds {})", MAX_DEPTH);
}
super.enterLogicalBinary(ctx);
super.enterEveryRule(ctx);
}

@Override
public void exitEveryRule(ParserRuleContext ctx) {
depthCounts.putOrAdd(ctx.getClass().getSimpleName(), 0, -1);
super.exitEveryRule(ctx);
}
}

Expand Down
55 changes: 49 additions & 6 deletions x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/parser/SqlParserTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -138,14 +138,57 @@ public void testMultiMatchQuery() {
assertThat(mmqp.optionMap(), hasEntry("fuzzy_rewrite", "scoring_boolean"));
}

public void testLimitToPreventStackOverflowFromLargeBooleanExpression() {
// 1000 elements is ok
new SqlParser().createExpression(Joiner.on(" OR ").join(nCopies(1000, "a = b")));
public void testLimitToPreventStackOverflowFromLargeUnaryBooleanExpression() {
// 100 elements is ok
new SqlParser().createExpression(
Joiner.on("NOT(").join(nCopies(100, "true")).concat(Joiner.on("").join(nCopies(99, ")"))));

// 1001 elements parser's "circuit breaker" is triggered
// 500 elements parser's "circuit breaker" is triggered
ParsingException e = expectThrows(ParsingException.class, () -> new SqlParser().createExpression(
Joiner.on("NOT(").join(nCopies(101, "true")).concat(Joiner.on("").join(nCopies(100, ")")))));
assertEquals("expression is too large to parse, (tree's depth exceeds 100)", e.getErrorMessage());
}

public void testLimitToPreventStackOverflowFromLargeBinaryBooleanExpression() {
// 100 elements is ok
new SqlParser().createExpression(Joiner.on(" OR ").join(nCopies(100, "true")));

// 101 elements parser's "circuit breaker" is triggered
ParsingException e = expectThrows(ParsingException.class, () ->
new SqlParser().createExpression(Joiner.on(" OR ").join(nCopies(1001, "a = b"))));
assertEquals("boolean expression is too large to parse, (exceeds 1000 elements)", e.getErrorMessage());
new SqlParser().createExpression(Joiner.on(" OR ").join(nCopies(101, "a = b"))));
assertEquals("expression is too large to parse, (tree's depth exceeds 100)", e.getErrorMessage());
}

public void testLimitToPreventStackOverflowFromLargeUnaryArithmeticExpression() {
// 100 elements is ok
new SqlParser().createExpression(
Joiner.on("abs(").join(nCopies(100, "i")).concat(Joiner.on("").join(nCopies(99, ")"))));

// 101 elements parser's "circuit breaker" is triggered
ParsingException e = expectThrows(ParsingException.class, () -> new SqlParser().createExpression(
Joiner.on("abs(").join(nCopies(101, "i")).concat(Joiner.on("").join(nCopies(100, ")")))));
assertEquals("expression is too large to parse, (tree's depth exceeds 100)", e.getErrorMessage());
}

public void testLimitToPreventStackOverflowFromLargeBinaryArithmeticExpression() {
// 100 elements is ok
new SqlParser().createExpression(Joiner.on(" + ").join(nCopies(100, "a")));

// 101 elements parser's "circuit breaker" is triggered
ParsingException e = expectThrows(ParsingException.class, () ->
new SqlParser().createExpression(Joiner.on(" + ").join(nCopies(101, "a"))));
assertEquals("expression is too large to parse, (tree's depth exceeds 100)", e.getErrorMessage());
}

public void testLimitToPreventStackOverflowFromLargeSubselectTree() {
// 100 elements is ok
new SqlParser().createStatement(
Joiner.on(" (").join(nCopies(100, "SELECT * FROM")).concat("t").concat(Joiner.on("").join(nCopies(99, ")"))));

// 101 elements parser's "circuit breaker" is triggered
ParsingException e = expectThrows(ParsingException.class, () -> new SqlParser().createStatement(
Joiner.on(" (").join(nCopies(101, "SELECT * FROM")).concat("t").concat(Joiner.on("").join(nCopies(100, ")")))));
assertEquals("expression is too large to parse, (tree's depth exceeds 100)", e.getErrorMessage());
}

private LogicalPlan parseStatement(String sql) {
Expand Down

0 comments on commit 7e09985

Please sign in to comment.