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 6e63ceaa218672f6d34af978cb5a9506026f497b Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 00:53:07 2026 -0400 fix: address six minor findings from Task 6 review - reject a bare ".." (or empty) dot-segment in quote_identifier, closing a directory-traversal gap in fetch_columns's file-plugin path - add direct unit tests for quote_identifier_path/quote_identifier rejecting backticks, and rename two tests that claimed coverage they did not provide - factor response.json() decoding into a _json() helper and guard all six sites (query, storage_plugins, cluster_status x2, profiles, profile) against a non-JSON body - profiles() tolerates a non-dict payload and clamps a negative limit to 0 - add a direct test for plugin_type - drop leftover ^...$ anchors from _IDENTIFIER for consistency with _FILE_IDENTIFIER --- drill_mcp/client_rest.py | 57 +++++++++++++++++++++-------- tests/test_client_rest.py | 93 +++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 132 insertions(+), 18 deletions(-) diff --git a/drill_mcp/client_rest.py b/drill_mcp/client_rest.py index d02fdf9..bed9076 100644 --- a/drill_mcp/client_rest.py +++ b/drill_mcp/client_rest.py @@ -47,7 +47,7 @@ from .redact import redact # empty string never matches, and `quote_literal_path` splits on "." and # validates every segment, so a lone "." (which splits into two empty # segments) is rejected too. -_IDENTIFIER = re.compile(r"^[A-Za-z0-9_$-]+$") +_IDENTIFIER = re.compile(r"[A-Za-z0-9_$-]+") # Drill's j_security_check returns HTTP 200 even on a wrong password; the only # signal is this marker string inside the HTML error page body. Matched @@ -135,6 +135,20 @@ def quote_identifier_path(value: str) -> str: _FILE_IDENTIFIER = re.compile(r"[A-Za-z0-9_$.-]+") +def _is_valid_file_identifier(value: str) -> bool: + """`_FILE_IDENTIFIER` plus a check that no "." segment is empty or "..". + + `_FILE_IDENTIFIER` alone permits a bare ".." (or a leading/trailing/ + doubled dot), which Drill treats as a directory reference -- e.g. + `columns("dfs.tmp", "..")` would otherwise reach the workspace's parent. + "/" is excluded from the character class, so this can never traverse more + than one level or reach a specific file, but it should still be rejected. + """ + if not _FILE_IDENTIFIER.fullmatch(value): + return False + return all(part not in ("", "..") for part in value.split(".")) + + def quote_identifier(value: str) -> str: """Quote a single identifier that may itself contain dots, e.g. a filename. @@ -145,7 +159,7 @@ def quote_identifier(value: str) -> str: quotes plugin.`workspace`.`file.ext` the same way). Reject rather than escape -- same trust boundary as `quote_identifier_path`. """ - if not _FILE_IDENTIFIER.fullmatch(value): + if not _is_valid_file_identifier(value): raise DrillError(f"invalid identifier: {value!r}") return f"`{value}`" @@ -159,6 +173,19 @@ def _check_query_id(query_id: str) -> str: return query_id +def _json(response: httpx.Response, url: str) -> Any: + """Decode a response body as JSON, converting a decode failure to `DrillError`. + + A 200 response carrying HTML -- exactly the auth-proxy scenario `_login` + exists to handle -- must never raise a raw `json.JSONDecodeError` out of + the client boundary. + """ + try: + return response.json() + except ValueError as exc: + raise DrillError(f"Drill at {url} returned a non-JSON response") from exc + + def _error_text(response: httpx.Response) -> str: """Drill's own error text is what a model needs to fix its SQL. Truncate it.""" try: @@ -250,7 +277,7 @@ def fetch_columns(query: Query, schema: str, table: str) -> list[dict[str, Any]] # `_FILE_IDENTIFIER` (not `_IDENTIFIER`) because file-plugin table # names are filenames and may contain a literal "." (e.g. "sales.csv") # as ONE identifier -- see `quote_identifier`. - if not _FILE_IDENTIFIER.fullmatch(table): + if not _is_valid_file_identifier(table): raise DrillError(f"invalid identifier: {table!r}") # Same split: file plugins have dynamic schemas and no @@ -391,12 +418,7 @@ class RestClient: "/query.json", json={"queryType": "SQL", "query": sql, "autoLimit": max_rows}, ) - try: - payload = response.json() - except ValueError as exc: - raise DrillError( - f"Drill at {self._config.url} returned a non-JSON response" - ) from exc + payload = _json(response, self._config.url) rows = payload.get("rows") or [] return QueryResult( columns=payload.get("columns") or [], @@ -427,12 +449,12 @@ class RestClient: # -- management -------------------------------------------------------- def storage_plugins(self) -> list[dict[str, Any]]: - payload = self._request("GET", "/storage.json").json() - return redact(payload) + response = self._request("GET", "/storage.json") + return redact(_json(response, self._config.url)) def cluster_status(self) -> dict[str, Any]: - cluster = self._request("GET", "/cluster.json").json() - status = self._request("GET", "/status.json").json() + cluster = _json(self._request("GET", "/cluster.json"), self._config.url) + status = _json(self._request("GET", "/status.json"), self._config.url) merged = dict(cluster) if isinstance(cluster, dict) else {"cluster": cluster} if isinstance(status, dict): merged.update(status) @@ -441,14 +463,19 @@ class RestClient: return merged def profiles(self, limit: int) -> list[dict[str, Any]]: - payload = self._request("GET", "/profiles.json").json() + limit = max(limit, 0) + response = self._request("GET", "/profiles.json") + payload = _json(response, self._config.url) + if not isinstance(payload, dict): + payload = {} running = payload.get("runningQueries") or [] finished = payload.get("finishedQueries") or [] return (list(running) + list(finished))[:limit] def profile(self, query_id: str) -> dict[str, Any]: _check_query_id(query_id) - return self._request("GET", f"/profiles/{query_id}.json").json() + response = self._request("GET", f"/profiles/{query_id}.json") + return _json(response, self._config.url) def cancel_query(self, query_id: str) -> str: _check_query_id(query_id) diff --git a/tests/test_client_rest.py b/tests/test_client_rest.py index 3d32e92..fcf2b4f 100644 --- a/tests/test_client_rest.py +++ b/tests/test_client_rest.py @@ -24,7 +24,14 @@ import httpx import pytest import respx -from drill_mcp.client_rest import DrillError, RestClient, quote_literal, quote_literal_path +from drill_mcp.client_rest import ( + DrillError, + RestClient, + quote_identifier, + quote_identifier_path, + quote_literal, + quote_literal_path, +) from drill_mcp.config import load_config BASE = "http://drill:8047"; @@ -83,6 +90,22 @@ class TestQuoting: with pytest.raises(DrillError, match="invalid identifier"): quote_literal_path(bad) + def test_identifier_path_rejects_a_backtick(self): + with pytest.raises(DrillError, match="invalid identifier"): + quote_identifier_path("dfs`x") + + def test_identifier_rejects_a_backtick(self): + with pytest.raises(DrillError, match="invalid identifier"): + quote_identifier("a`b") + + def test_identifier_rejects_a_bare_dot_dot_segment(self): + with pytest.raises(DrillError, match="invalid identifier"): + quote_identifier("..") + + def test_identifier_rejects_a_dot_dot_segment_within_a_longer_name(self): + with pytest.raises(DrillError, match="invalid identifier"): + quote_identifier("foo...bar") + class TestQuery: @respx.mock @@ -420,6 +443,16 @@ class TestMetadata: with pytest.raises(DrillError, match="invalid identifier"): make_client().tables("dfs'; DROP TABLE x --") + @respx.mock + def test_plugin_type_returns_the_schema_type(self): + route = respx.post(f"{BASE}/query.json").mock( + return_value=query_response( + ["SCHEMA_NAME", "TYPE"], [{"SCHEMA_NAME": "dfs.tmp", "TYPE": "file"}] + ) + ) + assert make_client().plugin_type("dfs.tmp") == "file" + assert b"'dfs.tmp'" in route.calls.last.request.read() + class TestFilePluginMetadata: """File plugins are absent from INFORMATION_SCHEMA; they need SHOW FILES.""" @@ -527,12 +560,21 @@ class TestFilePluginMetadata: assert b"`README`" in route.calls[1].request.read() @respx.mock - def test_columns_rejects_a_backtick_in_the_table_name(self): + def test_columns_table_name_guard_rejects_a_backtick_before_any_query_fires(self): + # The `_FILE_IDENTIFIER` guard in `fetch_columns` catches this before + # `plugin_type` (and thus `quote_identifier_path`) is ever reached. route = respx.post(f"{BASE}/query.json").mock(return_value=self._schemata("file")) with pytest.raises(DrillError, match="invalid identifier"): make_client().columns("dfs.tmp", "sales`; DROP TABLE x --.csv") assert not route.called + @respx.mock + def test_columns_rejects_a_bare_dot_dot_table_name(self): + route = respx.post(f"{BASE}/query.json").mock(return_value=self._schemata("file")) + with pytest.raises(DrillError, match="invalid identifier"): + make_client().columns("dfs.tmp", "..") + assert not route.called + @respx.mock def test_columns_uses_information_schema_for_a_non_file_plugin(self): route = respx.post(f"{BASE}/query.json").mock( @@ -559,7 +601,9 @@ class TestFilePluginMetadata: assert make_client().tables("nope") == [] @respx.mock - def test_show_files_path_still_rejects_injection(self): + def test_tables_schema_name_is_rejected_before_the_plugin_type_lookup(self): + # `plugin_type`'s own `quote_literal_path` call rejects this before + # `quote_identifier_path` (the SHOW FILES path) is ever reached. respx.post(f"{BASE}/query.json").mock(return_value=self._schemata("file")) with pytest.raises(DrillError, match="invalid identifier"): make_client().tables("dfs`; DROP TABLE x --") @@ -588,6 +632,16 @@ class TestManagement: assert plugins[0]["config"]["fs.s3a.secret.key"] == "***REDACTED***" assert plugins[0]["name"] == "s3" + @respx.mock + def test_storage_plugins_non_json_response_is_a_drill_error(self): + # A 200 response carrying HTML -- the auth-proxy scenario `_login` + # exists to handle -- must not raise a raw JSONDecodeError. + respx.get(f"{BASE}/storage.json").mock( + return_value=httpx.Response(200, text="<html>not json</html>") + ) + with pytest.raises(DrillError, match="non-JSON response"): + make_client().storage_plugins() + @respx.mock def test_cluster_status_merges_cluster_and_status(self): respx.get(f"{BASE}/cluster.json").mock( @@ -626,6 +680,31 @@ class TestManagement: ) assert make_client().profiles(limit=5)[0]["queryId"] == "live" + @respx.mock + def test_profiles_tolerates_a_non_dict_payload(self): + respx.get(f"{BASE}/profiles.json").mock( + return_value=httpx.Response(200, json=["unexpected", "list", "payload"]) + ) + assert make_client().profiles(limit=5) == [] + + @respx.mock + def test_profiles_clamps_a_negative_limit_to_zero(self): + respx.get(f"{BASE}/profiles.json").mock( + return_value=httpx.Response( + 200, + json={"runningQueries": [{"queryId": "live"}], "finishedQueries": []}, + ) + ) + assert make_client().profiles(limit=-5) == [] + + @respx.mock + def test_profiles_non_json_response_is_a_drill_error(self): + respx.get(f"{BASE}/profiles.json").mock( + return_value=httpx.Response(200, text="<html>not json</html>") + ) + with pytest.raises(DrillError, match="non-JSON response"): + make_client().profiles(limit=5) + @respx.mock def test_profile_fetches_one_query(self): respx.get(f"{BASE}/profiles/abc.json").mock( @@ -638,6 +717,14 @@ class TestManagement: with pytest.raises(DrillError, match="invalid"): make_client().profile("../../etc/passwd") + @respx.mock + def test_profile_non_json_response_is_a_drill_error(self): + respx.get(f"{BASE}/profiles/abc.json").mock( + return_value=httpx.Response(200, text="<html>not json</html>") + ) + with pytest.raises(DrillError, match="non-JSON response"): + make_client().profile("abc") + @respx.mock def test_cancel_query(self): route = respx.get(f"{BASE}/profiles/cancel/abc").mock(
