feat(safety): harden SQL validation (multi-CTE, recursive WITH, unicode normalization, precise errors, EXPLAIN gate)

nl2sql/safety.py

@@ -1,32 +1,56 @@
 from __future__ import annotations
+
 import re
 import time
+import unicodedata
 from nl2sql.types import StageResult, StageTrace
 
 # --- Regex utils ---
 _COMMENT_BLOCK = re.compile(r"/\*.*?\*/", re.DOTALL)
 _COMMENT_LINE = re.compile(r"--.*?$", re.MULTILINE)
-# string literals (single & double quotes), allow escaped quotes
+
+# String literals (single & double quotes), allow escaped quotes
 _STRING_SINGLE = re.compile(r"'([^'\\]|\\.)*'", re.DOTALL)
 _STRING_DOUBLE = re.compile(r'"([^"\\]|\\.)*"', re.DOTALL)
 
-# case-insensitive, word-boundary forbidden keywords
+# Case-insensitive, word-boundary forbidden keywords
 _FORBIDDEN = re.compile(
     r"\b(delete|update|insert|drop|create|alter|attach|pragma|reindex|vacuum|replace|grant|revoke|execute)\b",
     re.IGNORECASE,
 )
 
-# allow: SELECT ...   or   WITH <cte...> SELECT ...
-_ALLOW_SELECT = re.compile(r"^(?:WITH\b.*?\)\s*)?SELECT\b", re.IGNORECASE | re.DOTALL)
+# Allow: SELECT ...  or   WITH (one or many CTEs, optional RECURSIVE) ... SELECT ...
+_ALLOW_SELECT = re.compile(
+    r"^(?:WITH\s+(?:RECURSIVE\s+)?"
+    r".*?\)\s*(?:,\s*.*?\)\s*)*"
+    r")?SELECT\b",
+    re.IGNORECASE | re.DOTALL,
+)
+
+# Optional allowance: EXPLAIN SELECT ...
+_ALLOW_EXPLAIN_SELECT = re.compile(r"^EXPLAIN\s+SELECT\b", re.IGNORECASE | re.DOTALL)
 
-# --- New cleanup helpers ---
+# --- Cleanup helpers ---
 _FENCE_SQL = re.compile(r"```sql", re.IGNORECASE)
 _FENCE_ANY = re.compile(r"```")
 
 
+def _normalize_sql(sql: str) -> str:
+    """Normalize to NFKC and strip zero-width characters to prevent obfuscation."""
+    s = unicodedata.normalize("NFKC", sql)
+    # strip common zero-width spaces/joiners
+    return (
+        s.replace("\u200b", "")
+        .replace("\u200c", "")
+        .replace("\u200d", "")
+        .replace("\ufeff", "")
+    )
+
+
 def _sanitize_sql(sql: str) -> str:
-    """Remove markdown fences, comments, and surrounding junk."""
-    s = _FENCE_SQL.sub("", sql)
+    """Remove markdown fences, comments, and harmless trailing semicolons."""
+    s = _normalize_sql(sql)
+    s = _FENCE_SQL.sub("", s)
     s = _FENCE_ANY.sub("", s)
     s = _COMMENT_BLOCK.sub(" ", s)
     s = _COMMENT_LINE.sub(" ", s)
@@ -37,6 +61,7 @@ def _sanitize_sql(sql: str) -> str:
 
 
 def _mask_strings(s: str) -> str:
+    """Replace string literals so that inner semicolons/keywords don't affect checks."""
     s = _STRING_SINGLE.sub("'X'", s)
     s = _STRING_DOUBLE.sub('"X"', s)
     return s
@@ -44,64 +69,76 @@ def _mask_strings(s: str) -> str:
 
 def _split_statements(s: str) -> list[str]:
     """
-    Split only if there are real multiple statements,
-    ignoring harmless trailing semicolons or markdown.
+    Split on semicolons after string-masking. Ignore empties (e.g., trailing ';').
     """
     parts = [p.strip() for p in s.split(";")]
-    parts = [p for p in parts if p]
-    return parts
+    return [p for p in parts if p]
+
+
+def _ms(t0: float) -> int:
+    return int((time.perf_counter() - t0) * 1000)
 
 
 class Safety:
     name = "safety"
 
+    def __init__(self, allow_explain: bool = False) -> None:
+        """
+        :param allow_explain: If True, 'EXPLAIN SELECT ...' is allowed in addition to SELECT.
+        """
+        self.allow_explain = allow_explain
+
     def check(self, sql: str) -> StageResult:
         t0 = time.perf_counter()
-        print("🧩 SQL candidate:", sql)
 
-        # --- sanitize first ---
+        # 1) Sanitize and mask
         s = _sanitize_sql(sql)
         s = _mask_strings(s).strip()
 
+        # 2) Multiple statements check
         stmts = _split_statements(s)
         if len(stmts) != 1:
             return StageResult(
                 ok=False,
                 error=["Multiple statements detected"],
-                trace=StageTrace(
-                    stage=self.name, duration_ms=(time.perf_counter() - t0) * 1000
-                ),
+                trace=StageTrace(stage=self.name, duration_ms=_ms(t0)),
             )
 
         body = stmts[0]
 
-        if _FORBIDDEN.search(body):
+        # 3) Forbidden keyword check (report exact offending token)
+        m = _FORBIDDEN.search(body)
+        if m:
             return StageResult(
                 ok=False,
-                error=["Forbidden keyword detected"],
-                trace=StageTrace(
-                    stage=self.name, duration_ms=(time.perf_counter() - t0) * 1000
-                ),
+                error=[f"Forbidden keyword detected: '{m.group(0)}'"],
+                trace=StageTrace(stage=self.name, duration_ms=_ms(t0)),
             )
 
