mcp: validate data product names before interpolating into SQL

Author: bobbyiliev

src/environmentd/src/http/mcp.rs

@@ -768,6 +768,77 @@ async fn get_data_product_details(
     format_rows_response(rows, max_response_size)
 }
 
+/// Validates that a data product name is a safe SQL identifier.
+///
+/// The name is interpolated into `SELECT * FROM {name}`, so we parse a
+/// synthetic query and inspect the AST to ensure nothing beyond a bare
+/// table reference was smuggled in (preventing SQL injection).
+fn validate_data_product_name(name: &str) -> Result<(), McpRequestError> {
+    let name = name.trim();
+    if name.is_empty() {
+        return Err(McpRequestError::QueryValidationFailed(
+            "Data product name cannot be empty".to_string(),
+        ));
+    }
+
+    let make_err = || {
+        McpRequestError::QueryValidationFailed(format!(
+            "Invalid data product name: {}. Expected a valid object name, \
+             e.g. '\"database\".\"schema\".\"name\"' or 'my_view'",
+            name
+        ))
+    };
+
+    // Parse a synthetic query to verify the name is a valid table reference.
+    let probe = format!("SELECT 1 FROM {}", name);
+    let stmts = parse(&probe).map_err(|_| make_err())?;
+
+    if stmts.len() != 1 {
+        return Err(make_err());
+    }
+
+    // Inspect the AST to ensure the parsed query is a bare SELECT with
+    // exactly one table and no extra clauses (WHERE, ORDER BY, UNION, etc.)
+    // that would indicate injected SQL beyond the table name.
+    use mz_sql_parser::ast::{CteBlock, SetExpr, Statement};
+    let Statement::Select(select_stmt) = &stmts[0].ast else {
+        return Err(make_err());
+    };
+    let query = &select_stmt.query;
+
+    // No CTEs, ORDER BY, LIMIT, or OFFSET
+    let has_ctes = match &query.ctes {
+        CteBlock::Simple(ctes) => !ctes.is_empty(),
+        CteBlock::MutuallyRecursive(_) => true,
+    };
+    if has_ctes || !query.order_by.is_empty() || query.limit.is_some() || query.offset.is_some() {
+        return Err(make_err());
+    }
+
+    // Body must be a simple Select, not a UNION/EXCEPT/INTERSECT
+    let SetExpr::Select(select) = &query.body else {
+        return Err(make_err());
+    };
+
+    // No WHERE, GROUP BY, HAVING, QUALIFY, DISTINCT, or OPTIONS
+    if select.selection.is_some()
+        || !select.group_by.is_empty()
+        || select.having.is_some()
+        || select.qualify.is_some()
+        || select.distinct.is_some()
+        || !select.options.is_empty()
+    {
+        return Err(make_err());
+    }
+
+    // Exactly one table in FROM, no JOINs
+    if select.from.len() != 1 || !select.from[0].joins.is_empty() {
+        return Err(make_err());
+    }
+
+    Ok(())
+}
+
 /// Read rows from a data product. Issues a single read-only query.
 ///
 /// When `cluster_override` is provided, sets the cluster explicitly.
@@ -784,6 +855,9 @@ async fn read_data_product(
 ) -> Result<McpResult, McpRequestError> {
     debug!(name = %name, limit = limit, cluster_override = ?cluster_override, "Executing read_data_product");
 
+    // Validate the name is a safe identifier before interpolating into SQL.
+    validate_data_product_name(name)?;
+
     // Lightweight existence check: verify the data product is visible in the
     // catalog before running the read query. This gives a clean DataProductNotFound
     // error for missing or inaccessible products (including RBAC revocations)
@@ -1498,46 +1572,35 @@ mod tests {
         }
     }
 
-    // ── Response size cap tests ────────────────────────────────────────
+    // ── Data product name validation tests ─────────────────────────────
 
     #[mz_ore::test]
-    fn test_format_rows_response_within_limit() {
-        let rows = vec![vec![json!("a"), json!(1)], vec![json!("b"), json!(2)]];
-        let result = format_rows_response(rows, 1_000_000).unwrap();
-        let McpResult::ToolContent(content) = result else {
-            panic!("Expected ToolContent");
-        };
-        assert_eq!(content.content.len(), 1);
-        assert!(content.content[0].text.contains("\"a\""));
-        assert!(content.content[0].text.contains("\"b\""));
+    fn test_validate_data_product_name_valid() {
+        // Fully qualified quoted identifiers
+        assert!(validate_data_product_name(r#""materialize"."public"."my_view""#).is_ok());
+        assert!(validate_data_product_name(r#""my_db"."my_schema"."orders""#).is_ok());
+        // Two-part name
+        assert!(validate_data_product_name(r#""public"."my_view""#).is_ok());
+        // Unquoted simple name
+        assert!(validate_data_product_name("my_view").is_ok());
     }
 
     #[mz_ore::test]
-    fn test_format_rows_response_errors_when_over_limit() {
-        let rows: Vec<Vec<serde_json::Value>> = (0..100)
-            .map(|i| vec![json!(format!("row_{}", i)), json!(i)])
-            .collect();
-        let err = format_rows_response(rows, 500).unwrap_err();
-        let msg = err.to_string();
-        assert!(
-            msg.contains("exceeds the 500 byte limit"),
-            "Error should mention the size limit, got: {msg}"
-        );
-        assert!(
-            msg.contains("Use LIMIT or WHERE"),
-            "Error should suggest narrowing the query, got: {msg}"
-        );
+    fn test_validate_data_product_name_rejects_empty() {
+        assert!(validate_data_product_name("").is_err());
+        assert!(validate_data_product_name("   ").is_err());
     }
 
     #[mz_ore::test]
-    fn test_format_rows_response_empty_rows() {
-        let rows: Vec<Vec<serde_json::Value>> = vec![];
-        let result = format_rows_response(rows, 1000).unwrap();
-        let McpResult::ToolContent(content) = result else {
-            panic!("Expected ToolContent");
-        };
-        assert_eq!(content.content.len(), 1);
-        assert_eq!(content.content[0].text, "[]");
+    fn test_validate_data_product_name_rejects_sql_injection() {
+        // Attempted injection via semicolon
+        assert!(validate_data_product_name("my_view; DROP TABLE users").is_err());
+        // Attempted injection via subquery
+        assert!(validate_data_product_name("my_view UNION SELECT * FROM secrets").is_err());
+        // Multiple table references via comma
+        assert!(validate_data_product_name("my_view, secrets").is_err());
+        // SQL keywords
+        assert!(validate_data_product_name("my_view WHERE 1=1 --").is_err());
     }
 
     #[mz_ore::test]

src/environmentd/tests/server.rs

@@ -5353,8 +5353,8 @@ fn test_mcp_agents_with_data_product() {
         body["error"]["message"]
             .as_str()
             .unwrap()
-            .contains("Data product not found"),
-        "SQL injection in read_data_product name should be safely handled"
+            .contains("Invalid data product name"),
+        "SQL injection in read_data_product name should be rejected by validation"
     );
 
     // SQL injection attempt in cluster override.