379 lines
14 KiB
Python
379 lines
14 KiB
Python
|
"""
|
||
|
|
Copyright 2024-2026 ChatterMate
|
||
|
|
|
||
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
you may not use this file except in compliance with the License.
|
||
|
|
You may obtain a copy of the License at
|
||
|
|
|
||
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
|
||
|
|
Unless required by applicable law or agreed to in writing, software
|
||
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
See the License for the specific language governing permissions and
|
||
|
|
limitations under the License.
|
||
|
|
|
||
|
|
Injection corpus for the DB-connector SQL guardrails. Every mutation path,
|
||
|
|
smuggling trick and resource-abuse vector here must stay blocked.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from app.services.sql_guardrails import MASK_VALUE, mask_rows, validate_sql
|
||
|
|
|
||
|
|
ALLOWED = ["public.orders", "public.customers", "billing.invoices"]
|
||
|
|
|
||
|
|
|
||
|
|
def check(sql, allowed=ALLOWED, max_rows=100, engine="postgresql", masked=None):
|
||
|
|
return validate_sql(sql, allowed, max_rows, engine, masked_columns=masked)
|
||
|
|
|
||
|
|
|
||
|
|
class TestAllowedSelects:
|
||
|
|
def test_simple_select(self):
|
||
|
|
result = check("SELECT id, status FROM orders WHERE id = 42")
|
||
|
|
assert result.ok
|
||
|
|
assert "LIMIT 100" in result.sql
|
||
|
|
|
||
|
|
def test_schema_qualified(self):
|
||
|
|
assert check("SELECT * FROM billing.invoices").ok
|
||
|
|
|
||
|
|
def test_join_of_allowed_tables(self):
|
||
|
|
result = check(
|
||
|
|
"SELECT o.id FROM orders o JOIN customers c ON c.id = o.customer_id"
|
||
|
|
)
|
||
|
|
assert result.ok
|
||
|
|
|
||
|
|
def test_cte_over_allowed_tables(self):
|
||
|
|
result = check(
|
||
|
|
"WITH recent AS (SELECT * FROM orders WHERE created_at > now() - interval '1 day') "
|
||
|
|
"SELECT count(*) FROM recent"
|
||
|
|
)
|
||
|
|
assert result.ok
|
||
|
|
|
||
|
|
def test_subquery(self):
|
||
|
|
assert check(
|
||
|
|
"SELECT * FROM orders WHERE customer_id IN (SELECT id FROM customers)"
|
||
|
|
).ok
|
||
|
|
|
||
|
|
def test_existing_smaller_limit_kept(self):
|
||
|
|
result = check("SELECT id FROM orders LIMIT 5")
|
||
|
|
assert result.ok
|
||
|
|
assert "LIMIT 5" in result.sql
|
||
|
|
|
||
|
|
def test_oversized_limit_clamped(self):
|
||
|
|
result = check("SELECT id FROM orders LIMIT 500000")
|
||
|
|
assert result.ok
|
||
|
|
assert "LIMIT 100" in result.sql
|
||
|
|
assert "500000" not in result.sql
|
||
|
|
|
||
|
|
|
||
|
|
class TestMutationsBlocked:
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"sql",
|
||
|
|
[
|
||
|
|
"UPDATE orders SET status = 'paid'",
|
||
|
|
"DELETE FROM orders",
|
||
|
|
"INSERT INTO orders (id) VALUES (1)",
|
||
|
|
"DROP TABLE orders",
|
||
|
|
"TRUNCATE orders",
|
||
|
|
"CREATE TABLE x (id int)",
|
||
|
|
"ALTER TABLE orders ADD COLUMN x int",
|
||
|
|
"GRANT ALL ON orders TO public",
|
||
|
|
"MERGE INTO orders USING customers ON true WHEN MATCHED THEN UPDATE SET status='x'",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_blocked(self, sql):
|
||
|
|
assert not check(sql).ok
|
||
|
|
|
||
|
|
def test_multi_statement(self):
|
||
|
|
assert not check("SELECT 1; DROP TABLE orders").ok
|
||
|
|
|
||
|
|
def test_stacked_after_comment(self):
|
||
|
|
assert not check("SELECT id FROM orders; -- x\nDELETE FROM orders").ok
|
||
|
|
|
||
|
|
def test_select_into(self):
|
||
|
|
assert not check("SELECT * INTO evil FROM orders").ok
|
||
|
|
|
||
|
|
def test_select_for_update(self):
|
||
|
|
assert not check("SELECT * FROM orders FOR UPDATE").ok
|
||
|
|
|
||
|
|
|
||
|
|
class TestTableSmuggling:
|
||
|
|
def test_unlisted_table(self):
|
||
|
|
assert not check("SELECT * FROM users").ok
|
||
|
|
|
||
|
|
def test_unlisted_in_join(self):
|
||
|
|
assert not check("SELECT * FROM orders JOIN users ON true").ok
|
||
|
|
|
||
|
|
def test_unlisted_in_subquery(self):
|
||
|
|
assert not check(
|
||
|
|
"SELECT * FROM orders WHERE id IN (SELECT order_id FROM secrets)"
|
||
|
|
).ok
|
||
|
|
|
||
|
|
def test_unlisted_inside_cte_body(self):
|
||
|
|
# The classic smuggle: hide the real table inside a CTE definition.
|
||
|
|
assert not check(
|
||
|
|
"WITH x AS (SELECT * FROM pg_shadow) SELECT * FROM x"
|
||
|
|
).ok
|
||
|
|
|
||
|
|
def test_cte_alias_cannot_authorize_schema_reference(self):
|
||
|
|
assert not check(
|
||
|
|
"WITH users AS (SELECT 1) SELECT * FROM secret_schema.users"
|
||
|
|
).ok
|
||
|
|
|
||
|
|
def test_union_smuggle_blocked_at_top_level(self):
|
||
|
|
assert not check("SELECT id FROM orders UNION SELECT usename FROM pg_shadow").ok
|
||
|
|
|
||
|
|
def test_union_inside_subquery_still_walked(self):
|
||
|
|
assert not check(
|
||
|
|
"SELECT * FROM (SELECT id FROM orders UNION ALL SELECT oid FROM pg_shadow) t"
|
||
|
|
).ok
|
||
|
|
|
||
|
|
def test_system_catalog(self):
|
||
|
|
assert not check("SELECT * FROM information_schema.tables").ok
|
||
|
|
|
||
|
|
def test_bare_name_matches_public_schema_entry(self):
|
||
|
|
assert check("SELECT * FROM invoices", allowed=["public.invoices"]).ok
|
||
|
|
|
||
|
|
def test_wrong_schema_same_name(self):
|
||
|
|
assert not check("SELECT * FROM hidden.orders").ok
|
||
|
|
|
||
|
|
|
||
|
|
class TestFunctionDenylist:
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"sql",
|
||
|
|
[
|
||
|
|
"SELECT pg_sleep(10) FROM orders",
|
||
|
|
"SELECT pg_read_file('/etc/passwd') FROM orders",
|
||
|
|
"SELECT * FROM orders WHERE id = (SELECT pg_terminate_backend(1))",
|
||
|
|
"SELECT set_config('x', 'y', false) FROM orders",
|
||
|
|
"SELECT lo_import('/etc/passwd') FROM orders",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_postgres_functions(self, sql):
|
||
|
|
assert not check(sql).ok
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"sql",
|
||
|
|
[
|
||
|
|
"SELECT sleep(10) FROM orders",
|
||
|
|
"SELECT benchmark(1000000, md5('x')) FROM orders",
|
||
|
|
"SELECT load_file('/etc/passwd') FROM orders",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_mysql_functions(self, sql):
|
||
|
|
assert not check(sql, engine="mysql").ok
|
||
|
|
|
||
|
|
|
||
|
|
class TestMaskedColumns:
|
||
|
|
def test_direct_reference_blocked(self):
|
||
|
|
assert not check("SELECT email FROM customers", masked=["email"]).ok
|
||
|
|
|
||
|
|
def test_alias_smuggle_blocked(self):
|
||
|
|
assert not check("SELECT email AS harmless FROM customers", masked=["email"]).ok
|
||
|
|
|
||
|
|
def test_expression_smuggle_blocked(self):
|
||
|
|
assert not check(
|
||
|
|
"SELECT substr(email, 1, 3) FROM customers", masked=["email"]
|
||
|
|
).ok
|
||
|
|
|
||
|
|
def test_where_probe_blocked(self):
|
||
|
|
assert not check(
|
||
|
|
"SELECT id FROM customers WHERE email = 'a@b.c'", masked=["email"]
|
||
|
|
).ok
|
||
|
|
|
||
|
|
def test_unmasked_columns_fine(self):
|
||
|
|
assert check("SELECT id, full_name FROM customers", masked=["email"]).ok
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"sql",
|
||
|
|
[
|
||
|
|
# Whole-row serialization collapses masked fields into one value
|
||
|
|
# that name-based masking can't see — all must be blocked.
|
||
|
|
"SELECT to_jsonb(t) FROM customers t",
|
||
|
|
"SELECT row_to_json(customers) FROM customers",
|
||
|
|
"SELECT customers FROM customers",
|
||
|
|
"SELECT array_agg(customers) FROM customers",
|
||
|
|
"SELECT hstore(customers) FROM customers",
|
||
|
|
"SELECT CAST(customers AS text) FROM customers",
|
||
|
|
"SELECT customers::text FROM customers",
|
||
|
|
"SELECT to_json(c.*) FROM customers c",
|
||
|
|
"SELECT (c.*)::text FROM customers c",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_whole_row_serialization_blocked(self, sql):
|
||
|
|
assert not check(sql, masked=["email"]).ok
|
||
|
|
|
||
|
|
def test_whole_row_allowed_without_masking(self):
|
||
|
|
# No masked columns → whole-row serialization is fine.
|
||
|
|
assert check("SELECT to_jsonb(t) FROM customers t", masked=[]).ok
|
||
|
|
|
||
|
|
def test_star_and_count_still_allowed_with_masking(self):
|
||
|
|
# SELECT * / t.* expand to real column names (mask_rows redacts them);
|
||
|
|
# count(*) is a cardinality marker, not row data.
|
||
|
|
assert check("SELECT * FROM customers", masked=["email"]).ok
|
||
|
|
assert check("SELECT c.* FROM customers c", masked=["email"]).ok
|
||
|
|
assert check("SELECT count(*) FROM customers", masked=["email"]).ok
|
||
|
|
assert check("SELECT array_agg(id) FROM customers", masked=["email"]).ok
|
||
|
|
|
||
|
|
def test_mask_rows_output(self):
|
||
|
|
rows = mask_rows(
|
||
|
|
["id", "Email", "name"],
|
||
|
|
[(1, "a@b.c", "Ann"), (2, "c@d.e", "Cy")],
|
||
|
|
["email"],
|
||
|
|
)
|
||
|
|
assert rows[0] == (1, MASK_VALUE, "Ann")
|
||
|
|
assert rows[1][1] == MASK_VALUE
|
||
|
|
|
||
|
|
def test_mask_rows_no_masked_columns(self):
|
||
|
|
rows = [(1, "x")]
|
||
|
|
assert mask_rows(["id", "v"], rows, []) == rows
|
||
|
|
|
||
|
|
|
||
|
|
class TestEdgeCases:
|
||
|
|
def test_garbage(self):
|
||
|
|
assert not check("DRIBBLE FRUM x!!!").ok
|
||
|
|
|
||
|
|
def test_empty(self):
|
||
|
|
assert not check("").ok
|
||
|
|
|
||
|
|
def test_empty_allowlist(self):
|
||
|
|
assert not check("SELECT 1", allowed=[]).ok
|
||
|
|
|
||
|
|
def test_canonical_rerender(self):
|
||
|
|
# Whatever executes came from the AST, not the raw string.
|
||
|
|
result = check("select id \n from orders /* comment */ where id=1")
|
||
|
|
assert result.ok
|
||
|
|
assert "/*" not in result.sql
|
||
|
|
|
||
|
|
|
||
|
|
ROW_SCOPE = {"public.orders": "customer_email", "public.customers": "email"}
|
||
|
|
OWNER = "arun@chattermate.chat"
|
||
|
|
|
||
|
|
|
||
|
|
def scoped(sql, scope=None, value=OWNER, allowed=ALLOWED, masked=None):
|
||
|
|
return validate_sql(
|
||
|
|
sql, allowed, 100, "postgresql",
|
||
|
|
masked_columns=masked,
|
||
|
|
row_scope=ROW_SCOPE if scope is None else scope,
|
||
|
|
scope_value=value,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class TestRowScoping:
|
||
|
|
"""Scoped tables are rewritten into a pre-filtered subquery, so the rows
|
||
|
|
the agent can see are constrained by the relation, not by its own WHERE."""
|
||
|
|
|
||
|
|
def test_scoped_table_is_wrapped(self):
|
||
|
|
result = scoped("SELECT id, total FROM orders")
|
||
|
|
assert result.ok
|
||
|
|
assert "customer_email = '%s'" % OWNER in result.sql
|
||
|
|
assert "FROM (SELECT * FROM orders WHERE" in result.sql
|
||
|
|
|
||
|
|
def test_alias_is_preserved(self):
|
||
|
|
result = scoped("SELECT o.id FROM orders o WHERE o.total > 10")
|
||
|
|
assert result.ok
|
||
|
|
assert ") AS o" in result.sql
|
||
|
|
assert "o.total > 10" in result.sql
|
||
|
|
|
||
|
|
def test_unscoped_table_untouched(self):
|
||
|
|
result = validate_sql(
|
||
|
|
"SELECT id FROM billing.invoices", ALLOWED, 100, "postgresql",
|
||
|
|
row_scope=ROW_SCOPE, scope_value=OWNER,
|
||
|
|
)
|
||
|
|
assert result.ok
|
||
|
|
assert "WHERE" not in result.sql.upper().replace("LIMIT", "")
|
||
|
|
|
||
|
|
def test_or_true_cannot_widen_the_scope(self):
|
||
|
|
# The classic bypass: the outer predicate is irrelevant because the
|
||
|
|
# relation it selects from was already filtered.
|
||
|
|
result = scoped("SELECT id FROM orders WHERE 1=1 OR customer_email = 'victim@x.com'")
|
||
|
|
assert result.ok
|
||
|
|
assert "FROM (SELECT * FROM orders WHERE customer_email = '%s')" % OWNER in result.sql
|
||
|
|
|
||
|
|
def test_join_scopes_every_scoped_table(self):
|
||
|
|
result = scoped(
|
||
|
|
"SELECT o.id, c.name FROM orders o JOIN customers c ON c.id = o.customer_id"
|
||
|
|
)
|
||
|
|
assert result.ok
|
||
|
|
assert result.sql.count("customer_email = '%s'" % OWNER) == 1
|
||
|
|
assert "email = '%s'" % OWNER in result.sql
|
||
|
|
|
||
|
|
def test_cte_body_is_scoped(self):
|
||
|
|
result = scoped("WITH recent AS (SELECT * FROM orders) SELECT id FROM recent")
|
||
|
|
assert result.ok
|
||
|
|
assert "customer_email = '%s'" % OWNER in result.sql
|
||
|
|
|
||
|
|
def test_subquery_is_scoped(self):
|
||
|
|
result = scoped("SELECT id FROM orders WHERE id IN (SELECT id FROM orders)")
|
||
|
|
assert result.ok
|
||
|
|
assert result.sql.count("customer_email = '%s'" % OWNER) == 2
|
||
|
|
|
||
|
|
def test_fails_closed_without_a_scope_value(self):
|
||
|
|
# Never silently unscoped — that would read every customer's rows.
|
||
|
|
for value in (None, "", " "):
|
||
|
|
result = scoped("SELECT id FROM orders", value=value)
|
||
|
|
assert not result.ok
|
||
|
|
assert "no customer identity" in result.reason
|
||
|
|
|
||
|
|
def test_unscoped_table_still_queryable_without_a_value(self):
|
||
|
|
result = scoped("SELECT id FROM billing.invoices", value=None)
|
||
|
|
assert result.ok
|
||
|
|
|
||
|
|
def test_scope_value_is_escaped_not_interpolated(self):
|
||
|
|
result = scoped("SELECT id FROM orders", value="x' OR '1'='1")
|
||
|
|
assert result.ok
|
||
|
|
# Doubled quotes — the value stays one string literal.
|
||
|
|
assert "'x'' OR ''1''=''1'" in result.sql
|
||
|
|
|
||
|
|
def test_no_scope_config_is_a_no_op(self):
|
||
|
|
result = scoped("SELECT id FROM orders", scope={}, value=None)
|
||
|
|
assert result.ok
|
||
|
|
assert "SELECT id FROM orders" in result.sql
|
||
|
|
|
||
|
|
def test_bare_reference_matches_qualified_rule(self):
|
||
|
|
# Rule is on "public.orders"; the query says "orders".
|
||
|
|
result = scoped("SELECT id FROM public.orders")
|
||
|
|
assert result.ok
|
||
|
|
assert "customer_email = '%s'" % OWNER in result.sql
|
||
|
|
|
||
|
|
def test_masked_columns_still_enforced_under_scoping(self):
|
||
|
|
result = scoped("SELECT email FROM orders", masked=["email"])
|
||
|
|
assert not result.ok
|
||
|
|
assert "masked" in result.reason.lower()
|
||
|
|
|
||
|
|
def test_limit_applies_to_the_outer_query(self):
|
||
|
|
result = scoped("SELECT id FROM orders")
|
||
|
|
assert result.ok
|
||
|
|
assert result.sql.rstrip().endswith("LIMIT 100")
|
||
|
|
|
||
|
|
|
||
|
|
class TestRowScopeCaseInsensitivity:
|
||
|
|
"""Emails fold case; the same person may be Bob@x.com in one system and
|
||
|
|
bob@x.com in the other, and an exact match would return nothing."""
|
||
|
|
|
||
|
|
def _ci(self, value=OWNER):
|
||
|
|
return validate_sql(
|
||
|
|
"SELECT id FROM orders", ALLOWED, 100, "postgresql",
|
||
|
|
row_scope=ROW_SCOPE, scope_value=value, scope_case_insensitive=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
def test_predicate_folds_both_sides(self):
|
||
|
|
result = self._ci("Arun@ChatterMate.Chat")
|
||
|
|
assert result.ok
|
||
|
|
assert "LOWER(customer_email) = '%s'" % OWNER in result.sql
|
||
|
|
|
||
|
|
def test_value_is_still_escaped(self):
|
||
|
|
result = self._ci("x' OR '1'='1")
|
||
|
|
assert result.ok
|
||
|
|
assert "'x'' or ''1''=''1'" in result.sql.lower()
|
||
|
|
|
||
|
|
def test_exact_match_when_disabled(self):
|
||
|
|
result = scoped("SELECT id FROM orders", value="Arun@ChatterMate.Chat")
|
||
|
|
assert result.ok
|
||
|
|
assert "LOWER(" not in result.sql
|
||
|
|
assert "customer_email = 'Arun@ChatterMate.Chat'" in result.sql
|
||
|
|
|
||
|
|
def test_still_fails_closed(self):
|
||
|
|
result = self._ci(value=None)
|
||
|
|
assert not result.ok
|