160 lines
5.6 KiB
Python
160 lines
5.6 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from functools import lru_cache
|
|
|
|
import sqlglot
|
|
from sqlglot import exp
|
|
from sqlglot.errors import ParseError, TokenError
|
|
|
|
|
|
class SecurityError(Exception):
|
|
"""Base class for all query rejections."""
|
|
|
|
|
|
class SqlValidationError(SecurityError):
|
|
pass
|
|
|
|
|
|
class PermissionDeniedError(SecurityError):
|
|
pass
|
|
|
|
|
|
MAX_SQL_LENGTH = 200_000
|
|
# Doris first, MySQL as a fallback for constructs the Doris dialect misses.
|
|
PARSE_DIALECTS = ("doris", "mysql")
|
|
|
|
# exp.Query covers SELECT, set operations and subqueries.
|
|
_READ_ONLY_NODES = (exp.Query, exp.Show, exp.Describe, exp.Use)
|
|
_WRITE_NODES = (
|
|
exp.Insert,
|
|
exp.Update,
|
|
exp.Delete,
|
|
exp.Drop,
|
|
exp.Create,
|
|
exp.Alter,
|
|
exp.TruncateTable,
|
|
)
|
|
# Rejected regardless of allow_write: they read files, block workers or leak timing.
|
|
_FORBIDDEN_FUNCTIONS = frozenset({"load_file", "sleep", "benchmark", "outfile"})
|
|
|
|
|
|
@lru_cache(maxsize=256)
|
|
def parse_statement(sql: str) -> exp.Expression:
|
|
"""Parse a single statement, rejecting anything the dialects cannot model.
|
|
|
|
The returned tree is shared through the cache and must never be mutated.
|
|
"""
|
|
if not sql or not sql.strip():
|
|
raise SqlValidationError("Empty SQL statement.")
|
|
if len(sql) > MAX_SQL_LENGTH:
|
|
raise SqlValidationError(f"Statement exceeds {MAX_SQL_LENGTH} characters.")
|
|
if "\x00" in sql:
|
|
raise SqlValidationError("Statement contains a NUL byte.")
|
|
|
|
last_error: Exception | None = None
|
|
for dialect in PARSE_DIALECTS:
|
|
try:
|
|
trees = sqlglot.parse(sql, dialect=dialect)
|
|
except (ParseError, TokenError, RecursionError) as exc:
|
|
last_error = exc
|
|
continue
|
|
|
|
statements = [tree for tree in trees if tree is not None]
|
|
if not statements:
|
|
raise SqlValidationError("Empty SQL statement.")
|
|
if len(statements) > 1:
|
|
raise SqlValidationError("Multiple statements per request are not allowed.")
|
|
|
|
root = statements[0]
|
|
# sqlglot degrades DCL/ADMIN and other unmodelled syntax to Command.
|
|
if isinstance(root, exp.Command) or root.find(exp.Command) is not None:
|
|
raise SqlValidationError("Administrative or unsupported statement.")
|
|
return root
|
|
|
|
raise SqlValidationError(f"Statement could not be parsed: {last_error}")
|
|
|
|
|
|
def table_names(root: exp.Expression) -> set[str]:
|
|
"""Physical tables referenced by the tree, with CTE aliases removed."""
|
|
ctes = {cte.alias_or_name.lower() for cte in root.find_all(exp.CTE)}
|
|
names: set[str] = set()
|
|
for table in root.find_all(exp.Table):
|
|
name = table.name.lower()
|
|
if not name:
|
|
continue
|
|
parts = [part.lower() for part in (table.catalog, table.db) if part]
|
|
if not parts and name in ctes:
|
|
continue
|
|
names.add(".".join([*parts, name]))
|
|
return names
|
|
|
|
|
|
def extract_tables(sql: str) -> set[str]:
|
|
return table_names(parse_statement(sql))
|
|
|
|
|
|
def returns_rows(sql: str) -> bool:
|
|
return isinstance(parse_statement(sql), _READ_ONLY_NODES)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SqlGuard:
|
|
"""Statement level validation"""
|
|
|
|
allow_write: bool = False
|
|
max_rows: int = 1000
|
|
|
|
def validate(self, sql: str) -> str:
|
|
statement = sql.strip()
|
|
if statement.endswith(";"):
|
|
statement = statement[:-1].strip()
|
|
root = parse_statement(statement)
|
|
self._reject_dangerous(root)
|
|
|
|
if isinstance(root, _WRITE_NODES):
|
|
if not self.allow_write:
|
|
raise PermissionDeniedError(
|
|
f"{type(root).__name__.upper()} is rejected in read-only mode."
|
|
)
|
|
return statement
|
|
if isinstance(root, _READ_ONLY_NODES):
|
|
return self._apply_row_limit(statement, root)
|
|
raise SqlValidationError(f"Unsupported statement type: {type(root).__name__}.")
|
|
|
|
def _reject_dangerous(self, root: exp.Expression) -> None:
|
|
if root.find(exp.Into) is not None:
|
|
raise SqlValidationError("SELECT ... INTO is not allowed.")
|
|
for func in root.find_all(exp.Anonymous):
|
|
if func.name.lower() in _FORBIDDEN_FUNCTIONS:
|
|
raise SqlValidationError(f"Function '{func.name}' is not allowed.")
|
|
|
|
def _apply_row_limit(self, statement: str, root: exp.Expression) -> str:
|
|
"""Append a limit to the original text; the AST is only used to decide."""
|
|
if not isinstance(root, exp.Query) or root.args.get("limit") is not None:
|
|
return statement
|
|
return f"{statement} LIMIT {self.max_rows}"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PermissionPolicy:
|
|
"""Per-principal access rules, resolved before the statement runs."""
|
|
|
|
allowed_databases: frozenset[str] = frozenset()
|
|
denied_tables: frozenset[str] = frozenset()
|
|
allow_write: bool = False
|
|
|
|
def check(self, sql: str, default_database: str | None) -> None:
|
|
root = parse_statement(sql.strip().rstrip(";").strip())
|
|
if isinstance(root, _WRITE_NODES) and not self.allow_write:
|
|
raise PermissionDeniedError("Principal is not allowed to write.")
|
|
|
|
for table in table_names(root):
|
|
database, _, name = table.rpartition(".")
|
|
database = database or (default_database or "").lower()
|
|
if self.allowed_databases and database not in self.allowed_databases:
|
|
raise PermissionDeniedError(f"Database '{database}' is not accessible.")
|
|
qualified = f"{database}.{name}" if database else name
|
|
if name in self.denied_tables or qualified in self.denied_tables:
|
|
raise PermissionDeniedError(f"Table '{qualified}' is not accessible.")
|