Skip to content

Commit ae10b41

Browse files
fix: scan schema.table without a backtracking regex
CodeQL flagged QUALIFIED_TABLE as polynomial ReDoS on Agent Q&A (long runs of 'A'). Walk the string once instead. Co-authored-by: Venkat SF <venkatesh.sakamuri@stayflexi.com>
1 parent d3ff029 commit ae10b41

2 files changed

Lines changed: 75 additions & 9 deletions

File tree

backend/src/main/java/com/dbaagent/service/brain/core/BrainNoteIntentService.java

Lines changed: 66 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,6 @@
1919
@Service
2020
public class BrainNoteIntentService {
2121

22-
private static final Pattern QUALIFIED_TABLE = Pattern.compile(
23-
"(?<![\\w.])([a-zA-Z_][\\w]*)\\.([a-zA-Z_][\\w]*)(?![\\w.])"
24-
);
2522
private static final Pattern BACKTICK_IDENT = Pattern.compile("`([^`]+)`");
2623
private static final Pattern DEFINITION_CUE = Pattern.compile(
2724
"\\b(metric|definition|means|pinned|use this|correct|from|count|always|filter|join)\\b",
@@ -242,16 +239,76 @@ private Optional<String[]> resolveTarget(String text) {
242239
if (table != null) {
243240
return Optional.of(new String[] { table, column });
244241
}
245-
Matcher tables = QUALIFIED_TABLE.matcher(text);
246-
while (tables.find()) {
247-
String schema = tables.group(1);
248-
String name = tables.group(2);
242+
String[] prose = findFirstQualifiedTable(text);
243+
if (prose != null) {
244+
return Optional.of(new String[] { prose[0] + "." + prose[1], column });
245+
}
246+
return Optional.empty();
247+
}
248+
249+
/**
250+
* Linear scan for {@code schema.table} (not {@code a.b.c}). Avoids the
251+
* backtracking {@code [\w]*\.[\w]*} pattern CodeQL flags as ReDoS on
252+
* attacker-controlled Agent answers.
253+
*/
254+
static String[] findFirstQualifiedTable(String text) {
255+
if (text == null || text.isEmpty()) {
256+
return null;
257+
}
258+
int n = text.length();
259+
int i = 0;
260+
while (i < n) {
261+
char c = text.charAt(i);
262+
if (!isIdentStart(c)) {
263+
i++;
264+
continue;
265+
}
266+
if (i > 0) {
267+
char prev = text.charAt(i - 1);
268+
if (isIdentPart(prev) || prev == '.') {
269+
i++;
270+
continue;
271+
}
272+
}
273+
int schemaStart = i;
274+
i++;
275+
while (i < n && isIdentPart(text.charAt(i))) {
276+
i++;
277+
}
278+
if (i >= n || text.charAt(i) != '.') {
279+
continue;
280+
}
281+
int schemaEnd = i;
282+
i++;
283+
if (i >= n || !isIdentStart(text.charAt(i))) {
284+
continue;
285+
}
286+
int nameStart = i;
287+
i++;
288+
while (i < n && isIdentPart(text.charAt(i))) {
289+
i++;
290+
}
291+
if (i < n) {
292+
char next = text.charAt(i);
293+
if (isIdentPart(next) || next == '.') {
294+
continue;
295+
}
296+
}
297+
String schema = text.substring(schemaStart, schemaEnd);
249298
if (STOP_TABLES.contains(schema.toLowerCase(Locale.ROOT))) {
250299
continue;
251300
}
252-
return Optional.of(new String[] { schema + "." + name, column });
301+
return new String[] { schema, text.substring(nameStart, i) };
253302
}
254-
return Optional.empty();
303+
return null;
304+
}
305+
306+
private static boolean isIdentStart(char c) {
307+
return (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || c == '_';
308+
}
309+
310+
private static boolean isIdentPart(char c) {
311+
return isIdentStart(c) || (c >= '0' && c <= '9');
255312
}
256313

257314
/**

backend/src/test/java/com/dbaagent/service/brain/core/BrainNoteIntentServiceTest.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,15 @@ void proposeFromTurn_ignoresNonDefinitionAnswers() {
9494
assertThat(proposal).isEmpty();
9595
}
9696

97+
@Test
98+
void findFirstQualifiedTable_skipsTripleQualifiedAndCatalogSchemas() {
99+
assertThat(BrainNoteIntentService.findFirstQualifiedTable("see marts.dim_person.col then crm.accounts"))
100+
.containsExactly("crm", "accounts");
101+
assertThat(BrainNoteIntentService.findFirstQualifiedTable("pg_catalog.pg_class")).isNull();
102+
assertThat(BrainNoteIntentService.findFirstQualifiedTable("A".repeat(20_000) + " marts.dim_person"))
103+
.containsExactly("marts", "dim_person");
104+
}
105+
97106
@Test
98107
void mergeTexts_doesNotDuplicateTheSameSentence() {
99108
String merged = service.mergeTexts(

0 commit comments

Comments
 (0)