From 9cb058402337e60762e090f0e747f8619012f595 Mon Sep 17 00:00:00 2001 From: sumit Date: Mon, 17 Aug 2026 20:58:25 +0530 Subject: [PATCH 1/2] fix(editor): close SQL guard bypasses in the query editor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A CHAT_EDITOR (non-admin) user could delete or overwrite every row in any table by wrapping the write in a CTE. `WITH x AS (DELETE FROM t RETURNING *) SELECT * FROM x` parses as a Select, so classification returned read-only and returned before the admin check ever ran. PostgreSQL executes data-modifying CTEs for real. Verified end to end: 3 rows -> 0, success:true, no confirmation prompt, logged as an ordinary EDITOR_QUERY_EXECUTED / SUCCESS. - classifyStatement now walks the parse tree for data-modifying CTEs and SELECT ... INTO (detectSelectWrite), with a text backstop (detectHiddenWrite) so an unparseable variant fails closed instead of reaching isReadOnlyQuery, which reports anything starting with WITH as safe. - READ_ONLY_ONLY contexts now open read-only JDBC sessions, so the database refuses the write even when classification is wrong. Classification is a parser heuristic; this is what keeps the next parser gap from being data loss. HikariCP resets the flag on return to the pool, verified, so it cannot leak into an admin's later write. - Row caps are enforced with setMaxRows instead of a `\blimit\s+\d+` text match that hit inside comments, string literals and subqueries. An inner LIMIT returned 200k rows against a 1,000 cap, into an unbounded ArrayList and an unvirtualized table. - Cancel terminates the query instead of only aborting the HTTP request, which left the statement holding one of the pool's 10 connections. The client sends an executionId, RunningQueryRegistry maps it to the backend session pid, and the new cancel endpoint kills exactly that session, scoped to the connection and the user who started it. The previous UI behavior killed *every* active query on the connection, including other users' work. - Fix pg_terminate_backend binding: setLong sent bigint, so PostgreSQL found no matching overload and every kill failed, including the Active Queries screen's own button. - Editor timeout 600s -> 240s, under nginx's 300s proxy_read_timeout, so a slow query reports a real error rather than an opaque 504 while still running. - Rate-limit /api/connections/*/query (30r/m + burst 20, 429 on reject). QueryExecutionPolicyServiceTest stubbed isReadOnlyQuery to always return false — the opposite of what the shipped providers do for WITH — so it asserted behavior no deployment had, and withInsert_isTreatedAsMutation passed *because* of the stub. It now uses a real MySQLQueryExecutionProvider, plus 12 regression tests covering each bypass and the reads that must keep working. Co-Authored-By: Claude Opus 5 --- CLAUDE.md | 54 +++++++ .../dbaagent/controller/SchemaController.java | 61 ++++++++ .../java/com/dbaagent/model/QueryRequest.java | 3 + .../java/com/dbaagent/model/QueryResult.java | 19 +++ .../provider/api/QueryExecutionProvider.java | 12 ++ .../mysql/MySQLQueryExecutionProvider.java | 5 + .../PostgresQueryExecutionProvider.java | 5 + .../dbaagent/service/ActiveQueryService.java | 9 +- .../service/QueryExecutionPolicyService.java | 139 +++++++++++++++++- .../service/QueryExecutorService.java | 94 ++++++++++-- .../service/RunningQueryRegistry.java | 66 +++++++++ .../service/UserDataAccessPolicyService.java | 3 +- .../QueryExecutionPolicyServiceTest.java | 128 +++++++++++++++- .../QueryExecutorServiceCounterTest.java | 3 +- .../service/QueryExecutorServiceTest.java | 3 +- docker/nginx/default.conf | 27 ++++ src/components/tabs/Core/SqlRunnerTab.js | 42 +++--- src/lib/api/client.js | 12 ++ 18 files changed, 649 insertions(+), 36 deletions(-) create mode 100644 backend/src/main/java/com/dbaagent/service/RunningQueryRegistry.java diff --git a/CLAUDE.md b/CLAUDE.md index bc7dfc3..61b8417 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -262,6 +262,60 @@ broken. Assert the *outcome*, never the attempt: `No databases connected yet` because one `catch` covered both the connection fetch and decorative extras. An unreachable host must never look like an empty account. +### SQL Editor Guard Rules + +The Editor (`EditorSection` → `SqlRunnerTab` → `POST /connections/{id}/query` → +`QueryExecutionPolicyService` → `QueryExecutorService`) is the only surface where a +user submits arbitrary SQL. Everything below was a live bug, verified by executing +it against a real database — not a theoretical hardening pass. + +- **Never classify SQL by its leading keyword alone.** `WITH x AS (DELETE FROM t + RETURNING *) SELECT * FROM x` parses as a `Select` and *every* leading-keyword + check calls it read-only — including `isReadOnlyQuery`, which reports anything + starting with `WITH` as safe. PostgreSQL executes data-modifying CTEs for real, + so a **non-admin** wiped whole tables through the Editor with no confirmation + prompt, logged as an ordinary `EDITOR_QUERY_EXECUTED / SUCCESS`. `SELECT … INTO + newtab` is the same class of bug (it is DDL). `classifyStatement` now inspects + the parse tree (`detectSelectWrite`) **and** runs a text backstop + (`detectHiddenWrite`) so an unparseable variant fails closed instead of falling + through to the keyword path. +- **Read-only contexts open read-only JDBC sessions.** `QueryExecutorService` calls + `connection.setReadOnly(true)` whenever `mutationMode() == READ_ONLY_ONLY`, so + PostgreSQL refuses the write itself even if classification is wrong. Classification + is a parser heuristic; this is what keeps the *next* parser gap from being data + loss. A driver that rejects the hint raises rather than silently continuing + writable. HikariCP resets the flag on return to the pool (verified), so it cannot + leak into an admin's later write. +- **Row caps are enforced with `setMaxRows`, not by appending `LIMIT n`.** The old + check skipped its own LIMIT whenever the regex `\blimit\s+\d+` matched anywhere — + including inside a comment, a string literal, or a subquery. `WITH a AS (SELECT … + LIMIT 100) SELECT * FROM a` is ordinary analyst SQL and returned **200k rows** + against a 1,000 cap, straight into an unbounded `ArrayList` and then an + unvirtualized table. The SQL `LIMIT` is still appended for simple SELECTs, but + only as an optimization — correctness no longer depends on that text match. +- **Cancel must terminate the query, not just the HTTP request.** `abortController + .abort()` only closes the socket; the statement runs on holding one of the pool's + 10 connections for up to its timeout. The client now sends an `executionId`, + `RunningQueryRegistry` maps it to the backend session pid (via the dialect's + `getSessionPidQuery()`), and `POST /connections/{id}/query/{executionId}/cancel` + terminates exactly that session. The previous UI behavior was worse than nothing: + it killed **every** active query on the connection, including other users' work. + The cancel endpoint is scoped to the connection *and* the user who started the + run, so an execution id is not a kill primitive for someone else's query. +- **Keep the client timeout under the proxy's.** `docker/nginx/default.conf` gives + up at `proxy_read_timeout 300s`; the Editor used to ask for 600s, so a 6-minute + query returned an opaque 504 while still running. `QUERY_TIMEOUT_SECONDS = 240` + in `SqlRunnerTab.js` — change both together or not at all. +- **`/api/connections/*/query` is rate-limited in nginx** (`limit_req zone=sqlexec`, + 30r/m + burst 20, `429` on reject). It is the most expensive authenticated call + in the product. +- **Test the policy against the real providers.** `QueryExecutionPolicyServiceTest` + used to stub `isReadOnlyQuery` to always return `false` — the exact opposite of + what the shipped providers do for `WITH`. It asserted behavior no deployment had, + and `withInsert_isTreatedAsMutation` passed *because* of the stub. It now + constructs a real `MySQLQueryExecutionProvider`. Do not reintroduce a stubbed + dialect here; the mock is what let the blocker ship. + ### Data Model Rules - **`mcp_tokens.user_id` is a non-null FK with no cascade.** Deleting a user who holds diff --git a/backend/src/main/java/com/dbaagent/controller/SchemaController.java b/backend/src/main/java/com/dbaagent/controller/SchemaController.java index 8bab6c9..72204be 100644 --- a/backend/src/main/java/com/dbaagent/controller/SchemaController.java +++ b/backend/src/main/java/com/dbaagent/controller/SchemaController.java @@ -6,7 +6,9 @@ import com.dbaagent.service.CredentialService; import com.dbaagent.service.QueryExecutionContext; import com.dbaagent.service.QueryExecutionPolicyException; +import com.dbaagent.service.ActiveQueryService; import com.dbaagent.service.QueryExecutorService; +import com.dbaagent.service.RunningQueryRegistry; import com.dbaagent.service.SqlExecutionAuditService; import com.dbaagent.service.UserDataAccessPolicyException; import com.dbaagent.service.SchemaScannerService; @@ -36,6 +38,8 @@ public class SchemaController { private final QueryExecutorService queryExecutorService; private final AccessControlService accessControlService; private final SqlExecutionAuditService sqlExecutionAuditService; + private final RunningQueryRegistry runningQueryRegistry; + private final ActiveQueryService activeQueryService; @PostMapping("/scan") public ResponseEntity> scanSchema(@PathVariable String connectionId) { @@ -260,6 +264,63 @@ public ResponseEntity> executeQuery( } } + /** + * Terminates a query this caller started but abandoned. Aborting the HTTP + * request only closes the socket — the statement keeps running and holds a + * pooled connection until it finishes, so the client must ask for it to stop. + */ + @PostMapping("/query/{executionId}/cancel") + public ResponseEntity> cancelQuery( + @PathVariable String connectionId, + @PathVariable String executionId) { + Map response = new HashMap<>(); + try { + if (!credentialService.connectionExists(connectionId)) { + response.put("success", false); + response.put("message", "Connection not found"); + return ResponseEntity.status(HttpStatus.NOT_FOUND).body(response); + } + accessControlService.assertCanUseChatEditor(connectionId); + + var running = runningQueryRegistry.find(executionId); + if (running.isEmpty()) { + // Already finished, or never started. Nothing to cancel. + response.put("success", true); + response.put("cancelled", false); + response.put("message", "Query is no longer running"); + return ResponseEntity.ok(response); + } + + RunningQueryRegistry.RunningQuery target = running.get(); + // The execution id is a bearer token for a kill: scope it to this + // connection and to the user who started it, so one caller cannot + // terminate another's query by guessing an id. + String currentUser = accessControlService.getCurrentUsername(); + if (!connectionId.equals(target.connectionId()) + || (target.username() != null && currentUser != null && !target.username().equals(currentUser))) { + response.put("success", false); + response.put("message", "Query not found for this connection"); + return ResponseEntity.status(HttpStatus.NOT_FOUND).body(response); + } + + activeQueryService.killQuery(connectionId, target.sessionPid()); + runningQueryRegistry.unregister(executionId); + response.put("success", true); + response.put("cancelled", true); + response.put("message", "Query cancelled"); + return ResponseEntity.ok(response); + } catch (ResponseStatusException e) { + response.put("success", false); + response.put("message", e.getReason()); + return ResponseEntity.status(e.getStatusCode()).body(response); + } catch (Exception e) { + log.warn("Failed to cancel query {} on connection {}: {}", executionId, connectionId, e.getMessage()); + response.put("success", false); + response.put("message", "Failed to cancel query: " + e.getMessage()); + return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body(response); + } + } + // `{tableName:.+}` keeps schema-qualified ids (`crm.orders`) as one segment. @GetMapping("/tables/{tableName:.+}/indexes") public ResponseEntity> getTableIndexes( diff --git a/backend/src/main/java/com/dbaagent/model/QueryRequest.java b/backend/src/main/java/com/dbaagent/model/QueryRequest.java index 78a673e..a825075 100644 --- a/backend/src/main/java/com/dbaagent/model/QueryRequest.java +++ b/backend/src/main/java/com/dbaagent/model/QueryRequest.java @@ -11,6 +11,9 @@ public class QueryRequest { private Integer timeoutSeconds; // Optional per-query timeout override (null = use server default) private QueryExecutionOrigin executionOrigin = QueryExecutionOrigin.INTERNAL; private Boolean mutationConfirmed = Boolean.FALSE; + // Client-generated id for this run, used to cancel the query if the caller + // gives up before it finishes. + private String executionId; public QueryRequest(String query, Integer limit, Integer timeoutSeconds) { this.query = query; diff --git a/backend/src/main/java/com/dbaagent/model/QueryResult.java b/backend/src/main/java/com/dbaagent/model/QueryResult.java index a815fc7..f521c1e 100644 --- a/backend/src/main/java/com/dbaagent/model/QueryResult.java +++ b/backend/src/main/java/com/dbaagent/model/QueryResult.java @@ -17,4 +17,23 @@ public class QueryResult { private Boolean isLimited; // True when a row limit was applied to this query private Long executionTimeMs; private String query; + // Server-side session id this query ran on, so a client that gives up can ask + // the backend to terminate it instead of leaving it holding a connection. + private String sessionPid; + + /** + * Kept so adding {@code sessionPid} did not break every positional caller. + * Prefer the setter for new code — this class is a response DTO that grows. + */ + public QueryResult( + List columns, + List> rows, + Integer rowCount, + Long totalRowCount, + Boolean isLimited, + Long executionTimeMs, + String query + ) { + this(columns, rows, rowCount, totalRowCount, isLimited, executionTimeMs, query, null); + } } diff --git a/backend/src/main/java/com/dbaagent/provider/api/QueryExecutionProvider.java b/backend/src/main/java/com/dbaagent/provider/api/QueryExecutionProvider.java index c0780e3..9841ce0 100644 --- a/backend/src/main/java/com/dbaagent/provider/api/QueryExecutionProvider.java +++ b/backend/src/main/java/com/dbaagent/provider/api/QueryExecutionProvider.java @@ -51,4 +51,16 @@ public interface QueryExecutionProvider { * @return The query signature */ String generateQuerySignature(String query); + + /** + * SQL returning the server-side session id of the current connection, in the + * form accepted by this dialect's kill statement. Callers use it to cancel a + * still-running query after the client has gone away. + * + * @return a single-column, single-row query, or null when the dialect has no + * session identifier we can act on + */ + default String getSessionPidQuery() { + return null; + } } diff --git a/backend/src/main/java/com/dbaagent/provider/mysql/MySQLQueryExecutionProvider.java b/backend/src/main/java/com/dbaagent/provider/mysql/MySQLQueryExecutionProvider.java index 7350e86..92d82d5 100644 --- a/backend/src/main/java/com/dbaagent/provider/mysql/MySQLQueryExecutionProvider.java +++ b/backend/src/main/java/com/dbaagent/provider/mysql/MySQLQueryExecutionProvider.java @@ -123,4 +123,9 @@ public String generateQuerySignature(String query) { return String.valueOf(normalized.hashCode()); } } + + @Override + public String getSessionPidQuery() { + return "SELECT CONNECTION_ID()"; + } } diff --git a/backend/src/main/java/com/dbaagent/provider/postgres/PostgresQueryExecutionProvider.java b/backend/src/main/java/com/dbaagent/provider/postgres/PostgresQueryExecutionProvider.java index ac7e657..a2cf814 100644 --- a/backend/src/main/java/com/dbaagent/provider/postgres/PostgresQueryExecutionProvider.java +++ b/backend/src/main/java/com/dbaagent/provider/postgres/PostgresQueryExecutionProvider.java @@ -127,4 +127,9 @@ public String generateQuerySignature(String query) { return String.valueOf(normalized.hashCode()); } } + + @Override + public String getSessionPidQuery() { + return "SELECT pg_backend_pid()"; + } } diff --git a/backend/src/main/java/com/dbaagent/service/ActiveQueryService.java b/backend/src/main/java/com/dbaagent/service/ActiveQueryService.java index 49c2ff4..4c4ea00 100644 --- a/backend/src/main/java/com/dbaagent/service/ActiveQueryService.java +++ b/backend/src/main/java/com/dbaagent/service/ActiveQueryService.java @@ -203,8 +203,15 @@ public void killQuery(String connectionId, String pid) { String dbType = providerRegistry.getCanonicalName(connRequest.getDbType()); if ("postgres".equals(dbType)) { + // pg_terminate_backend takes an integer; binding a long makes + // the driver send bigint and PostgreSQL then finds no matching + // overload ("function pg_terminate_backend(bigint) does not + // exist"), so every kill failed. + if (backendPid > Integer.MAX_VALUE || backendPid < Integer.MIN_VALUE) { + throw new IllegalArgumentException("Not a valid PostgreSQL backend pid: " + backendPid); + } try (PreparedStatement ps = conn.prepareStatement("SELECT pg_terminate_backend(?)")) { - ps.setLong(1, backendPid); + ps.setInt(1, (int) backendPid); ps.execute(); } } else if ("mysql".equals(dbType)) { diff --git a/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java b/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java index 32127d8..6fd47cf 100644 --- a/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java +++ b/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java @@ -19,7 +19,12 @@ import net.sf.jsqlparser.statement.drop.Drop; import net.sf.jsqlparser.statement.insert.Insert; import net.sf.jsqlparser.statement.merge.Merge; +import net.sf.jsqlparser.schema.Table; +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.statement.truncate.Truncate; import net.sf.jsqlparser.statement.upsert.Upsert; import net.sf.jsqlparser.statement.update.Update; @@ -51,6 +56,17 @@ public class QueryExecutionPolicyService { Pattern.CASE_INSENSITIVE ); + // Text-level backstops for writes that hide inside a statement which reads as + // a SELECT. Applied to SQL with literals and comments already stripped. + private static final Pattern CTE_WRITE_PATTERN = Pattern.compile( + "\\bAS\\s*\\(\\s*(INSERT|UPDATE|DELETE|MERGE)\\b", + Pattern.CASE_INSENSITIVE + ); + private static final Pattern SELECT_INTO_PATTERN = Pattern.compile( + "\\bSELECT\\b[\\s\\S]*?\\bINTO\\s+(?!STRICT\\b)[\"`\\w]", + Pattern.CASE_INSENSITIVE + ); + private final DatabaseProviderRegistry providerRegistry; public PolicyDecision enforce( @@ -163,9 +179,23 @@ private StatementClassification classifyStatement(String statement, QueryExecuti return new StatementClassification("USE", false, false, false, false, true); } + // A pre-parse scan runs first so a statement the parser mis-models — or + // rejects outright — still fails closed. `WITH ... DELETE` parses cleanly + // as a Select, and the keyword fallback below treats anything starting + // with WITH as read-only, so without this a write reaches the database + // classified as a read. + String hiddenWrite = detectHiddenWrite(trimmed); + try { Statement parsed = CCJSqlParserUtil.parse(trimmed); - if (parsed instanceof Select) { + if (parsed instanceof Select select) { + String writeKind = hiddenWrite != null ? hiddenWrite : detectSelectWrite(select); + if (writeKind != null) { + // requiresWhereClause=false: the WHERE guard reads the top-level + // statement, which is a SELECT here. A data-modifying CTE always + // needs explicit confirmation instead. + return new StatementClassification(writeKind, false, true, false, true, false); + } return new StatementClassification("SELECT", true, false, false, false, false); } if (parsed instanceof ExplainStatement) { @@ -220,6 +250,12 @@ private StatementClassification classifyStatement(String statement, QueryExecuti return new StatementClassification(explainWrappedMutation, false, true, requiresWhere, hasWhere, false); } + // The parser failed or produced a type we don't model. `isReadOnlyQuery` + // only looks at the leading keyword, so a hidden write must veto it. + if (hiddenWrite != null) { + return new StatementClassification(hiddenWrite, false, true, false, true, false); + } + String queryType = QueryNormalizer.detectQueryType(trimmed); boolean readOnly = executionProvider.isReadOnlyQuery(trimmed); boolean mutating = !readOnly && !"UNKNOWN".equalsIgnoreCase(queryType); @@ -228,6 +264,107 @@ private StatementClassification classifyStatement(String statement, QueryExecuti return new StatementClassification(queryType, readOnly, mutating, requiresWhere, hasWhere, false); } + /** + * Detects a write hidden inside a statement that parses as a {@link Select}: + * a data-modifying CTE ({@code WITH x AS (DELETE ...) SELECT ...}, which + * PostgreSQL executes for real) or {@code SELECT ... INTO newtable}, which + * creates a table. Returns the write's query type, or null when the select + * really is read-only. + */ + private String detectSelectWrite(Select select) { + String fromCte = detectWriteInWithItems(select.getWithItemsList(), 0); + if (fromCte != null) { + return fromCte; + } + return detectSelectInto(select); + } + + private String detectWriteInWithItems(List> withItems, int depth) { + // CTEs can nest; bound the walk so a pathological query can't spin here. + if (withItems == null || withItems.isEmpty() || depth > 10) { + return null; + } + for (WithItem item : withItems) { + if (item == null) { + continue; + } + if (item.getDelete() != null) return "DELETE (in CTE)"; + if (item.getUpdate() != null) return "UPDATE (in CTE)"; + if (item.getInsert() != null) return "INSERT (in CTE)"; + + // getSelect() throws when the item holds a non-select statement, which + // the getters above already covered. A select here may carry its own + // WITH list, so recurse to catch a write nested one level deeper. + try { + Select inner = item.getSelect(); + if (inner != null) { + String innerWrite = detectSelectWriteQuietly(inner, depth + 1); + if (innerWrite != null) { + return innerWrite; + } + } + } catch (Exception ignored) { + // Not a select-bearing CTE; nothing further to inspect. + } + } + return null; + } + + private String detectSelectWriteQuietly(Select select, int depth) { + if (depth > 10) { + return null; + } + String fromCte = detectWriteInWithItems(select.getWithItemsList(), depth); + return fromCte != null ? fromCte : detectSelectInto(select); + } + + private String detectSelectInto(Select select) { + if (select instanceof PlainSelect plain) { + List into = plain.getIntoTables(); + if (into != null && !into.isEmpty()) { + return "SELECT INTO"; + } + } + if (select instanceof ParenthesedSelect parenthesed) { + Select inner = parenthesed.getSelect(); + return inner == null ? null : detectSelectInto(inner); + } + if (select instanceof SetOperationList setOps && setOps.getSelects() != null) { + for (Select part : setOps.getSelects()) { + String found = detectSelectInto(part); + if (found != null) { + return found; + } + } + } + return null; + } + + /** + * Text-level backstop for the same two shapes, used when the parser is not + * available or disagrees. Quoted literals and comments are stripped first so + * the word DELETE inside a string can't trigger a false positive. + */ + private String detectHiddenWrite(String sql) { + if (sql == null || sql.isBlank()) { + return null; + } + String scrubbed = stripLeadingComments(stripComments(stripQuotedLiterals(sql))); + if (!scrubbed.regionMatches(true, 0, "WITH", 0, 4)) { + return SELECT_INTO_PATTERN.matcher(scrubbed).find() ? "SELECT INTO" : null; + } + if (CTE_WRITE_PATTERN.matcher(scrubbed).find()) { + return "WRITE (in CTE)"; + } + return SELECT_INTO_PATTERN.matcher(scrubbed).find() ? "SELECT INTO" : null; + } + + private String stripComments(String sql) { + return sql + .replaceAll("(?s)/\\*.*?\\*/", " ") + .replaceAll("(?m)--[^\\n]*", " "); + } + private boolean isUseStatement(String sql) { return sql != null && USE_ANY_PATTERN.matcher(stripLeadingComments(sql)).find(); } diff --git a/backend/src/main/java/com/dbaagent/service/QueryExecutorService.java b/backend/src/main/java/com/dbaagent/service/QueryExecutorService.java index 0c14cc9..110eaa7 100644 --- a/backend/src/main/java/com/dbaagent/service/QueryExecutorService.java +++ b/backend/src/main/java/com/dbaagent/service/QueryExecutorService.java @@ -29,10 +29,34 @@ public class QueryExecutorService { private final UserDataAccessPolicyService userDataAccessPolicyService; private final com.dbaagent.service.telemetry.TelemetryCounters telemetryCounters; private final KeyColumnAnalysisRepository keyColumnAnalysisRepository; + private final RunningQueryRegistry runningQueryRegistry; @Value("${db.query-timeout-seconds:30}") private int queryTimeoutSeconds; + /** + * Reads the connection's server-side session id via the dialect's own query. + * Never fatal: without it the query simply cannot be cancelled early. + */ + private String resolveSessionPid(Connection connection, String dbType) { + String pidQuery; + try { + pidQuery = providerRegistry.getDialect(dbType).queryExecution().getSessionPidQuery(); + } catch (Exception e) { + return null; + } + if (pidQuery == null) { + return null; + } + try (Statement stmt = connection.createStatement(); + ResultSet rs = stmt.executeQuery(pidQuery)) { + return rs.next() ? rs.getString(1) : null; + } catch (SQLException e) { + log.debug("Could not resolve session pid: {}", e.getMessage()); + return null; + } + } + public List getDatabaseObjects(String connectionId) throws SQLException { // Try to get from cache, but handle cache failures gracefully Cache cache = null; @@ -547,6 +571,23 @@ public QueryResult executeQuery( ); try (Connection connection = connectionService.getConnection(connectionId, connRequest)) { + // Defense in depth: when the caller may not mutate, ask the database + // to enforce that too. Classification is a parser heuristic, so a + // read-only session is what keeps a future gap from becoming data + // loss — PostgreSQL then refuses the write itself. + boolean readOnlySession = + executionContext.mutationMode() == QueryExecutionContext.MutationMode.READ_ONLY_ONLY; + if (readOnlySession) { + try { + connection.setReadOnly(true); + } catch (SQLException e) { + // A driver that refuses the hint must not silently downgrade + // to a writable session. + throw new SQLException( + "Could not open a read-only database session for this connection: " + e.getMessage(), e); + } + } + // Per-query timeout override: use the request's timeout if provided, // otherwise fall back to the server default (db.query-timeout-seconds). int maxAllowedTimeoutSeconds = 600; // 10 minutes ceiling @@ -554,6 +595,14 @@ public QueryResult executeQuery( ? Math.min(queryRequest.getTimeoutSeconds(), maxAllowedTimeoutSeconds) : queryTimeoutSeconds; + // Publish the server-side session id so a client that abandons this + // request can terminate the query rather than leaving it to run out + // its timeout holding a pooled connection. + String sessionPid = resolveSessionPid(connection, connRequest.getDbType()); + result.setSessionPid(sessionPid); + runningQueryRegistry.register( + queryRequest.getExecutionId(), connectionId, sessionPid, executionContext.actorUsername()); + // Execute all preamble statements (SET @var, USE db, etc.) on the same connection // so that session variables are in scope for the final SELECT. for (int si = 0; si < statements.size() - 1; si++) { @@ -573,15 +622,23 @@ public QueryResult executeQuery( stmt.setQueryTimeout(effectiveTimeout); } - // Apply limit if specified + // Row cap. setMaxRows is the enforcement: the driver stops handing + // back rows regardless of the query's shape, so a LIMIT inside a CTE, + // a comment, or a string literal can no longer defeat the cap. We ask + // for one extra row purely to detect truncation. + // + // Appending " LIMIT n" is kept as an optimization for simple SELECTs + // (it lets the database stop early rather than stream rows we drop), + // but correctness no longer depends on that textual check. + boolean limitRequested = queryRequest.getLimit() != null && queryRequest.getLimit() > 0; + int rowLimit = limitRequested ? queryRequest.getLimit() : 0; boolean limitApplied = false; String baseQuery = null; // original query without limit, for COUNT - if (queryRequest.getLimit() != null && queryRequest.getLimit() > 0) { + if (limitRequested) { + stmt.setMaxRows(rowLimit + 1); + String queryTrimmed = finalQuery.trim(); String queryLower = queryTrimmed.toLowerCase(); - // Only add LIMIT to SELECT queries, not to SHOW/DESCRIBE/EXPLAIN commands. - // Use regex to check for actual LIMIT clause (avoids false positives from - // column names, comments, or identifiers that contain the word "limit"). boolean isSelect = queryLower.startsWith("select"); boolean alreadyHasLimit = queryLower.matches("(?s).*\\blimit\\s+\\d+.*"); if (isSelect && !alreadyHasLimit) { @@ -589,7 +646,10 @@ public QueryResult executeQuery( if (baseQuery.endsWith(";")) { baseQuery = baseQuery.substring(0, baseQuery.length() - 1); } - finalQuery = baseQuery + " LIMIT " + queryRequest.getLimit(); + // Ask for one more than we will return, matching setMaxRows, + // so an exactly-at-the-cap result is still detectable as + // truncated rather than looking complete. + finalQuery = baseQuery + " LIMIT " + (rowLimit + 1); limitApplied = true; } } @@ -620,9 +680,16 @@ public QueryResult executeQuery( } result.setColumns(columns); - // Get rows + // Get rows. setMaxRows let one extra row through so we can + // tell "exactly at the cap" from "there were more"; it is + // counted but never returned. List> rows = new ArrayList<>(); + boolean truncated = false; while (rs.next()) { + if (limitRequested && rows.size() >= rowLimit) { + truncated = true; + break; + } List row = new ArrayList<>(); for (int i = 1; i <= columnCount; i++) { Object value = rs.getObject(i); @@ -632,11 +699,14 @@ public QueryResult executeQuery( } result.setRows(rows); result.setRowCount(rows.size()); - - // If a limit was applied, mark result as limited and fetch total row count. - // The COUNT query runs on the same connection so @session variables are in scope. - if (limitApplied && baseQuery != null) { + if (truncated) { result.setIsLimited(true); + } + + // Fetch the true total only when we appended the LIMIT + // ourselves, since that is the one case where baseQuery is + // a valid standalone statement to wrap in COUNT(*). + if (truncated && limitApplied && baseQuery != null) { try (Statement countStmt = connection.createStatement()) { countStmt.setQueryTimeout(effectiveTimeout > 0 ? effectiveTimeout : 60); // Strip trailing ORDER BY before wrapping in COUNT subquery, @@ -661,6 +731,8 @@ public QueryResult executeQuery( result.setRows(Arrays.asList(Arrays.asList(updateCount))); } } + } finally { + runningQueryRegistry.unregister(queryRequest.getExecutionId()); } result.setExecutionTimeMs(System.currentTimeMillis() - startTime); diff --git a/backend/src/main/java/com/dbaagent/service/RunningQueryRegistry.java b/backend/src/main/java/com/dbaagent/service/RunningQueryRegistry.java new file mode 100644 index 0000000..ff1f805 --- /dev/null +++ b/backend/src/main/java/com/dbaagent/service/RunningQueryRegistry.java @@ -0,0 +1,66 @@ +package com.dbaagent.service; + +import lombok.extern.slf4j.Slf4j; +import org.springframework.stereotype.Service; + +import java.time.Duration; +import java.time.Instant; +import java.util.Iterator; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Tracks in-flight queries by a client-supplied execution id so a client that + * gives up can terminate the query it started. Aborting the HTTP request only + * closes the socket; without this the statement keeps running and holds a pooled + * connection until it finishes on its own. + * + *

