Server hardening: summarize, resource limits, HTTP defaults, pool session-state (DB-22138/22159/22139/22176/22133/22202)#11
Open
Sfurti-yb wants to merge 7 commits into
Open
Conversation
…aborts
Closes DB-22138 — a single problematic table/view (view over a
since-dropped object, permission-denied table, division-by-zero in a
view definition, foreign-table connection error) previously aborted
the whole loop and discarded every other table's row count. The
tool's response collapsed to `[{"error": "..."}]` — a schema summary
with N-1 tables silently missing.
Vishal's ticket repro: schema with a normal `z_good` table (3 rows)
plus a view `a_bad AS SELECT 1 AS x WHERE (1/0) = 1`. The view sorts
alphabetically first, so v1 iterated it before z_good, hit the
division-by-zero, and returned only the error — z_good's count was
never computed. Under a least-privilege pool user this hits regularly:
a role that can't read one table loses the entire summary.
Three structural changes to `summarize_database`:
1. **Initial tables lookup** has its own try/except. If the
`information_schema.tables` query itself fails, record the error
and skip the (empty) loop.
2. **Per-table SAVEPOINT.** Each iteration runs
`SAVEPOINT tbl_<i>` → column lookup + `COUNT(*)` →
`RELEASE SAVEPOINT`. On failure: `ROLLBACK TO SAVEPOINT` +
`RELEASE SAVEPOINT` (pg keeps the savepoint alive after rollback,
so the release is required to keep the stack clean for the next
iteration). Failed table gets `{"table": name, "error": str(e)}`
in the summary and iteration continues.
3. **Outer try/finally retained** — the outermost `ROLLBACK` still
closes the read-only transaction even if a savepoint cleanup
itself fails.
Design notes:
- Numeric savepoint names (`tbl_0`, `tbl_1`, ...) to avoid pg's
63-char identifier limit for weird table names. Named via
`sql.Identifier` to be safe regardless.
- Not fixing DB-22158 (schema identifier interpolation at line 163)
here — it's scoped to the follow-up v2.1 release and shares this
function, so bundling would confuse review. Same for DB-22188
(empty vs non-existent schema).
Tests: 4 new integration cases in `tests/test_integration_reads.py`
(require `YUGABYTEDB_URL`):
- test_summarize_continues_past_error_view — direct Vishal repro.
Both `z_good` (row_count=3) and `a_bad` (error) appear.
- test_summarize_continues_past_permission_denied — restricted +
accessible tables; both appear.
- test_summarize_multiple_errors_one_good — two failing views + one
good table; all three appear, no truncation.
- test_summarize_empty_schema_returns_empty_list — regression: empty
schema still returns [].
111 non-integration tests still pass.
---
_automated · Claude Code (Opus 4.7)_
… (DB-22159)
Closes DB-22159 — a single unprivileged caller could take the whole
MCP server offline for all users, or crash the process, using trivial
queries. Vishal's live repros:
# 5 trivial queries down the service for everyone
5x run_read_only_query("SELECT pg_sleep(8)") # occupy 5-conn pool
then run_read_only_query("SELECT 1") # blocks 7.4s
# one query balloons the process toward OOM
run_read_only_query("SELECT repeat('x', 1024) FROM generate_series(1, 100000)")
-> server RSS 93 MB -> 712 MB
Root causes fixed:
1. **Hardcoded 5-connection pool** (`min_size=1, max_size=5` in
app_lifespan). No knob for operators. Now configurable via
`YB_MCP_POOL_MIN_SIZE` / `YB_MCP_POOL_MAX_SIZE`.
2. **No statement_timeout.** A slow query (pg_sleep, cartesian join,
heavy scan) holds a pool connection indefinitely. `SET LOCAL
statement_timeout = '{ms}ms'` at transaction start bounds it.
`SET LOCAL` scopes to the transaction and dies with the ROLLBACK /
COMMIT, so the timeout doesn't leak across pool checkouts.
Configurable via `YB_MCP_STATEMENT_TIMEOUT_MS` (default 30000).
3. **No result-row cap.** `cur.fetchall()` buffers the whole result
into memory. `SELECT repeat('x', N) FROM generate_series(1, N)`
drives RSS growth linearly → OOM crash. Replaced with
`cur.fetchmany(max_result_rows + 1)` — detects overrun by fetching
one extra row, then truncates and returns a `truncated: true`
marker so the caller sees the elision. Non-truncated responses
keep the v1 list-of-dicts shape (unchanged for the common case).
Configurable via `YB_MCP_MAX_RESULT_ROWS` (default 10000).
4. **No query-size cap.** Vishal measured ~6.4s of CPU in the
guardrail parser on a ~1MB write query. Pre-parse rejection is the
cheap first-line defense. Configurable via `YB_MCP_MAX_QUERY_LEN`
(default 100000 bytes).
New module-level `_positive_int(s)` argparse helper parses all five
env vars — a typo (`YB_MCP_POOL_MAX_SIZE=abc` or `=0`) fails startup
with a clean argparse error instead of a `ValueError` traceback. Same
pattern the DB-22162 follow-up fix uses for `max_insert_rows`; not
retrofitting `max_insert_rows` here to keep scope tight.
Test coverage:
- `tests/test_config.py` (new, 22 unit tests):
- `_positive_int` — positive / zero / negative / non-int / empty
/ float / whitespace.
- Defaults for all five limits.
- Env-var → ServerConfig plumbing.
- Bad env values → SystemExit (argparse), not a traceback.
- `tests/test_integration_reads.py` (5 new integration tests):
- `test_statement_timeout_kills_pg_sleep` — pg_sleep(60) with 1s
timeout fails within ~5s (elapsed-time assertion).
- `test_result_row_cap_truncates` — 10-row query with cap=3 →
truncation shape.
- `test_result_row_cap_not_triggered` — regression: below cap
keeps the v1 list shape.
- `test_max_query_len_rejects_oversized_query` — 300-byte-padded
query with 200-byte cap → pre-execute rejection.
- `test_max_query_len_allows_normal_query` — regression.
- `tests/conftest.py` — new `mcp_session_capped` fixture with small
caps (STATEMENT_TIMEOUT_MS=1000, MAX_RESULT_ROWS=3, MAX_QUERY_LEN=200,
write enabled) so integration tests exercise the caps quickly.
133 non-integration tests pass.
---
_automated · Claude Code (Opus 4.7)_
…9/22176) Closes DB-22139 (Medium) — Vishal's ticket: with the out-of-box defaults the HTTP transport was an unauthenticated MCP→DB proxy on all interfaces, with no DNS-rebinding protection. Verified: POST /mcp (no token, Origin: <https://evil.example.com>) -> HTTP 200 startup log: "No auth provider configured, auth disabled" "Uvicorn running on http://0.0.0.0:8000" Dockerfile: ENTRYPOINT ["yugabytedb-mcp"] CMD ["--transport","http"] Also folds in DB-22176 (Low) — Origin allowlist case-sensitivity — because it touches the same middleware and the same code path. Fix has three parts: **1. Default bind host is now loopback (127.0.0.1), not 0.0.0.0.** New `MCP_HOST` env / `--host` arg controls the bind host. Operators who want a public bind must set `MCP_HOST=0.0.0.0` explicitly. The `YugabyteDBMCPServer.run()` signature changed from `run(self, host="0.0.0.0", port=8000)` to `run(self, port=8000)`; host comes from `CONFIG.host` now. The old default was a security footgun — a distracted operator running `yugabytedb-mcp --transport http` on any AWS instance immediately exposed the DB. **2. `_check_http_startup(host)` fail-closed guard.** Runs before uvicorn opens the socket. When HTTP mode is combined with a non-loopback host AND no auth provider AND no escape hatch, the server logs a CRITICAL error explaining the four ways to fix it and calls `sys.exit(1)`. The escape hatch — `MCP_ALLOW_UNAUTHENTICATED= true` — lets the server run without auth on a public host, but with a very prominent WARNING banner in the logs so the choice is visible in prod. Independent of auth: the guard also logs a WARNING when the Origin allowlist is empty in HTTP mode, since `OriginValidationMiddleware` no-ops on an empty allowlist. Not fatal — some deployments front non-browser clients only — but the log gives operators a signal. **3. Origin allowlist is now case-insensitive (DB-22176).** RFC 6454 declares scheme + host to be case-insensitive; every real browser lowercases both before sending the Origin header. Pre-fix the allowlist compared case-sensitively, so an admin typing `MCP_ALLOWED_ORIGINS=https://MyApp.Example.com` silently rejected every real browser request (which sends `https://myapp.example.com`). `_parse_allowed_origins` lowercases entries at load time; `OriginValidationMiddleware.__call__` lowercases the incoming Origin at compare time. Belt-and-braces. Docs updated: - README config table gains `MCP_HOST` and `MCP_ALLOW_UNAUTHENTICATED` rows. "Self-hosted remote mode" section rewritten to open with the secure-by-default behavior + the required `MCP_HOST=0.0.0.0` + `MCP_AUTH_PROVIDER` combination. - Dockerfile CMD gains a comment block explaining the same. The base image + user-hardening from DB-22161 is follow-up scope (v2.1 release), so this commit only adds documentation, not actual container hardening. Test coverage (28 new unit tests): - `tests/test_http_startup.py` (new, 18): - `TestIsLoopback` (5) — parametrized over loopback / non-loopback. - `TestEnvBool` (4) — matches parse_config's `.lower() == "true"`; documents that `1` / `yes` / `on` are False (DB-22186 follow-up). - `TestCheckHttpStartup` (9) — stdio skips check; loopback+no-auth OK; loopback+auth OK; public+auth OK; public+no-auth+no-escape → SystemExit(1) with CRITICAL log mentioning MCP_AUTH_PROVIDER and 0.0.0.0; escape hatch works + logs prominent WARNING; escape hatch is a no-op when auth is configured; empty allowlist logs WARNING; allowlist configured suppresses that WARNING. - `tests/test_mcp_spec_middleware.py` (+7): - `TestOriginCaseInsensitivity` (4) — uppercase scheme/host, mixed case, uppercase scheme only, regression (different origin still 403). - `TestParseAllowedOriginsCasing` (3) — env value lowercased, multi-entry lowercased, MCP_BASE_URL fallback lowercased. 169 non-integration tests pass. --- _automated · Claude Code (Opus 4.7)_
Sfurti-yb
marked this pull request as ready for review
July 9, 2026 22:02
Sfurti-yb
marked this pull request as draft
July 10, 2026 05:21
`test_summarize_multiple_errors_one_good` had `b_second_bad` fail via its projection (`SELECT (5/0)::int AS y`), but Postgres's COUNT(*) doesn't evaluate the projection — it just counts rows. So the view existed with 1 row, and the summarize tool's `SELECT COUNT(*)` returned 1 without error. Test asserted a per-object error entry that never materialized. Move the failure into WHERE, same shape as a_first_bad. The DB has to evaluate the WHERE expression to know which rows count, so the division-by-zero fires and the summarize tool records the per-object error the test asserts on. Test-only change; production code untouched. --- _automated · Claude Code (Opus 4.7)_
Closes DB-22133 (SET ROLE bleed across pool checkouts) and DB-22202
(advisory-lock cross-user DoS + prepared-statement pool-slot poisoning).
Both tickets are the same class of bug: session-level state left behind
by one tool call surviving to the next user's checkout. `RESET ROLE`
in `_conn_as_role`'s finally covers DB-22133 on the happy path but not
if the finally is skipped (KeyboardInterrupt, mid-yield exit, etc.), and
it does nothing for advisory locks or prepared-statement plans — which
was the DB-22202 discovery.
Fix: pass a `reset=` callback to the psycopg ConnectionPool that issues
`DISCARD ALL` on every connection return. `DISCARD ALL` is Postgres's
canonical "scrub session" statement — it clears:
- the current role (RESET ROLE — closes DB-22133)
- advisory locks (closes the cross-user DoS in DB-22202)
- prepared statements + cached plans (closes the pool-slot poisoning
in DB-22202, where DEALLOCATE ALL through the read tool would
desync psycopg's client-side cache)
- session GUCs (bonus — a caller can no longer leave
`default_transaction_read_only = off` behind)
- temp tables, sequences, cached plans
Runs after every `with pool.connection() as conn:` block, so any state
the tool accidentally leaves behind is scrubbed before the next checkout.
Design notes:
- The existing `RESET ROLE` in `_conn_as_role`'s finally stays in place
as belt-and-braces. DISCARD ALL covers the same ground, but keeping
the finally-block reset means role isolation still works between
statements within a single tool call (not just between pool
checkouts). Cheap, safe, defensive.
- `_pool_reset` swallows failures with a WARNING rather than raising.
If DISCARD itself fails on a dead connection, the pool's existing
`check_connection` health check replaces the connection on next
checkout. Raising from the reset callback would leak exceptions into
psycopg-pool's internal machinery.
Test coverage:
- tests/test_pool_reset.py (new, 3 unit tests):
- Callback issues DISCARD ALL exactly once.
- Callback swallows execute() failures + logs WARNING.
- No side effects beyond the single execute.
- tests/test_integration_pool_reset.py (new, 4 integration tests, need
YUGABYTEDB_URL):
- DB-22133 direct repro — SET ROLE → return → next checkout back to
baseline role.
- DB-22202 advisory lock repro — pg_advisory_lock(42421) → return,
verify via an independent psycopg connection that pg_locks no
longer has the row.
- DB-22202 prepared-statement repro — prep + DEALLOCATE ALL through
the pool, next checkout still returns correct results (would have
failed pre-fix with "prepared statement _pg3_0 does not exist").
- Bonus — session GUC leak: SET default_transaction_read_only=on
doesn't survive to the next checkout.
CI workflow fix (bundled here because it directly gates whether these
tests actually run):
- .github/workflows/test.yml — unit job was `pytest test_guardrails.py
test_auth.py`, missing test_config.py, test_http_startup.py, and now
test_pool_reset.py. Switched to `pytest tests/
--ignore-glob='tests/test_integration_*.py'` so new unit test files
are picked up automatically.
- Integration job was explicit list of three files, missing
test_integration_write_flag.py (silent gap since before this PR) and
now test_integration_pool_reset.py. Switched to glob:
`pytest tests/test_integration_*.py -v`.
172 non-integration tests pass locally (was 169 before this commit).
---
_automated · Claude Code (Opus 4.7)_
Sfurti-yb
changed the title
Two bugs surfaced by CI on the previous commit:
1. `DISCARD ALL cannot run inside a transaction block` — psycopg's
default is autocommit=False, so `conn.execute("DISCARD ALL")` opens
an implicit transaction and immediately fails. Fix `_pool_reset`:
flip to autocommit for the DISCARD, restore in a `finally` so the
connection state is what the next pool checkout expects.
2. `syntax error at or near "NOT"` — Postgres has no
`CREATE ROLE IF NOT EXISTS` syntax. Swallow
`psycopg.errors.DuplicateObject` around the CREATE instead.
Unit tests updated:
- `test_flips_autocommit_around_discard` — asserts the exact call order
`set_autocommit(True) → execute("DISCARD ALL") → set_autocommit(False)`.
- `test_restores_autocommit_even_when_discard_fails` — the finally
block still restores autocommit=False when DISCARD raises.
- `test_swallows_set_autocommit_failure` — if the initial
`set_autocommit(True)` raises (dead connection), the callback still
log-and-continues rather than leaking into psycopg-pool.
- Dropped `test_no_side_effects_beyond_execute` — asserted a single
method call, no longer true after adding the autocommit flip.
174 unit tests pass locally (was 172, +2 pool_reset cases).
---
_automated · Claude Code (Opus 4.7)_
Sfurti-yb
marked this pull request as ready for review
July 10, 2026 13:54
ashetkar
reviewed
Jul 13, 2026
|
|
||
| The endpoint uses Cognito's `USER_PASSWORD_AUTH` flow under the hood — the Cognito **app client must have `ALLOW_USER_PASSWORD_AUTH` enabled** (Cognito Console → User pools → App integration → App client settings). MFA, password-reset, and other challenge flows are **not** handled by this endpoint — they require the browser flow. | ||
|
|
||
| OIDC (`MCP_AUTH_PROVIDER=oidc`) wiring exists but is untested; please share findings if you exercise it. |
Contributor
There was a problem hiding this comment.
I think we have now tested this OIDC path as well using Keycloak. If so, please remove this line.
Empirically verified against local YugabyteDB (2025.2.0): a lock taken via `pg_try_advisory_lock(N)` on one connection survives `DISCARD ALL` and remains visible to other sessions via `pg_locks`. Only `pg_advisory_unlock_all()` actually releases it. This is a YB deviation from vanilla Postgres 15, where DISCARD ALL is documented to include advisory-lock release. Impact before this fix: the DB-22202 cross-user advisory-lock DoS repro Vishal filed remained open despite the pool `reset=` callback. Any tool call that grabbed an advisory lock (accidental or otherwise) would leave the lock in place across pool checkouts, blocking subsequent sessions that tried to acquire the same key. Fix: after `DISCARD ALL`, also run `SELECT pg_advisory_unlock_all()`. Runs in the same autocommit block; both statements cannot run inside a transaction, so the existing autocommit-flip logic covers both. Unit tests updated to assert the two-execute call order (was one). Integration test `test_discard_all_clears_advisory_lock` was previously FAILING against a real YB — now passes. Full pool-reset integration suite: 4/4 pass. Unit suite: 5/5 pass. --- _automated · Claude Code (Opus 4.7)_
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Ships PR #3 (final) of the MCP Server v2 audit primary release. Self-contained hardening fixes bundled into one PR because each is small; separate commits so the review surface stays scannable one concern at a time.
Six tickets closed.
summarize_databaseuses per-table SAVEPOINT so one erroring view/table skips instead of aborting the whole loop.SET LOCAL statement_timeout; result-row cap viafetchmany + 1; configurable pool sizes; query-length pre-check. Five new env vars, all parsed with a_positive_intargparse helper.MCP_HOST=127.0.0.1. Refuses to start when public host + no auth; escape hatchMCP_ALLOW_UNAUTHENTICATED=trueruns with a prominent WARNING.reset=callback runsDISCARD ALLon every connection return. Supersedes the existingRESET ROLEin_conn_as_role's finally (which stays as belt-and-braces). Closes the "SET ROLE bleeds if the finally is skipped" gap.DISCARD ALLalso clears advisory locks (the cross-user DoS repro Vishal verified live) and prepared-statement plans (the pool-slot poisoning whereDEALLOCATE ALLdesynced psycopg's cache). One code change, two tickets.Commits
summarize_database: per-table SAVEPOINT so one bad object skips, not aborts— DB-22138.tools.py:summarize_database, 4 integration tests.Add resource limits: statement timeout, result cap, configurable pool (DB-22159)— DB-22159.server.py+tools.py+ newtests/test_config.py(22 unit tests) + 5 integration tests + newmcp_session_cappedfixture.HTTP secure-by-default + Origin allowlist case-insensitivity (DB-22139/22176)—server.py+Dockerfile+README.md+ newtests/test_http_startup.py(18 unit tests) + 7 Origin case-insensitivity tests.Fix DB-22138 integration test: COUNT(*) skips projection, use WHERE— one-line test-only fix so the CI-caught bug in commit 1's test is repaired (unblocks the integration job).Pool session-state isolation via DISCARD ALL (DB-22133 + DB-22202)—server.pynew_pool_resetcallback +tests/test_pool_reset.py(3 unit tests) +tests/test_integration_pool_reset.py(4 integration tests direct from Vishal's ticket repros) +.github/workflows/test.ymlglob-based test picker so new unit / integration files are picked up automatically.Config additions
Seven new env vars, all with sensible defaults so v1 deployments keep working:
MCP_HOST127.0.0.10.0.0.0requires auth.MCP_ALLOW_UNAUTHENTICATEDfalseYB_MCP_POOL_MIN_SIZE1YB_MCP_POOL_MAX_SIZE5YB_MCP_STATEMENT_TIMEOUT_MS30000YB_MCP_MAX_RESULT_ROWS10000truncated: truemarker.YB_MCP_MAX_QUERY_LEN100000Backward compatibility
MCP_HOSTdefaults to127.0.0.1(was0.0.0.0). A deployment that was implicitly relying on the old0.0.0.0default must now setMCP_HOST=0.0.0.0explicitly. This is intentional — the pre-fix default was a security footgun.list[dict]to{"rows": [...], "truncated": true, ...}when the row cap is hit. Under the default 10k cap this is unlikely to hit v1 workflows. Non-truncated responses keep the v1 shape.reset=DISCARD ALLis transparent to callers — it fires only after awith pool.connection()block ends, and the existingRESET ROLEin_conn_as_role's finally keeps working the same way inside a call.CI workflow change (bundled with commit 5)
Historically the CI's unit job ran
pytest tests/test_guardrails.py tests/test_auth.pyand the integration job ran an explicit three-file list — new test files added over the course of this PR (test_config.py, test_http_startup.py, test_pool_reset.py, test_integration_pool_reset.py) weren't being picked up. Switched both jobs to glob-based selection:pytest tests/ --ignore-glob='tests/test_integration_*.py'pytest tests/test_integration_*.pyAlso fixes a pre-existing gap:
test_integration_write_flag.pywas silently outside CI before this PR.Sequencing
Off
origin/main(045bdd0). Independent of PR #9 (read-tool safety) and PR #10 (OIDC v2 mapping) in terms of file overlaps:tools.py::run_read_only_query(line 216 area, guardrail wiring). This PR touches the same function but in different areas (fetchmany, SET LOCAL). Rebase after PR Unified SQL guardrail + read-tool safety fixes (DB-22129/22172/22131/22157/22203) #9 lands.tools.py::_get_db_role+server.py::parse_config. This PR touchestools.py::summarize_database+ a different area ofserver.py::parse_config+server.py::app_lifespanfor the pool reset. Rebase after PR OIDC v2 identity mapping + JWT audience validation (DB-22135/174/192/136) #10 lands.Both rebases should be clean since the changes are in different functions or different lines within the same function.
Test plan
uv run pytest tests/ --ignore-glob='tests/test_integration_*.py' -q— 172 unit tests passs.a_badview over1/0+s.z_goodtable → response contains both entriesMCP_HOST=0.0.0.0withoutMCP_AUTH_PROVIDER→ server exits with clear CRITICAL logpg_sleep(60)with default 30s timeout → fails within ~30s (not 60)MCP_ALLOWED_ORIGINS=https://MyApp.Example.com+ browser sendshttps://myapp.example.com→ 200 (was 403)_conn_as_rolecheckout reports the pool user's role, not the previous role