Skip to content

fix(security): enforce the chat schema allowlist over the whole statement #70

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
geekypunk wants to merge 3 commits into
base: main
Choose a base branch
from
+310 −23
Open

fix(security): enforce the chat schema allowlist over the whole statement #70

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a547a07
fix(security): enforce the chat schema allowlist over the whole state…
geekypunk Aug 19, 2026
70a1472
fix(security): fail closed when a protected table sits where inspecti…
geekypunk Aug 19, 2026
bb8b895
fix(security): stop the fail-closed check conflating same-named table…
geekypunk Aug 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
197 changes: 174 additions & 23 deletions backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,9 @@
import net.sf.jsqlparser.statement.select.ParenthesedSelect;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SetOperationList;
import net.sf.jsqlparser.statement.select.WithItem;
import net.sf.jsqlparser.util.TablesNamesFinder;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
Expand Down Expand Up @@ -117,9 +120,19 @@ public QueryGuardDecision enforcePreExecution(

try {
Statement parsed = CCJSqlParserUtil.parse(queryRequest.getQuery());
if (parsed instanceof Select select && select.getPlainSelect() != null) {
enforceAllowedSchemas(select.getPlainSelect(), policy.allowedSchemas());
QueryInspection inspection = inspectPlainSelect(select.getPlainSelect(), protectedObjects);
if (parsed instanceof Select select) {
// Enumerate over the WHOLE statement, not just FROM/JOIN of the
// outermost PlainSelect. An allowlist is only sound if the walk is
// total: any node left unvisited is implicitly permitted, which is
// how a subquery, UNION branch, or CTE body reached a schema
// outside the caller's scope.
enforceAllowedSchemas(parsed, policy.allowedSchemas());
assertProtectedTablesAreInspectable(parsed, collectPlainSelects(select), protectedObjects);
// Likewise inspect every branch. Gating this on getPlainSelect()
// != null skipped protection entirely for a SetOperationList,
// because a UNION's body is not a PlainSelect.
for (PlainSelect branch : collectPlainSelects(select)) {
QueryInspection inspection = inspectPlainSelect(branch, protectedObjects);
if (inspection.selectsWildcardFromProtectedTable || inspection.rawProtectedColumnsSelected) {
logPolicyEvent(SecurityEventType.CHAT_ACCESS_POLICY_BLOCKED, executionContext.actorUsername(), connectionId, "sql_blocked", Map.of(
"query", truncate(queryRequest.getQuery()),
Expand All @@ -132,6 +145,7 @@ public QueryGuardDecision enforcePreExecution(
"POLICY_SQL_BLOCKED"
);
}
}
}
} catch (UserDataAccessPolicyException e) {
throw e;
Expand Down Expand Up @@ -324,16 +338,167 @@ private boolean mentionsProtectedTables(String normalized, ConnectionChatAccessP
return policy.impactedTables().stream().anyMatch(table -> normalized.contains(table.toLowerCase(Locale.ROOT)));
}

private void enforceAllowedSchemas(PlainSelect select, Set<String> allowedSchemas) {
if (allowedSchemas == null || allowedSchemas.isEmpty()) {
/**
* Fails closed on statement shapes inspection cannot reach.
*
* TablesNamesFinder sees every table in the statement; collectPlainSelects
* deliberately does not descend into a select nested inside FROM/JOIN/WHERE/
* HAVING, because enumerating arbitrary expression trees correctly is the very
* thing that went wrong here the first time. So instead of trying harder to
* walk, compare the two: when a protected table is referenced somewhere the
* column inspection could not examine, refuse the query.
*
* A syntax form we failed to enumerate must never become an implicit permit --
* that is exactly how a subquery, a UNION branch and a CTE body each evaded
* the schema allowlist. Refusing costs a conservative block on some safe
* nested aggregates; allowing costs the data.
*/
private void assertProtectedTablesAreInspectable(
Statement statement,
List<PlainSelect> inspectedBranches,
Map<String, ConnectionChatAccessPolicyService.ProtectionDescriptor> protectedObjects
) {
if (protectedObjects == null || protectedObjects.isEmpty()) {
return;
}
Map<String, String> aliasToTable = buildAliasMap(select);
Set<String> referencedSchemas = new LinkedHashSet<>();
collectReferencedSchemas(select.getFromItem(), aliasToTable, referencedSchemas);
Set<String> referenced = new LinkedHashSet<>();
for (String name : new TablesNamesFinder<>().getTables(statement)) {
if (name != null && !name.isBlank()) {
referenced.add(normalizeName(name));
}
}
Set<String> inspected = new LinkedHashSet<>();
for (PlainSelect branch : inspectedBranches) {
collectDirectTables(branch, inspected);
}
for (ConnectionChatAccessPolicyService.ProtectionDescriptor descriptor : protectedObjects.values()) {
String protectedName = descriptor.qualifiedTableName();
boolean isReferenced = referenced.stream().anyMatch(name -> namesMatch(protectedName, name));
boolean wasInspected = inspected.stream().anyMatch(name -> namesMatch(protectedName, name));
if (isReferenced && !wasInspected) {
throw new UserDataAccessPolicyException(
"This query reaches restricted data through a nested query DeepSQL cannot fully verify, so it was blocked before execution.",
"POLICY_SQL_BLOCKED"
);
}
}
}

/**
* Does a query's table reference name the protected table?
*
* Asymmetric on purpose, because the two sides carry different information.
* ConnectionChatAccessPolicyService.qualifyTable() drops the schema when it is
* "public", so a bare PROTECTED name means public.<table> -- it is not unknown.
* A bare REFERENCE in a query is genuinely unknown: it resolves through the
* session search_path and could be any schema.
*
* reference unqualified -> match on bare name. Ambiguous, so block; the
* search_path may well point at the protected table.
* protected public -> the qualified reference must actually say public.
* marts.customer_profiles is a different table, and
* treating it as protected refused every other
* schema's copy -- which this product's own
* multi-schema fixtures (crm/sales/finance/hr) hit.
* both qualified -> exact match.
*/
private boolean namesMatch(String protectedName, String referencedName) {
String protectedNorm = normalizeName(protectedName);
String referencedNorm = normalizeName(referencedName);
if (protectedNorm.isEmpty() || referencedNorm.isEmpty()) {
return false;
}
if (!referencedNorm.contains(".")) {
return bareName(protectedNorm).equals(referencedNorm);
}
if (!protectedNorm.contains(".")) {
return referencedNorm.equals("public." + protectedNorm);
}
return protectedNorm.equals(referencedNorm);
}

private String bareName(String normalizedName) {
int dot = normalizedName.lastIndexOf('.');
return dot > 0 && dot < normalizedName.length() - 1
? normalizedName.substring(dot + 1)
: normalizedName;
}

/** Tables named directly in this branch's FROM/JOIN -- what inspection actually saw. */
private void collectDirectTables(PlainSelect select, Set<String> out) {
if (select.getFromItem() instanceof Table table) {
out.add(normalizeName(table.getFullyQualifiedName()));
}
if (select.getJoins() != null) {
for (Join join : select.getJoins()) {
collectReferencedSchemas(join.getRightItem(), aliasToTable, referencedSchemas);
if (join.getRightItem() instanceof Table table) {
out.add(normalizeName(table.getFullyQualifiedName()));
}
}
}
}

/**
* Collects every PlainSelect in a statement: the top level, each branch of a
* set operation (UNION/INTERSECT/EXCEPT), parenthesised selects, and every
* CTE body. Callers that inspect only the outermost select leave the rest
* unprotected.
*/
private List<PlainSelect> collectPlainSelects(Select select) {
List<PlainSelect> found = new ArrayList<>();
collectPlainSelects(select, found);
return found;
}

private void collectPlainSelects(Select select, List<PlainSelect> found) {
if (select == null) {
return;
}
if (select.getWithItemsList() != null) {
for (WithItem<?> item : select.getWithItemsList()) {
if (item != null) {
collectPlainSelects(item.getSelect(), found);
}
}
}
if (select instanceof PlainSelect plain) {
found.add(plain);
} else if (select instanceof SetOperationList setOps) {
if (setOps.getSelects() != null) {
for (Select branch : setOps.getSelects()) {
collectPlainSelects(branch, found);
}
}
} else if (select instanceof ParenthesedSelect parenthesed) {
collectPlainSelects(parenthesed.getSelect(), found);
}
}

/**
* Enforces the schema allowlist over every table reference in the statement.
*
* Uses JSqlParser's TablesNamesFinder rather than a hand-rolled walk of
* FROM and JOIN. The distinction is the whole fix: enumeration has to be
* exhaustive by construction, because an allowlist implemented as a partial
* walk implicitly permits every syntax position the walker forgot -- here a
* subquery in WHERE/HAVING/SELECT, a UNION branch, and a CTE body.
*
* Bare (unqualified) names stay unchecked, exactly as before: they resolve
* through the session search_path, and CTE names are not tables at all.
*/
private void enforceAllowedSchemas(Statement statement, Set<String> allowedSchemas) {
if (allowedSchemas == null || allowedSchemas.isEmpty()) {
return;
}
Set<String> referencedSchemas = new LinkedHashSet<>();
for (String qualifiedName : new TablesNamesFinder<>().getTables(statement)) {
if (qualifiedName == null) {
continue;
}
String[] parts = qualifiedName.split("\\.");
if (parts.length >= 2) {
// db.schema.table and schema.table both put the schema second-to-last.
referencedSchemas.add(parts[parts.length - 2]);
}
}
for (String schema : referencedSchemas) {
Expand All @@ -346,20 +511,6 @@ private void enforceAllowedSchemas(PlainSelect select, Set<String> allowedSchema
}
}

private void collectReferencedSchemas(FromItem fromItem, Map<String, String> aliasToTable, Set<String> schemas) {
if (fromItem instanceof Table table) {
String schema = table.getSchemaName();
if (schema != null && !schema.isBlank()) {
schemas.add(schema);
return;
}
String qualified = aliasToTable.get(normalizeName(table.getName()));
if (qualified != null && qualified.contains(".")) {
schemas.add(qualified.substring(0, qualified.indexOf('.')));
}
}
}

private boolean containsDangerousProtectedReference(
String normalizedSql,
Map<String, ConnectionChatAccessPolicyService.ProtectionDescriptor> protectedObjects
Expand Down
136 changes: 136 additions & 0 deletions backend/src/test/java/com/dbaagent/service/UserDataAccessPolicyServiceTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,142 @@ void enforcePreExecution_blocksQueriesOutsideAllowedSchema() {
assertThat(exception.getErrorCode()).isEqualTo("POLICY_SCHEMA_BLOCKED");
}

private ConnectionChatAccessPolicyService.EffectivePolicy martsOnlyPolicy() {
return new ConnectionChatAccessPolicyService.EffectivePolicy(
true, "conn-1", "analyst",
Set.of(), Set.of(), Set.of(), Set.of("marts"),
true, true, "Only schema marts", List.of(), List.of()
);
}

private UserDataAccessPolicyException assertSchemaScopeBlocks(String sql) {
ConnectionChatAccessPolicyService.EffectivePolicy schemaPolicy = martsOnlyPolicy();
when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(schemaPolicy);
lenient().when(policyService.buildProtectionDescriptors(schemaPolicy)).thenReturn(Map.of());
return assertThrows(
UserDataAccessPolicyException.class,
() -> service.enforcePreExecution(
"conn-1",
new QueryRequest(sql, null, null),
new QueryExecutionContext(QueryExecutionOrigin.CHAT, QueryExecutionContext.MutationMode.READ_ONLY_ONLY, "analyst", false, false)
)
);
}

// The schema allowlist walked only FROM and JOIN, so a forbidden schema
// reached through any other syntax position was never enumerated and the
// allowlist silently permitted it.

// Column inspection cannot reach a select nested inside FROM/JOIN/WHERE, so a
// protected table referenced there must be refused rather than implicitly allowed.

@Test
void enforcePreExecution_blocksProtectedTableInsideDerivedTable() {
when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(policy());

UserDataAccessPolicyException exception = assertThrows(
UserDataAccessPolicyException.class,
() -> service.enforcePreExecution(
"conn-1",
new QueryRequest("SELECT t.email FROM (SELECT email FROM customer_profiles) t", null, null),
new QueryExecutionContext(QueryExecutionOrigin.CHAT, QueryExecutionContext.MutationMode.READ_ONLY_ONLY, "analyst", false, false)
)
);

assertThat(exception.getErrorCode()).isEqualTo("POLICY_SQL_BLOCKED");
}

@Test
void enforcePreExecution_blocksProtectedTableInsideWhereSubquery() {
when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(policy());

UserDataAccessPolicyException exception = assertThrows(
UserDataAccessPolicyException.class,
() -> service.enforcePreExecution(
"conn-1",
new QueryRequest("SELECT id FROM orders WHERE id IN (SELECT email FROM customer_profiles)", null, null),
new QueryExecutionContext(QueryExecutionOrigin.CHAT, QueryExecutionContext.MutationMode.READ_ONLY_ONLY, "analyst", false, false)
)
);

assertThat(exception.getErrorCode()).isEqualTo("POLICY_SQL_BLOCKED");
}

// A qualified protection names exactly one table. marts.customer_profiles is a
// different table from public.customer_profiles and must not be caught by it.
@Test
void enforcePreExecution_allowsSameNamedTableInAnotherSchemaWhenProtectionIsQualified() {
when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(policy());
when(policyService.buildProtectionDescriptors(any(ConnectionChatAccessPolicyService.EffectivePolicy.class)))
.thenReturn(Map.of("public.customer_profiles",
descriptor("public", "customer_profiles", false, "email")));

service.enforcePreExecution(
"conn-1",
new QueryRequest("SELECT id FROM (SELECT id FROM marts.customer_profiles) t", null, null),
new QueryExecutionContext(QueryExecutionOrigin.CHAT, QueryExecutionContext.MutationMode.READ_ONLY_ONLY, "analyst", false, false)
);
}

// The genuinely ambiguous direction is an unqualified REFERENCE, not an
// unqualified protection: qualifyTable() stores public.<t> as bare <t>, so a
// bare protected name means public, while a bare reference in a query
// resolves through search_path and could be any schema. Block that one.
@Test
void enforcePreExecution_blocksUnqualifiedReferenceBecauseSearchPathIsUnknown() {
when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(policy());

UserDataAccessPolicyException exception = assertThrows(
UserDataAccessPolicyException.class,
() -> service.enforcePreExecution(
"conn-1",
new QueryRequest("SELECT id FROM (SELECT id FROM customer_profiles) t", null, null),
new QueryExecutionContext(QueryExecutionOrigin.CHAT, QueryExecutionContext.MutationMode.READ_ONLY_ONLY, "analyst", false, false)
)
);

assertThat(exception.getErrorCode()).isEqualTo("POLICY_SQL_BLOCKED");
}

@Test
void enforcePreExecution_blocksForbiddenSchemaInsideWhereSubquery() {
assertThat(assertSchemaScopeBlocks(
"SELECT id FROM marts.orders WHERE total = (SELECT MAX(salary) FROM hr.salaries)"
).getErrorCode()).isEqualTo("POLICY_SCHEMA_BLOCKED");
}

@Test
void enforcePreExecution_blocksForbiddenSchemaInUnionBranch() {
assertThat(assertSchemaScopeBlocks(
"SELECT id FROM marts.orders UNION ALL SELECT ssn FROM hr.salaries"
).getErrorCode()).isEqualTo("POLICY_SCHEMA_BLOCKED");
}

// The same fail-open gate skipped protected-column inspection entirely, so a
// UNION reached restricted columns even inside an allowed schema.
@Test
void enforcePreExecution_blocksProtectedColumnsInUnionBranch() {
when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(policy());

UserDataAccessPolicyException exception = assertThrows(
UserDataAccessPolicyException.class,
() -> service.enforcePreExecution(
"conn-1",
new QueryRequest("SELECT 1 AS x UNION ALL SELECT email FROM customer_profiles", null, null),
new QueryExecutionContext(QueryExecutionOrigin.CHAT, QueryExecutionContext.MutationMode.READ_ONLY_ONLY, "analyst", false, false)
)
);

assertThat(exception.getErrorCode()).isEqualTo("POLICY_SQL_BLOCKED");
}

@Test
void enforcePreExecution_blocksForbiddenSchemaInsideCte() {
assertThat(assertSchemaScopeBlocks(
"WITH leaked AS (SELECT ssn FROM hr.salaries) SELECT * FROM leaked"
).getErrorCode()).isEqualTo("POLICY_SCHEMA_BLOCKED");
}

@Test
void enforcePreExecution_allowsMartsQueriesWhenOtherSchemasHaveProtectedColumns() {
ConnectionChatAccessPolicyService.EffectivePolicy schemaPolicy = new ConnectionChatAccessPolicyService.EffectivePolicy(
Expand Down
Loading