fix(security): reject subqueries in count_rows where clause

[moizpgedge]

count_rows's where parameter was interpolated directly into the generated SELECT COUNT(*) statement with no check beyond validateReadOnlyQuery, which only recognises constructs that escape read-only mode. A subquery in where is ordinary legal SQL, so it was invisible to that guard, and let a caller run a boolean- or error-based blind injection oracle against any table the connected role could read, not only the table named in the call.

Add validateCountRowsWhereClause, which rejects a bare SELECT or TABLE keyword in the clause's comment-stripped, literal-blanked residue (reusing stripSQLNoise from readonly_guard.go). TABLE is included because "TABLE tablename" is PostgreSQL shorthand for "SELECT * FROM tablename" and reads identical data without the word SELECT ever appearing; VALUES, the grammar's third row-returning form, is deliberately left unblocked, since it admits no FROM clause and so can never read a table, and unlike SELECT/TABLE it is not fully reserved and can legitimately be an unquoted column name. query_database and execute_explain are unaffected, since running arbitrary SQL, including subqueries, is their entire purpose.

Verified directly against a live server: every payload from the report is now rejected (boolean-blind and error-based cross-table oracles, EXISTS, the TABLE-shorthand variant), a legitimate simple predicate still succeeds, VALUES is correctly left unblocked, and the pre-existing stacked-query rejection is untouched.

Fixes #200

Summary by CodeRabbit

• Security

  • Strengthened count_rows validation to reject subqueries in where filters, including SELECT and TABLE patterns.
  • Continued allowing supported VALUES expressions and blocking stacked queries.
  • Improved rejection messages when a filter contains a prohibited subquery.

• Documentation

  • Added unreleased changelog notes describing the security mitigation and its scope.

[coderabbitai]

Walkthrough

count_rows now rejects SELECT and TABLE subqueries in its where clause after SQL-noise stripping. Unit and integration tests cover accepted predicates, evasion patterns, VALUES, stacked queries, and rejection messages. The changelog documents the mitigation and its scope.

Changes

count_rows security validation

Layer / File(s)	Summary
Where-clause validation and handler wiring internal/tools/count_rows.go, internal/tools/count_rows_test.go, docs/changelog.md	count_rows uses dedicated validation that strips comments and literals, rejects SELECT/TABLE constructs, updates tool descriptions and logging, and documents the mitigation. Unit tests cover valid predicates and evasion patterns.
End-to-end injection regression coverage test/count_rows_injection_test.go	JSON-RPC integration tests verify subquery rejection, accepted VALUES expressions, preserved stacked-query rejection, and the expected error text.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• pgEdge/pgedge-postgres-mcp#196: Also tightens count_rows where-clause validation.

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant MCPServer
  participant CountRowsHandler
  participant PostgreSQL
  Client->>MCPServer: tools/call count_rows
  MCPServer->>CountRowsHandler: dispatch request
  CountRowsHandler->>CountRowsHandler: reject SELECT/TABLE subqueries
  CountRowsHandler->>PostgreSQL: execute accepted predicate
  PostgreSQL-->>CountRowsHandler: count result
  CountRowsHandler-->>Client: JSON-RPC response

Loading

Suggested reviewers: dpage

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main security fix in count_rows and is concise.
Linked Issues check	✅ Passed	The change blocks subquery-based reads in count_rows.where while preserving simple predicates and stacked-query rejection.
Out of Scope Changes check	✅ Passed	The added docs and tests all support the reported count_rows injection fix and appear in scope.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix-issue-200-count-rows-where-injection

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codacy-production]

Up to standards ✅

🟢 Issues 0 issues

Results:
0 new issues

View in Codacy

🟢 Metrics 36 complexity · 2 duplication

Metric	Results
Complexity	36
Duplication	2

View in Codacy

_{NEW Get contextual insights on your PRs based on Codacy's metrics, along with PR and Jira context, without leaving GitHub. Enable AI reviewer
TIP This summary will be updated as you push new changes.}

[coderabbitai]

🧹 Nitpick comments (1)

internal/tools/count_rows_test.go (1)

54-87: 🔒 Security & Privacy | 🔵 Trivial

Solid subquery-rejection coverage; consider adding nested-comment and dollar-quote evasion cases.

Given the guard's correctness rests on stripSQLNoise's noise-removal fidelity (see the request_verification comment on internal/tools/count_rows.go), it would strengthen this suite to add cases like "1=1 AND (/* /* */ SELECT 1 FROM secret_table */) IS NOT NULL" (nested comments) and a dollar-quoted-string variant, mirroring the existing SEL/**/ECT test's intent.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@internal/tools/count_rows_test.go` around lines 54 - 87, The subquery
rejection test lacks coverage for nested-comment and dollar-quoted-string
evasion paths in stripSQLNoise. Extend
TestValidateCountRowsWhereClause_RejectsSubqueries with representative cases
that hide or surround SELECT using nested comments and dollar-quoted strings,
mirroring the existing comment-splitting evasion coverage, and assert each is
rejected.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@internal/tools/count_rows_test.go`:
- Around line 54-87: The subquery rejection test lacks coverage for
nested-comment and dollar-quoted-string evasion paths in stripSQLNoise. Extend
TestValidateCountRowsWhereClause_RejectsSubqueries with representative cases
that hide or surround SELECT using nested comments and dollar-quoted strings,
mirroring the existing comment-splitting evasion coverage, and assert each is
rejected.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 83e5d014-a517-4873-bcfc-87918bd92c1a

📥 Commits

Reviewing files that changed from the base of the PR and between 56c515e and 0020acc.

📒 Files selected for processing (4)

• docs/changelog.md
• internal/tools/count_rows.go
• internal/tools/count_rows_test.go
• test/count_rows_injection_test.go