Entries are short-lived: registered when execution begins and removed in a + * finally block. The sweep on write is a backstop for a process that died between + * those two points. + */ +@Service +@Slf4j +public class RunningQueryRegistry { + + private static final Duration MAX_AGE = Duration.ofHours(1); + private static final int SWEEP_THRESHOLD = 256; + + public record RunningQuery(String connectionId, String sessionPid, String username, Instant startedAt) {} + + private final Map running = new ConcurrentHashMap<>(); + + public void register(String executionId, String connectionId, String sessionPid, String username) { + if (executionId == null || executionId.isBlank() || sessionPid == null || sessionPid.isBlank()) { + return; + } + if (running.size() > SWEEP_THRESHOLD) { + sweepExpired(); + } + running.put(executionId, new RunningQuery(connectionId, sessionPid, username, Instant.now())); + } + + public void unregister(String executionId) { + if (executionId != null && !executionId.isBlank()) { + running.remove(executionId); + } + } + + public Optional find(String executionId) { + if (executionId == null || executionId.isBlank()) { + return Optional.empty(); + } + return Optional.ofNullable(running.get(executionId)); + } + + private void sweepExpired() { + Instant cutoff = Instant.now().minus(MAX_AGE); + Iterator> it = running.entrySet().iterator(); + while (it.hasNext()) { + if (it.next().getValue().startedAt().isBefore(cutoff)) { + it.remove(); + } + } + } +} diff --git a/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java b/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java index 6ef3e66..a65a9f6 100644 --- a/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java +++ b/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java @@ -216,7 +216,8 @@ public QueryResult redactResult( result.getTotalRowCount(), result.getIsLimited(), result.getExecutionTimeMs(), - result.getQuery() + result.getQuery(), + result.getSessionPid() ); logPolicyEvent(SecurityEventType.CHAT_ACCESS_POLICY_REDACTED, executionContext.actorUsername(), connectionId, "result_redacted", Map.of( "query", truncate(result.getQuery()), diff --git a/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java b/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java index ef78d7f..659cfb5 100644 --- a/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java +++ b/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java @@ -4,6 +4,7 @@ import com.dbaagent.provider.DatabaseProviderRegistry; import com.dbaagent.provider.api.DatabaseDialect; import com.dbaagent.provider.api.QueryExecutionProvider; +import com.dbaagent.provider.mysql.MySQLQueryExecutionProvider; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -23,16 +24,20 @@ class QueryExecutionPolicyServiceTest { @Mock private DatabaseProviderRegistry providerRegistry; @Mock private DatabaseDialect databaseDialect; - @Mock private QueryExecutionProvider queryExecutionProvider; private QueryExecutionPolicyService service; @BeforeEach void setUp() { + // The real provider, not a stub. A stubbed isReadOnlyQuery() that always + // answered "false" made these tests assert the opposite of production: + // MySQLQueryExecutionProvider reports anything starting with WITH as + // read-only, which is how `WITH x AS (DELETE ...) SELECT` reached the + // database classified as a read. + QueryExecutionProvider realProvider = new MySQLQueryExecutionProvider(); when(providerRegistry.getDialect(anyString())).thenReturn(databaseDialect); - when(databaseDialect.queryExecution()).thenReturn(queryExecutionProvider); + when(databaseDialect.queryExecution()).thenReturn(realProvider); lenient().when(providerRegistry.getCanonicalName(anyString())).thenReturn("mysql"); - lenient().when(queryExecutionProvider.isReadOnlyQuery(anyString())).thenReturn(false); service = new QueryExecutionPolicyService(providerRegistry); } @@ -296,4 +301,121 @@ void internalMutation_isStillAllowed() { assertThat(decision.mutating()).isTrue(); assertThat(decision.primaryQueryType()).isEqualTo("TRUNCATE"); } + + // --- Writes hidden inside a statement that reads as a SELECT --------------- + // PostgreSQL executes data-modifying CTEs, and the parser models them as a + // Select. Each of these deleted or rewrote a whole table from a non-admin + // account before the classifier learned to look inside. + + private QueryExecutionPolicyException assertBlockedForViewer(String sql) { + return assertThrows( + QueryExecutionPolicyException.class, + () -> service.enforce( + new QueryRequest(sql, null, null), + QueryExecutionContext.editor("viewer", false, false), + "postgresql" + ) + ); + } + + @Test + void cteDelete_isBlockedForNonAdmin() { + QueryExecutionPolicyException e = + assertBlockedForViewer("WITH x AS (DELETE FROM orders RETURNING *) SELECT * FROM x"); + assertThat(e.getErrorCode()).isEqualTo(QueryExecutionPolicyException.EDITOR_MUTATION_FORBIDDEN); + } + + @Test + void cteUpdate_isBlockedForNonAdmin() { + assertBlockedForViewer("WITH u AS (UPDATE orders SET total = 0 RETURNING *) SELECT * FROM u"); + } + + @Test + void cteInsert_isBlockedForNonAdmin() { + assertBlockedForViewer("WITH i AS (INSERT INTO audit(id) VALUES (1) RETURNING *) SELECT * FROM i"); + } + + @Test + void cteWriteInLaterPosition_isBlockedForNonAdmin() { + assertBlockedForViewer( + "WITH a AS (SELECT 1), b AS (DELETE FROM orders RETURNING *) SELECT * FROM a"); + } + + @Test + void nestedCteWrite_isBlockedForNonAdmin() { + assertBlockedForViewer( + "WITH o AS (WITH i AS (DELETE FROM orders RETURNING *) SELECT * FROM i) SELECT * FROM o"); + } + + @Test + void selectInto_isBlockedForNonAdmin() { + assertBlockedForViewer("SELECT * INTO exfiltrated FROM customers"); + } + + @Test + void cteWrite_requiresConfirmationForAdmin() { + QueryExecutionPolicyException e = assertThrows( + QueryExecutionPolicyException.class, + () -> service.enforce( + new QueryRequest("WITH x AS (DELETE FROM orders RETURNING *) SELECT * FROM x", null, null), + QueryExecutionContext.editor("admin", true, false), + "postgresql" + ) + ); + assertThat(e.getErrorCode()) + .isEqualTo(QueryExecutionPolicyException.EDITOR_MUTATION_CONFIRMATION_REQUIRED); + } + + @Test + void cteWrite_isAllowedForConfirmedAdmin() { + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("WITH x AS (DELETE FROM orders RETURNING *) SELECT * FROM x", null, null), + QueryExecutionContext.editor("admin", true, true), + "postgresql" + ); + assertThat(decision.mutating()).isTrue(); + } + + // --- The guard must not swallow legitimate reads -------------------------- + + @Test + void readOnlyCte_remainsAllowedForNonAdmin() { + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("WITH recent AS (SELECT * FROM orders LIMIT 10) SELECT * FROM recent", null, null), + QueryExecutionContext.editor("viewer", false, false), + "postgresql" + ); + assertThat(decision.mutating()).isFalse(); + assertThat(decision.primaryQueryType()).isEqualTo("SELECT"); + } + + @Test + void writeKeywordInsideStringLiteral_isStillAReadForNonAdmin() { + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("SELECT 'WITH x AS (DELETE FROM t)' AS example", null, null), + QueryExecutionContext.editor("viewer", false, false), + "postgresql" + ); + assertThat(decision.mutating()).isFalse(); + } + + @Test + void writeKeywordInsideComment_isStillAReadForNonAdmin() { + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("SELECT 1 -- WITH x AS (DELETE FROM t)\n", null, null), + QueryExecutionContext.editor("viewer", false, false), + "postgresql" + ); + assertThat(decision.mutating()).isFalse(); + } + + @Test + void insertIntoSelect_isStillClassifiedAsInsert() { + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("INSERT INTO archive SELECT * FROM orders", null, null), + QueryExecutionContext.editor("admin", true, true), + "postgresql" + ); + assertThat(decision.primaryQueryType()).isEqualTo("INSERT"); + } } diff --git a/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceCounterTest.java b/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceCounterTest.java index 87b5d3e..ed23806 100644 --- a/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceCounterTest.java +++ b/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceCounterTest.java @@ -58,7 +58,8 @@ void setUp() { queryExecutionPolicyService, userDataAccessPolicyService, telemetryCounters, - keyColumnAnalysisRepository + keyColumnAnalysisRepository, + new RunningQueryRegistry() ); } diff --git a/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceTest.java b/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceTest.java index c23d805..a26d479 100644 --- a/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceTest.java +++ b/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceTest.java @@ -56,7 +56,8 @@ void setUp() { queryExecutionPolicyService, userDataAccessPolicyService, telemetryCounters, - keyColumnAnalysisRepository + keyColumnAnalysisRepository, + new RunningQueryRegistry() ); } diff --git a/docker/nginx/default.conf b/docker/nginx/default.conf index e69243f..b4e2f7d 100644 --- a/docker/nginx/default.conf +++ b/docker/nginx/default.conf @@ -1,3 +1,9 @@ +# Ad-hoc SQL execution is the most expensive thing an authenticated user can do: +# each request can occupy a database connection for minutes. Cap the sustained +# rate per client, with a burst so a normal person clicking Run repeatedly is +# never throttled. +limit_req_zone $binary_remote_addr zone=sqlexec:10m rate=30r/m; + server { listen 80; server_name _; @@ -26,6 +32,27 @@ server { add_header Expires "0"; } + # Ad-hoc SQL execution, rate-limited. Declared before the general /api/ rule + # so it wins for this path; everything else is identical to it. + # `nodelay` lets the burst through immediately rather than queueing it, so + # normal interactive use is unaffected and only sustained abuse is rejected. + location ~ ^/api/connections/[^/]+/query$ { + limit_req zone=sqlexec burst=20 nodelay; + limit_req_status 429; + + proxy_pass http://backend:8080$request_uri; + proxy_http_version 1.1; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_read_timeout 300s; + proxy_send_timeout 300s; + proxy_connect_timeout 10s; + proxy_buffering off; + proxy_cache off; + } + # API proxy to the backend service # Allows the frontend to call /api/... without CORS issues location /api/ { diff --git a/src/components/tabs/Core/SqlRunnerTab.js b/src/components/tabs/Core/SqlRunnerTab.js index 9eedffd..376ea19 100644 --- a/src/components/tabs/Core/SqlRunnerTab.js +++ b/src/components/tabs/Core/SqlRunnerTab.js @@ -37,7 +37,6 @@ import { queryPerformanceAPI, explainAPI, brainAPI, - activeQueryAPI, } from "@/lib/api/client"; import ExplainAnalysisPanel from "./ExplainAnalysisPanel"; import QueryOptimizePanel from "./QueryOptimizePanel"; @@ -51,6 +50,10 @@ import { connectionHasMultipleSchemas, } from "@/lib/schemaNames"; +// Kept below the 300s proxy_read_timeout in docker/nginx/default.conf so a slow +// query fails with a real message rather than an opaque 504. +const QUERY_TIMEOUT_SECONDS = 240; + // Constants for diagram layout const DIAGRAM_NODE_WIDTH = 240; const DIAGRAM_NODE_HEIGHT = 120; @@ -294,6 +297,7 @@ export default function SqlRunnerTab({ connectionId }) { const autocompleteRegisteredRef = useRef(false); const dbObjectsRef = useRef([]); const abortControllerRef = useRef(null); + const executionIdRef = useRef(null); // Note: savedQueriesPanelRef kept for potential future use, but panel UI is now in modal const savedQueriesPanelRef = useRef(null); const hasRowCount = (value) => value !== null && value !== undefined; @@ -1256,19 +1260,27 @@ export default function SqlRunnerTab({ connectionId }) { const abortController = new AbortController(); abortControllerRef.current = abortController; + // Identifies this run so cancelling can terminate it on the database. + const executionId = + globalThis.crypto?.randomUUID?.() ?? + `exec-${Date.now()}-${Math.random().toString(16).slice(2)}`; + executionIdRef.current = executionId; try { - // Use extended timeout (10 min) for SQL editor queries to support long-running queries. - // The server default is 30s which is too short for analytical queries. + // The server default of 30s is too short for analytical queries, but the + // nginx proxy in front of the API gives up at 300s (docker/nginx/default.conf). + // Staying under that means a slow query surfaces as a real error here + // instead of an opaque gateway timeout. const response = await queryAPI.executeQuery( connectionId, trimmedQuery, 1000, - 600, + QUERY_TIMEOUT_SECONDS, abortController.signal, { executionOrigin: "EDITOR", mutationConfirmed, + executionId, }, ); @@ -1344,29 +1356,25 @@ export default function SqlRunnerTab({ connectionId }) { } } finally { abortControllerRef.current = null; + executionIdRef.current = null; setIsRunning(false); } }; const handleStopQuery = () => { + const executionId = executionIdRef.current; if (abortControllerRef.current) { abortControllerRef.current.abort(); } setIsRunning(false); setError(null); - // Best-effort: capture active queries and kill any running ones on the DB side. - // Fire-and-forget — don't block the UI on this. - if (connectionId) { - activeQueryAPI.capture(connectionId) - .then((queries) => { - const running = queries.filter((q) => q.state === "active" || q.state === "running"); - running.forEach((q) => { - if (q.pid) { - activeQueryAPI.kill(connectionId, q.pid).catch(() => {}); - } - }); - }) - .catch(() => {}); + // Aborting above only drops the HTTP response; the statement keeps running + // and holds a pooled connection. Ask the server to terminate this specific + // execution. Previously this killed every active query on the connection, + // which could take out other users' work and DeepSQL's own background jobs. + if (connectionId && executionId) { + executionIdRef.current = null; + queryAPI.cancelQuery(connectionId, executionId).catch(() => {}); } }; diff --git a/src/lib/api/client.js b/src/lib/api/client.js index 9ab49ee..2694008 100644 --- a/src/lib/api/client.js +++ b/src/lib/api/client.js @@ -1386,6 +1386,11 @@ export const queryAPI = { if (requestOptions.mutationConfirmed != null) { payload.mutationConfirmed = requestOptions.mutationConfirmed; } + // Lets the caller cancel this run server-side; aborting the request alone + // leaves the query running on the database. + if (requestOptions.executionId) { + payload.executionId = requestOptions.executionId; + } // Use a longer axios timeout for queries with extended timeouts const axiosTimeout = timeoutSeconds != null @@ -1402,6 +1407,13 @@ export const queryAPI = { return response.data; }, + cancelQuery: async (connectionId, executionId) => { + const response = await apiClient.post( + `/api/connections/${connectionId}/query/${encodeURIComponent(executionId)}/cancel`, + ); + return response.data; + }, + getTableIndexes: async (connectionId, tableName) => { // Encode so schema-qualified ids (`crm.orders`) survive the path segment. const tableId = encodeURIComponent(String(tableName || "")); From a258b5036b303d11cae4fbeebed215d845e823bb Mon Sep 17 00:00:00 2001 From: sumit Date: Tue, 18 Aug 2026 17:18:56 +0530 Subject: [PATCH 2/2] fix(security): remove ReDoS in the Editor SQL guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeQL flagged three java/polynomial-redos alerts (140, 141, 142) on the text backstop added in 9cb0584. Both were real and reachable by any authenticated Editor user, measured on JDK 25: 224KB of repeated "SELECT " -> 43.8s of CPU in SELECT_INTO_PATTERN 96KB of repeated "a/*" -> 14.1s of CPU in stripComments The work happened in classifyStatement, before the query reached the database, so a single request stalled a request thread for the better part of a minute. The nginx sqlexec limiter (30r/m + burst 20) bounds the rate but not the per-request cost. - SELECT_INTO_PATTERN dropped its leading `\bSELECT\b[\s\S]*?` wildcard and now anchors on `\bINTO\s+...` alone. The lazy wildcard was what backtracked; it also bought nothing, since the caller already knows the statement shape. Because a bare INTO also appears in `INSERT INTO t SELECT ...`, detectHiddenWrite now consults the pattern only for statements that actually start with SELECT or WITH — otherwise an INSERT would be relabelled SELECT INTO. - stripComments replaced its two regexes with one linear scan. The block form matched a lazy wildcard between delimiters and degraded on an unterminated comment. After: both payloads classify in ~15ms, and 1.4MB of the same input stays at ~16ms — linear, not quadratic. Verified against the live endpoint: both hostile payloads return 400 in ~0.12s. stripQuotedLiterals on the same path was measured and left alone; Java matches its alternation linearly (27 chars of unterminated literal: <1ms). Adds two regression tests: an adversarial-input timing bound, and insertIntoSelect_isNotMisreadAsSelectInto covering the widened pattern. Full guard behavior re-verified live under real non-admin auth — every CTE write variant and SELECT INTO still blocked, reads and confirmed admin writes unaffected. Co-Authored-By: Claude Opus 5 --- .../service/QueryExecutionPolicyService.java | 49 +++++++++++++++---- .../QueryExecutionPolicyServiceTest.java | 37 ++++++++++++++ 2 files changed, 77 insertions(+), 9 deletions(-) diff --git a/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java b/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java index 6fd47cf..d330e6a 100644 --- a/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java +++ b/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java @@ -62,8 +62,13 @@ public class QueryExecutionPolicyService { "\\bAS\\s*\\(\\s*(INSERT|UPDATE|DELETE|MERGE)\\b", Pattern.CASE_INSENSITIVE ); + // Anchored on INTO alone, with no leading wildcard. The previous form + // (\bSELECT\b[\s\S]*?\bINTO ...) backtracked quadratically: 224KB of + // repeated "SELECT " took ~44s of CPU in the guard before the query reached + // the database. The caller already knows the statement is a SELECT, so the + // prefix bought nothing. private static final Pattern SELECT_INTO_PATTERN = Pattern.compile( - "\\bSELECT\\b[\\s\\S]*?\\bINTO\\s+(?!STRICT\\b)[\"`\\w]", + "\\bINTO\\s+(?!STRICT\\b)[\"`\\w]", Pattern.CASE_INSENSITIVE ); @@ -350,19 +355,45 @@ private String detectHiddenWrite(String sql) { return null; } String scrubbed = stripLeadingComments(stripComments(stripQuotedLiterals(sql))); - if (!scrubbed.regionMatches(true, 0, "WITH", 0, 4)) { - return SELECT_INTO_PATTERN.matcher(scrubbed).find() ? "SELECT INTO" : null; - } - if (CTE_WRITE_PATTERN.matcher(scrubbed).find()) { + boolean startsWithWith = scrubbed.regionMatches(true, 0, "WITH", 0, 4); + if (startsWithWith && CTE_WRITE_PATTERN.matcher(scrubbed).find()) { return "WRITE (in CTE)"; } - return SELECT_INTO_PATTERN.matcher(scrubbed).find() ? "SELECT INTO" : null; + // SELECT_INTO_PATTERN matches a bare INTO target, so it must only be + // consulted for statements that actually read as a SELECT — otherwise + // "INSERT INTO t SELECT ..." would be relabelled SELECT INTO. + if (startsWithWith || scrubbed.regionMatches(true, 0, "SELECT", 0, 6)) { + return SELECT_INTO_PATTERN.matcher(scrubbed).find() ? "SELECT INTO" : null; + } + return null; } + // Removes block and line comments with a single linear scan. The regex form + // matched a lazy wildcard between block-comment delimiters and backtracked + // quadratically on an unterminated block comment: 96KB of input cost ~14s + // of CPU. private String stripComments(String sql) { - return sql - .replaceAll("(?s)/\\*.*?\\*/", " ") - .replaceAll("(?m)--[^\\n]*", " "); + StringBuilder out = new StringBuilder(sql.length()); + int i = 0; + int len = sql.length(); + while (i < len) { + char c = sql.charAt(i); + if (c == '/' && i + 1 < len && sql.charAt(i + 1) == '*') { + int end = sql.indexOf("*/", i + 2); + i = end < 0 ? len : end + 2; + out.append(' '); + continue; + } + if (c == '-' && i + 1 < len && sql.charAt(i + 1) == '-') { + int nl = sql.indexOf('\n', i + 2); + i = nl < 0 ? len : nl; + out.append(' '); + continue; + } + out.append(c); + i++; + } + return out.toString(); } private boolean isUseStatement(String sql) { diff --git a/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java b/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java index 659cfb5..566ab0d 100644 --- a/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java +++ b/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java @@ -409,6 +409,43 @@ void writeKeywordInsideComment_isStillAReadForNonAdmin() { assertThat(decision.mutating()).isFalse(); } + @Test + void hiddenWriteScan_staysLinearOnAdversarialInput() { + // The text backstop used a lazy wildcard (\bSELECT\b[\s\S]*?\bINTO) + // and a regex block-comment strip, both of which backtracked + // quadratically: 224KB of repeated "SELECT " burned ~44s of CPU inside + // the guard, before the query ever reached the database. Any + // authenticated Editor user could stall a request thread with it. + String repeatedSelect = "SELECT " + "SELECT ".repeat(32_000); + String unterminatedBlockComment = "SELECT 1 /*" + "a/*".repeat(32_000); + + for (String hostile : List.of(repeatedSelect, unterminatedBlockComment)) { + long startedAt = System.currentTimeMillis(); + service.enforce( + new QueryRequest(hostile, null, null), + QueryExecutionContext.editor("viewer", false, false), + "postgresql" + ); + long elapsed = System.currentTimeMillis() - startedAt; + assertThat(elapsed) + .as("classification of a %d char statement must not backtrack", hostile.length()) + .isLessThan(5_000L); + } + } + + @Test + void insertIntoSelect_isNotMisreadAsSelectInto() { + // SELECT_INTO_PATTERN matches a bare INTO target now that the SELECT + // prefix is gone, so the caller must gate it on the statement actually + // reading as a SELECT. + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("INSERT INTO archive SELECT * FROM orders", null, null), + QueryExecutionContext.editor("admin", true, true), + "postgresql" + ); + assertThat(decision.primaryQueryType()).isEqualTo("INSERT"); + } + @Test void insertIntoSelect_isStillClassifiedAsInsert() { QueryExecutionPolicyService.PolicyDecision decision = service.enforce(