Skip to content

Commit

Permalink
Enforce term limits during EdgeQueries (NationalSecurityAgency#2274)
Browse files Browse the repository at this point in the history
* Enforce term limits during EdgeQueries
* Update error message
* Fix EdgeQuery maxQueryTerms setting
  • Loading branch information
jschmidt10 authored Feb 8, 2024
1 parent 45623c7 commit c5287da
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 14 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ public EdgeQueryConfiguration(EdgeQueryLogic configuredLogic, Query query) {
setModelName(configuredLogic.getModelName());
setModelTableName(configuredLogic.getModelTableName());
setEdgeQueryModel(configuredLogic.getEdgeQueryModel());
setMaxQueryTerms(configuredLogic.getMaxQueryTerms());
}

public List<? extends Type<?>> getDataTypes() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -168,11 +168,6 @@ public Object visit(ASTReferenceExpression node, Object data) {
*/
@Override
public Object visit(ASTAndNode node, Object data) {
if (termCount > maxTerms) {
log.error("Query has too many terms");
throw new IllegalArgumentException("Too many search terms " + termCount);
}

// run the visitor against all of the children
List<List<? extends EdgeContext>> childContexts = new ArrayList<>(node.jjtGetNumChildren());
for (int i = 0; i < node.jjtGetNumChildren(); i++) {
Expand Down Expand Up @@ -274,11 +269,6 @@ public Object visit(ASTAndNode node, Object data) {
*/
@Override
public Object visit(ASTOrNode node, Object data) {
if (termCount > maxTerms) {
log.error("Query has too many terms");
throw new IllegalArgumentException("Too many search terms " + termCount);
}

// run the visitor against all of the children
List<List<? extends EdgeContext>> childContexts = new ArrayList<>(node.jjtGetNumChildren());
for (int i = 0; i < node.jjtGetNumChildren(); i++) {
Expand Down Expand Up @@ -410,10 +400,18 @@ private void checkNotExclusion(IdentityContext context, String msg) {
*/
@Override
public Object visit(ASTEQNode node, Object data) {
termCount++;
incrementTermCountAndCheck();
return visitExpresionNode(node, EQUALS);
}

private void incrementTermCountAndCheck() {
if (++termCount > maxTerms) {
String message = "Exceeded max term limit of " + maxTerms;
log.error(message);
throw new IllegalArgumentException(message);
}
}

/**
* Equals node (=~) should have exactly two children (reference nodes) one child will have the identifier eg: SOURCE ... the other child will have the
* string literal eg: 'searchTerm'
Expand All @@ -428,7 +426,7 @@ public Object visit(ASTEQNode node, Object data) {
*/
@Override
public Object visit(ASTERNode node, Object data) {
termCount++;
incrementTermCountAndCheck();
List<IdentityContext> contexts = (List<IdentityContext>) visitExpresionNode(node, EQUALS_REGEX);
if (contexts.get(0).getIdentity().equals(EDGE_SOURCE)) {
sawEquivalenceRegexSource = true;
Expand All @@ -440,13 +438,13 @@ public Object visit(ASTERNode node, Object data) {

@Override
public Object visit(ASTNRNode node, Object data) {
termCount++;
incrementTermCountAndCheck();
return visitExpresionNode(node, NOT_EQUALS_REGEX);
}

@Override
public Object visit(ASTNENode node, Object data) {
termCount++;
incrementTermCountAndCheck();
return visitExpresionNode(node, NOT_EQUALS);
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
package datawave.query.jexl.visitors;

import static java.util.Collections.emptyList;

import static org.junit.Assert.assertThrows;

import org.apache.commons.jexl3.JexlFeatures;
import org.apache.commons.jexl3.parser.ASTJexlScript;
import org.apache.commons.jexl3.parser.Parser;
import org.apache.commons.jexl3.parser.StringProvider;
import org.junit.Before;
import org.junit.Test;

import datawave.query.tables.edge.EdgeQueryLogic;

public class EdgeTableRangeBuildingVisitorTest {

private int termLimit = 3;
private Parser parser;

private EdgeTableRangeBuildingVisitor visitor;

@Before
public void setup() {
parser = new Parser(new StringProvider(";"));

visitor = new EdgeTableRangeBuildingVisitor(false, emptyList(), termLimit, emptyList());
}

@Test
public void shouldEnforceTermLimit() {
ASTJexlScript parsedQuery = parseQuery("TYPE == 'like it' OR TYPE == 'love it' OR TYPE == 'gotta have it' OR TYPE == 'hand it over or else'");

assertThrows(IllegalArgumentException.class, () -> parsedQuery.jjtAccept(visitor, null));
}

private ASTJexlScript parseQuery(String query) {
return parser.parse(null, new JexlFeatures(), EdgeQueryLogic.fixQueryString(query), null);
}
}

0 comments on commit c5287da

Please sign in to comment.