Skip to content

fix(source/cockroachdb): detect existing LIMIT regardless of whitespace #3554

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
anxkhn wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(source/cockroachdb): detect existing LIMIT regardless of whitespace #3554

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions internal/sources/cockroachdb/cockroachdb.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -299,6 +299,10 @@ func (s *Source) CanExecuteWrite(sql string) error {
return nil
}

// limitClauseRegexp matches a LIMIT keyword regardless of surrounding
// whitespace, so multiline or tab-formatted queries are detected too.
var limitClauseRegexp = regexp.MustCompile(`(?i)\bLIMIT\b`)
Comment on lines +302 to +304

@gemini-code-assist gemini-code-assist Bot Jun 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

security-high high

To prevent security bypasses where a client includes the word LIMIT inside a comment or a string literal, we should define a regex to match and strip SQL comments and string literals before checking for the LIMIT clause.

Suggested change
// limitClauseRegexp matches a LIMIT keyword regardless of surrounding
// whitespace, so multiline or tab-formatted queries are detected too.
var limitClauseRegexp = regexp.MustCompile(`(?i)\bLIMIT\b`)
// sqlTokenRegexp matches SQL string literals, double-quoted identifiers,
// multi-line comments, and single-line comments.
var sqlTokenRegexp = regexp.MustCompile(`'[^'\\]*(?:\\.[^'\\]*)*'|"[^"\\]*(?:\\.[^"\\]*)*"|/\*[^*]*\*+(?:[^/*][^*]*\*+)*/|--.*`)
// limitClauseRegexp matches a LIMIT keyword regardless of surrounding
// whitespace, so multiline or tab-formatted queries are detected too.
var limitClauseRegexp = regexp.MustCompile(`(?i)\bLIMIT\b`)


// ApplyQueryLimits applies row limits to a SQL query for MCP security compliance.
// Context timeout management is the responsibility of the caller (following Go best practices).
// Returns potentially modified SQL with LIMIT clause for SELECT queries.
Expand All @@ -308,8 +312,7 @@ func (s *Source) ApplyQueryLimits(sql string) (string, error) {
// Apply row limit only to SELECT queries
if sqlType == SQLTypeSelect && s.MaxRowLimit > 0 {
// Check if query already has LIMIT clause
normalized := strings.ToUpper(sql)
if !strings.Contains(normalized, " LIMIT ") {
if !limitClauseRegexp.MatchString(sql) {
// Add LIMIT clause - trim trailing whitespace and semicolon
sql = strings.TrimSpace(sql)
sql = strings.TrimSuffix(sql, ";")
Expand Down
21 changes: 21 additions & 0 deletions internal/sources/cockroachdb/security_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -218,6 +218,27 @@ func TestApplyQueryLimits(t *testing.T) {
expectedSQL: "SELECT * FROM users LIMIT 50",
shouldAddLimit: false,
},
{
name: "SELECT with multiline existing LIMIT",
sql: "SELECT *\nFROM users\nLIMIT 50",
maxRowLimit: 100,
expectedSQL: "SELECT *\nFROM users\nLIMIT 50",
shouldAddLimit: false,
},
{
name: "SELECT with lowercase multiline existing LIMIT",
sql: "SELECT * FROM users\nlimit 50",
maxRowLimit: 100,
expectedSQL: "SELECT * FROM users\nlimit 50",
shouldAddLimit: false,
},
{
name: "SELECT with tab-separated existing LIMIT",
sql: "SELECT a\tLIMIT 5",
maxRowLimit: 100,
expectedSQL: "SELECT a\tLIMIT 5",
shouldAddLimit: false,
},
Comment on lines +221 to +241

@gemini-code-assist gemini-code-assist Bot Jun 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Add test cases to verify that LIMIT detection is not bypassed by comments or string literals, and that trailing single-line comments do not comment out the appended LIMIT clause.

		{
			name:           "SELECT with multiline existing LIMIT",
			sql:            "SELECT *\nFROM users\nLIMIT 50",
			maxRowLimit:    100,
			expectedSQL:    "SELECT *\nFROM users\nLIMIT 50",
			shouldAddLimit: false,
		},
		{
			name:           "SELECT with lowercase multiline existing LIMIT",
			sql:            "SELECT * FROM users\nlimit 50",
			maxRowLimit:    100,
			expectedSQL:    "SELECT * FROM users\nlimit 50",
			shouldAddLimit: false,
		},
		{
			name:           "SELECT with tab-separated existing LIMIT",
			sql:            "SELECT a\tLIMIT 5",
			maxRowLimit:    100,
			expectedSQL:    "SELECT a\tLIMIT 5",
			shouldAddLimit: false,
		},
		{
			name:           "SELECT with LIMIT in comment",
			sql:            "SELECT * FROM users -- LIMIT is not here",
			maxRowLimit:    100,
			expectedSQL:    "SELECT * FROM users -- LIMIT is not here\nLIMIT 100",
			shouldAddLimit: true,
		},
		{
			name:           "SELECT with LIMIT in string literal",
			sql:            "SELECT * FROM users WHERE name = 'LIMIT'",
			maxRowLimit:    100,
			expectedSQL:    "SELECT * FROM users WHERE name = 'LIMIT' LIMIT 100",
			shouldAddLimit: true,
		},
		{
			name:           "SELECT with trailing single-line comment",
			sql:            "SELECT * FROM users -- trailing comment",
			maxRowLimit:    100,
			expectedSQL:    "SELECT * FROM users -- trailing comment\nLIMIT 100",
			shouldAddLimit: true,
		}

{
name: "SELECT without LIMIT and semicolon",
sql: "SELECT * FROM users;",
Expand Down