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 aa9432ef7f5ca5f14e2936d8f59292b9a8fc9597 Author: cgivre <[email protected]> AuthorDate: Tue Aug 11 19:02:28 2026 -0400 fix: harden client_rest per review (fullmatch, non-JSON guard, kerberos/close tests, tag-stripped credential check) - quote_literal/quote_literal_path use fullmatch instead of match: `$` matches before a trailing newline under match(), so 'foo\n' and 'dfs.tmp\n' silently passed the identifier allowlist despite the module's own comment claiming whitespace is rejected. Not exploitable (no quote/backslash/semicolon in the allowed charset), but wrong for a trust-boundary validator. Added regression cases to both parametrized rejection tests. - query() now wraps response.json() in try/except ValueError and raises DrillError naming the URL, instead of letting a bare JSONDecodeError escape when Drill or a fronting proxy/SSO gateway returns a 200 HTML page. - truncated is now max_rows > 0 and len(rows) >= max_rows, so max_rows=0 no longer reports a false truncation. - Dropped dead 'if not parts' check in quote_literal_path (str.split never returns an empty list; the any(...) clause already covers it). - _INVALID_CREDENTIALS is now matched against the login response body with HTML tags stripped first, so a marker split across tags (e.g. 'Invalid<br>username/password credentials') is still detected instead of silently passing as a successful login. - Added kerberos auth tests (missing extra -> DrillError naming the extra; extra present -> auth object reaches httpx.Client) and a close() test, removing the prior pragma: no cover now that both branches are exercised. - test_gives_up_after_one_retry now asserts login/query call counts, so an unbounded retry loop would fail the test instead of hanging. - Documented in test_login_rejects_200_with_invalid_credentials_body why it deliberately omits a /query.json mock (that omission is what makes the test non-vacuous). - Moved _error_text above RestClient so Task 6's metadata/management methods can append to the end of the file instead of inserting mid-file. --- drill_mcp/client_rest.py | 62 ++++++++++++++++----------- tests/test_client_rest.py | 104 +++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 136 insertions(+), 30 deletions(-) diff --git a/drill_mcp/client_rest.py b/drill_mcp/client_rest.py index 602efaf..3138c7a 100644 --- a/drill_mcp/client_rest.py +++ b/drill_mcp/client_rest.py @@ -49,9 +49,16 @@ _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 -# case-insensitively, with flexible whitespace, since it arrives embedded in -# markup rather than as the whole body. +# case-insensitively, with flexible whitespace, against the body with HTML +# tags stripped first -- the marker can arrive with tags inside the phrase +# (e.g. "Invalid<br>username/password credentials"), which a plain regex +# search against the raw markup would miss. _INVALID_CREDENTIALS = re.compile(r"invalid\s+username\s*/\s*password\s+credentials", re.IGNORECASE) +_HTML_TAG = re.compile(r"<[^>]+>") + + +def _strip_tags(html: str) -> str: + return _HTML_TAG.sub(" ", html) class DrillError(Exception): @@ -74,7 +81,7 @@ def quote_literal(value: str) -> str: parameters, so anything outside the safe character set is rejected rather than escaped. """ - if not _IDENTIFIER.match(value): + if not _IDENTIFIER.fullmatch(value): raise DrillError(f"invalid identifier: {value!r}") return f"'{value}'" @@ -82,11 +89,28 @@ def quote_literal(value: str) -> str: def quote_literal_path(value: str) -> str: """Same as `quote_literal`, but permits a dotted schema path like `dfs.tmp`.""" parts = value.split(".") - if not parts or any(not _IDENTIFIER.match(part) for part in parts): + if any(not _IDENTIFIER.fullmatch(part) for part in parts): raise DrillError(f"invalid identifier: {value!r}") return f"'{value}'" +def _error_text(response: httpx.Response) -> str: + """Drill's own error text is what a model needs to fix its SQL. Truncate it.""" + try: + payload = response.json() + except ValueError: + payload = None + message = "" + if isinstance(payload, dict): + message = payload.get("errorMessage") or payload.get("message") or "" + if not message: + message = response.text + message = " ".join(message.split()) + if len(message) > 2000: + message = message[:2000] + " ... [truncated]" + return message or f"Drill returned HTTP {response.status_code}" + + # -- client -------------------------------------------------------------- @@ -98,7 +122,7 @@ class RestClient: if config.auth == "kerberos": try: from httpx_gssapi import HTTPSPNEGOAuth - except ImportError as exc: # pragma: no cover - exercised in Task 7 style + except ImportError as exc: raise DrillError( "auth: kerberos requires the kerberos extra: pip install drill-mcp[kerberos]" ) from exc @@ -143,7 +167,7 @@ class RestClient: f"authentication endpoint at {self._config.url} returned " f"HTTP {response.status_code}" ) - if _INVALID_CREDENTIALS.search(response.text): + if _INVALID_CREDENTIALS.search(_strip_tags(response.text)): raise DrillError( f"authentication failed for user {self._config.user!r} at {self._config.url}" ) @@ -189,28 +213,16 @@ class RestClient: "/query.json", json={"queryType": "SQL", "query": sql, "autoLimit": max_rows}, ) - payload = response.json() + try: + payload = response.json() + except ValueError as exc: + raise DrillError( + f"Drill at {self._config.url} returned a non-JSON response" + ) from exc rows = payload.get("rows") or [] return QueryResult( columns=payload.get("columns") or [], rows=rows, query_id=payload.get("queryId"), - truncated=len(rows) >= max_rows, + truncated=max_rows > 0 and len(rows) >= max_rows, ) - - -def _error_text(response: httpx.Response) -> str: - """Drill's own error text is what a model needs to fix its SQL. Truncate it.""" - try: - payload = response.json() - except ValueError: - payload = None - message = "" - if isinstance(payload, dict): - message = payload.get("errorMessage") or payload.get("message") or "" - if not message: - message = response.text - message = " ".join(message.split()) - if len(message) > 2000: - message = message[:2000] + " ... [truncated]" - return message or f"Drill returned HTTP {response.status_code}" diff --git a/tests/test_client_rest.py b/tests/test_client_rest.py index 202b420..80c8e74 100644 --- a/tests/test_client_rest.py +++ b/tests/test_client_rest.py @@ -17,6 +17,9 @@ # License. # +import sys +import types + import httpx import pytest import respx @@ -46,7 +49,17 @@ class TestQuoting: @pytest.mark.parametrize( "bad", - ["foo'bar", "foo;DROP", "foo bar", "foo\nbar", "", "foo\\bar", "foo`bar", "dfs.tmp"], + [ + "foo'bar", + "foo;DROP", + "foo bar", + "foo\nbar", + "", + "foo\\bar", + "foo`bar", + "dfs.tmp", + "foo\n", # trailing newline: `$` matches before it under .match(), not under .fullmatch() + ], ) def test_literal_rejects_dangerous_input(self, bad): with pytest.raises(DrillError, match="invalid identifier"): @@ -54,7 +67,17 @@ class TestQuoting: @pytest.mark.parametrize( "bad", - ["foo'bar", "foo;DROP", "foo bar", "foo\nbar", "", "..", "foo\\bar", "foo`bar"], + [ + "foo'bar", + "foo;DROP", + "foo bar", + "foo\nbar", + "", + "..", + "foo\\bar", + "foo`bar", + "dfs.tmp\n", # trailing newline on the last segment + ], ) def test_literal_path_rejects_dangerous_input(self, bad): with pytest.raises(DrillError, match="invalid identifier"): @@ -88,6 +111,24 @@ class TestQuery: ) assert make_client().query("SELECT 1", max_rows=2).truncated is True + @respx.mock + def test_not_truncated_when_max_rows_is_zero(self): + # 0 >= 0 would be a false "truncated" without the max_rows > 0 guard. + respx.post(f"{BASE}/query.json").mock( + return_value=httpx.Response(200, json={"columns": [], "rows": []}) + ) + assert make_client().query("SELECT 1", max_rows=0).truncated is False + + @respx.mock + def test_non_json_response_is_reported_as_drill_error(self): + # A 200 HTML page (e.g. from an SSO gateway or an undetected auth + # failure in front of Drill) must not surface a bare JSONDecodeError. + respx.post(f"{BASE}/query.json").mock( + return_value=httpx.Response(200, text="<html><body>not json</body></html>") + ) + with pytest.raises(DrillError, match="non-JSON"): + make_client().query("SELECT 1", max_rows=10) + @respx.mock def test_drill_error_text_is_surfaced(self): respx.post(f"{BASE}/query.json").mock( @@ -153,10 +194,15 @@ class TestBasicAuth: @respx.mock def test_gives_up_after_one_retry(self): - respx.post(f"{BASE}/j_security_check").mock(return_value=httpx.Response(200)) - respx.post(f"{BASE}/query.json").mock(return_value=httpx.Response(401)) + login = respx.post(f"{BASE}/j_security_check").mock(return_value=httpx.Response(200)) + query = respx.post(f"{BASE}/query.json").mock(return_value=httpx.Response(401)) with pytest.raises(DrillError, match="authentication"): make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + # Without bounded call counts, an unbounded retry loop would hang + # instead of failing -- these assertions are what actually prove the + # retry is bounded to exactly one attempt. + assert login.call_count == 2 + assert query.call_count == 2 @respx.mock def test_login_failure_is_reported(self): @@ -169,7 +215,14 @@ class TestBasicAuth: """Drill's j_security_check returns HTTP 200 even on a wrong password; the failure is only visible in the HTML error page body. This is the regression test: without checking the body, a wrong password is - silently treated as a successful login.""" + silently treated as a successful login. + + Deliberately no /query.json mock is registered: if login wrongly + succeeds, the client proceeds to query() and respx raises + AllMockedAssertionError instead of DrillError, failing this test. + That's what makes this test non-vacuous -- do not add a query mock + here, it would silently hollow out the regression coverage. + """ respx.post(f"{BASE}/j_security_check").mock( return_value=httpx.Response( 200, @@ -180,6 +233,21 @@ class TestBasicAuth: make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) assert "s3cret" not in str(exc.value) + @respx.mock + def test_login_rejects_200_with_tags_inside_the_marker_phrase(self): + """The invalid-credentials marker can arrive with HTML tags inside the + phrase itself (e.g. a <br> mid-sentence), not just surrounding it. + Matching against the raw markup would miss this.""" + respx.post(f"{BASE}/j_security_check").mock( + return_value=httpx.Response( + 200, + text="<html><body>Invalid<br>username/password credentials</body></html>", + ) + ) + with pytest.raises(DrillError, match="authentication") as exc: + make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + assert "s3cret" not in str(exc.value) + @respx.mock def test_login_succeeds_on_200_with_ordinary_body(self): login = respx.post(f"{BASE}/j_security_check").mock( @@ -221,3 +289,29 @@ class TestBasicAuth: ) make_client().query("SELECT 1", max_rows=1) assert not login.called + + +class TestClose: + def test_close_closes_the_underlying_http_client(self): + client = make_client() + assert client._http.is_closed is False + client.close() + assert client._http.is_closed is True + + +class TestKerberosAuth: + def test_missing_extra_raises_a_clear_error(self, monkeypatch): + monkeypatch.setitem(sys.modules, "httpx_gssapi", None) + with pytest.raises(DrillError, match=r"drill-mcp\[kerberos\]"): + make_client(auth="kerberos") + + def test_extra_present_wires_the_auth_object_into_the_http_client(self, monkeypatch): + class _FakeSpnegoAuth(httpx.Auth): + def auth_flow(self, request): + yield request + + sentinel = _FakeSpnegoAuth() + stub = types.SimpleNamespace(HTTPSPNEGOAuth=lambda: sentinel) + monkeypatch.setitem(sys.modules, "httpx_gssapi", stub) + client = make_client(auth="kerberos") + assert client._http.auth is sentinel
