This is an automated email from the ASF dual-hosted git repository. cgivre pushed a commit to branch feat/drill-mcp-server in repository https://gitbox.apache.org/repos/asf/drill-mcp.git
commit 0413ffd40a25f6cd74bae14822320efe039cdbf5 Author: cgivre <[email protected]> AuthorDate: Tue Aug 11 15:47:28 2026 -0400 fix: close guard gaps found in review — TokenError, embedded writes, EXPLAIN recursion --- .../specs/2026-08-11-drill-mcp-design.md | 15 +++++- drill_mcp/guard.py | 55 +++++++++++++++++++-- tests/test_guard.py | 57 ++++++++++++++++++++++ 3 files changed, 121 insertions(+), 6 deletions(-) diff --git a/docs/superpowers/specs/2026-08-11-drill-mcp-design.md b/docs/superpowers/specs/2026-08-11-drill-mcp-design.md index 9aaa8eb..7196861 100644 --- a/docs/superpowers/specs/2026-08-11-drill-mcp-design.md +++ b/docs/superpowers/specs/2026-08-11-drill-mcp-design.md @@ -106,7 +106,15 @@ writable_plugins: [] # e.g. [dfs.tmp] `guard.py` parses each statement with `sqlglot` and applies: - Exactly one statement per call. Multiple statements are rejected. -- `SELECT`, `SHOW`, `DESCRIBE`, `EXPLAIN`, and `WITH ... SELECT` are allowed. +- `SELECT`, `SHOW`, `DESCRIBE`, and `WITH ... SELECT` are allowed. The read + branch sweeps the whole subtree, not just the root node, so a read-rooted + statement that writes deeper down (`WITH x AS (INSERT ... RETURNING *) + SELECT ...`, `SELECT ... INTO ...`) is still rejected. +- `EXPLAIN` is allowed, but the guard strips the leading `EXPLAIN` / + `EXPLAIN PLAN FOR` and re-checks the remainder, so the write allowlist and + hidden-schema rules apply to the explained statement. Allowing `EXPLAIN` on + the strength of its leading keyword alone would let `EXPLAIN PLAN FOR CREATE + TABLE ...` reach Drill without the allowlist ever running. - `CREATE TABLE AS`, `CREATE VIEW`, `CREATE TEMPORARY TABLE AS`, `DROP TABLE`, and `DROP VIEW` are allowed **only** when the target's leading identifier (the plugin, or plugin plus workspace) matches an entry in @@ -114,7 +122,10 @@ writable_plugins: [] # e.g. [dfs.tmp] `dfs` permits `dfs.tmp.foo`; an entry of `dfs.tmp` does not permit `dfs.raw.foo`. - Everything else — `INSERT`, `ALTER`, `SET`, `USE`, `REFRESH`, anything unrecognized — is rejected. -- A parse failure is a rejection, not a pass-through. +- A parse failure is a rejection, not a pass-through. This covers tokenizer + failures (`sqlglot.errors.TokenError`) as well as parse failures — they are + siblings under `SqlglotError`, not parent and child, and unterminated string + literals are exactly the input an adversarial caller produces. `sqlglot` rather than regex, deliberately. A regex guard is defeated by `-- CREATE TABLE` in a comment or `'DROP TABLE'` inside a string literal, and by diff --git a/drill_mcp/guard.py b/drill_mcp/guard.py index f2cdeeb..520850a 100644 --- a/drill_mcp/guard.py +++ b/drill_mcp/guard.py @@ -21,10 +21,23 @@ from sqlglot import exp DIALECT = "postgres" # closest available fit for Drill's Calcite SQL # Commands sqlglot does not model as expressions, but which cannot write. -_SAFE_COMMANDS = {"SHOW", "DESCRIBE", "DESC", "EXPLAIN"} +# EXPLAIN is handled separately (see _check_write): it is not blanket-safe +# because its body can itself be a write. +_SAFE_COMMANDS = {"SHOW", "DESCRIBE", "DESC"} _READ_TYPES = (exp.Select, exp.Union, exp.Intersect, exp.Except, exp.Subquery, exp.Describe) +# Node types that indicate a write is embedded somewhere inside a statement +# whose root node is a read type (e.g. `WITH x AS (INSERT ...) SELECT * FROM x`, +# or Postgres-dialect `SELECT ... INTO`). Checking only the root type is not +# enough: the safety property must not depend on Drill's parser being any +# narrower than sqlglot's Postgres dialect. +_EMBEDDED_WRITE_TYPES = (exp.Insert, exp.Update, exp.Delete, exp.Merge, exp.Create, exp.Drop, exp.Into) + +# EXPLAIN unwraps its body and re-checks it recursively; this bounds +# `EXPLAIN EXPLAIN EXPLAIN ...` so a malicious input cannot blow the stack. +_MAX_EXPLAIN_DEPTH = 5 + class PolicyError(Exception): """Raised when a statement is not permitted. The message is shown to the caller.""" @@ -60,12 +73,19 @@ def matches_prefix(qualified: str, entries: Iterable[str]) -> bool: def check(sql: str, policy: Policy) -> None: """Return None if `sql` is permitted under `policy`; raise PolicyError otherwise.""" + _check(sql, policy, depth=0) + + +def _check(sql: str, policy: Policy, depth: int) -> None: if not sql or not sql.strip(): raise PolicyError("empty SQL statement") try: statements = [s for s in sqlglot.parse(sql, read=DIALECT) if s is not None] - except sqlglot.ParseError as exc: + except sqlglot.errors.SqlglotError as exc: + # Catches both ParseError and TokenError (siblings under SqlglotError, + # not parent/child) — malformed input such as an unterminated string + # literal must be a PolicyError, never an uncaught crash. raise PolicyError( f"could not parse SQL, so it cannot be checked against policy and is rejected: {exc}" ) from exc @@ -77,15 +97,28 @@ def check(sql: str, policy: Policy) -> None: statement = statements[0] _check_hidden(statement, policy) - _check_write(statement, policy) + _check_write(statement, policy, depth) -def _check_write(statement: exp.Expression, policy: Policy) -> None: +def _check_write(statement: exp.Expression, policy: Policy, depth: int) -> None: if isinstance(statement, _READ_TYPES): + embedded = statement.find(*_EMBEDDED_WRITE_TYPES) + if embedded is not None: + raise PolicyError( + f"statement contains an embedded {embedded.key.upper()}, which is not permitted" + ) return if isinstance(statement, exp.Command): keyword = str(statement.this or "").upper() + if keyword == "EXPLAIN": + if depth >= _MAX_EXPLAIN_DEPTH: + raise PolicyError("too many nested EXPLAIN statements") + body = _explain_body(statement) + if not body.strip(): + raise PolicyError("EXPLAIN with no statement body is not permitted") + _check(body, policy, depth + 1) + return if keyword in _SAFE_COMMANDS: return raise PolicyError(f"statement type {keyword or 'UNKNOWN'} is not permitted") @@ -114,6 +147,20 @@ def _check_write(statement: exp.Expression, policy: Policy) -> None: raise PolicyError(f"statement type {statement.key.upper()} is not permitted") +def _explain_body(statement: exp.Command) -> str: + """Extract the SQL text following `EXPLAIN` (and an optional `PLAN FOR`). + + sqlglot has no Drill EXPLAIN grammar, so the whole statement falls back to + `exp.Command`, with everything after the leading keyword left as raw, + unparsed text. We recurse `check()` over that text rather than trusting + the leading keyword alone — otherwise `EXPLAIN PLAN FOR CREATE TABLE ...` + would bypass the write-target allowlist entirely. + """ + remainder = statement.args.get("expression") + text = remainder.this if isinstance(remainder, exp.Literal) else str(remainder or "") + return re.sub(r"^\s*PLAN\s+FOR\s+", "", text, flags=re.IGNORECASE) + + def _write_target(statement: exp.Expression) -> exp.Table | None: target = statement.this if isinstance(target, exp.Schema): diff --git a/tests/test_guard.py b/tests/test_guard.py index 42d4ff5..f5099d6 100644 --- a/tests/test_guard.py +++ b/tests/test_guard.py @@ -78,6 +78,9 @@ class TestWritesAreDeniedByDefault: "ALTER SESSION SET `store.format` = 'json'", "USE dfs.tmp", "REFRESH TABLE METADATA dfs.tmp.foo", + "UPDATE dfs.tmp.foo SET a = 1", + "DELETE FROM dfs.tmp.foo", + "MERGE INTO dfs.tmp.foo USING dfs.tmp.src ON true WHEN MATCHED THEN DELETE", ], ) def test_rejected_with_no_writable_plugins(self, sql): @@ -144,6 +147,13 @@ class TestInjectionAttempts: with pytest.raises(PolicyError, match="parse"): check("SELECT FROM WHERE ((", CLOSED) + def test_tokenizer_failure_is_a_policy_error_not_a_crash(self): + # An unterminated string literal raises sqlglot's TokenError, a sibling + # of ParseError under SqlglotError, not a subclass of it. A guard that + # only catches ParseError lets this escape uncaught. + with pytest.raises(PolicyError, match="parse"): + check("SELECT * FROM t WHERE note = 'it''s", CLOSED) + class TestMatchesPrefix: def test_exact_match(self): @@ -166,3 +176,50 @@ class TestMatchesPrefix: def test_empty_entries_never_match(self): assert not matches_prefix("dfs.tmp", []) + + +class TestEmbeddedWritesInsideReadRoots: + """A statement whose root node is a read type can still contain a write in + its subtree. Checking only the root type is not enough — the safety + property must not depend on Drill's parser being narrower than sqlglot's + Postgres dialect. Neither of these is executable Drill SQL, but the guard + must not rely on that. + """ + + def test_write_inside_a_cte_is_rejected(self): + with pytest.raises(PolicyError): + check( + "WITH x AS (INSERT INTO dfs.tmp.foo VALUES (1) RETURNING *) SELECT * FROM x", + CLOSED, + ) + + def test_select_into_is_rejected(self): + with pytest.raises(PolicyError): + check("SELECT * INTO dfs.tmp.newt FROM dfs.raw.src", CLOSED) + + +class TestExplainRecursesIntoItsBody: + """EXPLAIN is not blanket-safe: sqlglot parses it as an opaque exp.Command, + so the guard must strip the leading EXPLAIN / EXPLAIN PLAN FOR keywords and + re-run the full check on what remains, rather than trusting the keyword + alone. + """ + + def test_explain_of_a_read_is_permitted(self): + check("EXPLAIN PLAN FOR SELECT * FROM dfs.tmp.t", CLOSED) + + def test_explain_of_a_write_into_disallowed_plugin_is_rejected(self): + with pytest.raises(PolicyError): + check("EXPLAIN PLAN FOR CREATE TABLE s3.out AS SELECT 1", CLOSED) + + def test_explain_of_a_write_into_allowed_plugin_is_permitted(self): + check("EXPLAIN PLAN FOR CREATE TABLE dfs.tmp.out AS SELECT 1", OPEN) + + def test_explain_of_a_hidden_schema_is_rejected(self): + hidden = Policy(hidden_schemas=("sys",)) + with pytest.raises(PolicyError): + check("EXPLAIN PLAN FOR SELECT * FROM sys.options", hidden) + + def test_deeply_nested_explain_is_rejected_not_a_stack_overflow(self): + with pytest.raises(PolicyError): + check("EXPLAIN " * 10 + "SELECT 1", CLOSED)