-        if not _ALLOW_SELECT.match(body):
+        # 4) Allow only SELECT (or optionally EXPLAIN SELECT)
+        allowed = bool(_ALLOW_SELECT.match(body))
+        if not allowed and self.allow_explain:
+            allowed = bool(_ALLOW_EXPLAIN_SELECT.match(body))
+
+        if not allowed:
             return StageResult(
                 ok=False,
                 error=["Non-SELECT statement"],
-                trace=StageTrace(
-                    stage=self.name, duration_ms=(time.perf_counter() - t0) * 1000
-                ),
+                trace=StageTrace(stage=self.name, duration_ms=_ms(t0)),
             )
 
+        # 5) Success
         return StageResult(
             ok=True,
             data={
                 "sql": body,
-                "rationale": "Statement validated as SELECT-only (strings/comments/markdown ignored).",
+                "rationale": (
+                    "Statement validated as SELECT-only (strings/comments/markdown ignored)."
+                    + (" EXPLAIN SELECT allowed." if self.allow_explain else "")
+                ),
             },
-            trace=StageTrace(
-                stage=self.name, duration_ms=(time.perf_counter() - t0) * 1000
-            ),
+            trace=StageTrace(stage=self.name, duration_ms=_ms(t0)),
         )
 
+    # Backward-compat alias
     run = check

tests/test_safety.py

@@ -119,3 +119,124 @@ def test_safety_blocks_multiple_nonempty_statements_even_if_second_is_comment():
     sql_bad = "SELECT 1;  /* spacer */  DROP TABLE x;"
     assert s.check(sql).ok
     assert not s.check(sql_bad).ok
+
+
+def test_safety_allows_multiple_ctes():
+    s = Safety()
+    sql = """
+    WITH a AS (SELECT 1 AS x),
+         b AS (SELECT 2 AS y)
+    SELECT a.x, b.y FROM a CROSS JOIN b;
+    """
+    assert s.check(sql).ok
+
+
+def test_safety_allows_with_recursive():
+    s = Safety()
+    sql = """
+    WITH RECURSIVE cnt(x) AS (
+      SELECT 1 UNION ALL SELECT x+1 FROM cnt WHERE x < 3
+    )
+    SELECT * FROM cnt;
+    """
+    assert s.check(sql).ok
+
+
+def test_safety_blocks_zero_width_obfuscation_in_keyword():
+    s = Safety()
+    # "DROP" با zero-width joiner وسط حروف
+    bad = "DR\u200dOP TABLE users;"
+    r = s.check(bad)
+    assert not r.ok
+
+
+def test_safety_ignores_markdown_fences():
+    s = Safety()
+    sql = "```sql\nSELECT 1;\n```"
+    assert s.check(sql).ok
+
+
+def test_safety_semicolon_inside_string_literal_is_ignored():
+    s = Safety()
+    sql = "SELECT 'a; b; c' AS sample;"
+    assert s.check(sql).ok
+
+
+def test_safety_forbidden_keyword_inside_string_literal_ok():
+    s = Safety()
+    sql = "SELECT 'DROP TABLE x' AS note, 'delete from y' AS text;"
+    assert s.check(sql).ok
+
+
+def test_safety_reports_offending_token_in_error_message():
+    s = Safety()
+    r = s.check("  \n  ReIndex  users;")
+    assert not r.ok
+    assert any("reindex" in e.lower() for e in (r.error or []))
+
+
+def test_safety_multiple_statements_with_masked_strings_is_blocked():
+    s = Safety()
+    sql = "SELECT 'abc'; SELECT 1;"
+    r = s.check(sql)
+    assert not r.ok
+
+
+def test_safety_duration_ms_is_int():
+    s = Safety()
+    r = s.check("SELECT 1;")
+    assert isinstance(r.trace.duration_ms, int)
+
+
+def test_safety_allows_explain_select_when_enabled():
+    s = Safety(allow_explain=True)
+    r = s.check("EXPLAIN SELECT * FROM users;")
+    assert r.ok
+
+
+def test_safety_blocks_explain_select_when_disabled():
+    s = Safety(allow_explain=False)
+    r = s.check("EXPLAIN SELECT * FROM users;")
+    assert not r.ok
+
+
+def test_safety_blocks_forbidden_inside_cte_body():
+    s = Safety()
+    sql = """
+    WITH bad AS (DELETE FROM users)
+    SELECT * FROM users;
+    """
+    assert not s.check(sql).ok
+
+
+def test_safety_permits_with_comments_and_newlines_complex():
+    s = Safety()
+    sql = """
+    /* head */ WITH a AS (SELECT 1 /*x*/ AS x) -- inline
+    , b AS (SELECT 2 AS y) /* tail */
+    SELECT a.x, b.y FROM a JOIN b; -- end
+    """
+    assert s.check(sql).ok
+
+
+def test_safety_blocks_bom_prefixed_forbidden():
+    s = Safety()
+    sql = "\ufeffDROP TABLE x;"
+    assert not s.check(sql).ok
+
+
+def test_safety_allows_trailing_double_semicolon():
+    s = Safety()
+    assert s.check("SELECT 1;;").ok
+
+
+@pytest.mark.parametrize("q", ["explain   select 1;", "EXPLAIN\nSELECT 1;"])
+def test_safety_explain_various_spacing_when_enabled(q):
+    s = Safety(allow_explain=True)
+    assert s.check(q).ok
+
+
+def test_safety_stage_name_constant():
+    s = Safety()
+    r = s.check("SELECT 1;")
+    assert r.trace.stage == "safety"