diff --git a/open_notebook/domain/base.py b/open_notebook/domain/base.py index 45459d45..d2a74c0d 100644 --- a/open_notebook/domain/base.py +++ b/open_notebook/domain/base.py @@ -1,3 +1,4 @@ +import re from datetime import datetime from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar, Union, cast @@ -35,6 +36,39 @@ class ObjectModel(BaseModel): created: Optional[datetime] = None updated: Optional[datetime] = None + @classmethod + def _validate_order_by(cls, order_by: str) -> str: + """Validate and normalize an ORDER BY clause to prevent SurrealQL injection. + + Supports: "field", "field direction", "field1 direction, field2 direction". + Any subclass that builds its own query around `order_by` (instead of + delegating to `get_all()`) must route through this so the allowlist + can't silently drift between call sites. + """ + allowed_field_pattern = re.compile(r"^[a-z_][a-z0-9_]*$") + allowed_directions = {"asc", "desc"} + + clauses = [c.strip() for c in order_by.split(",")] + validated_clauses = [] + for clause in clauses: + parts = clause.strip().split() + if len(parts) == 1: + if not allowed_field_pattern.match(parts[0].lower()): + raise InvalidInputError(f"Invalid order_by field: '{parts[0]}'") + validated_clauses.append(parts[0].lower()) + elif len(parts) == 2: + if not allowed_field_pattern.match( + parts[0].lower() + ) or parts[1].lower() not in allowed_directions: + raise InvalidInputError( + f"Invalid order_by clause: '{clause.strip()}'" + ) + validated_clauses.append(f"{parts[0].lower()} {parts[1].lower()}") + else: + raise InvalidInputError(f"Invalid order_by clause: '{clause.strip()}'") + + return ", ".join(validated_clauses) + @classmethod async def get_all(cls: Type[T], order_by=None) -> List[T]: try: @@ -48,39 +82,7 @@ class ObjectModel(BaseModel): "get_all() must be called from a specific model class" ) if order_by: - # Validate order_by to prevent SurrealQL injection - # Supports: "field", "field direction", "field1 direction, field2 direction" - import re - - allowed_field_pattern = re.compile(r"^[a-z_][a-z0-9_]*$") - allowed_directions = {"asc", "desc"} - - clauses = [c.strip() for c in order_by.split(",")] - validated_clauses = [] - for clause in clauses: - parts = clause.strip().split() - if len(parts) == 1: - if not allowed_field_pattern.match(parts[0].lower()): - raise InvalidInputError( - f"Invalid order_by field: '{parts[0]}'" - ) - validated_clauses.append(parts[0].lower()) - elif len(parts) == 2: - if not allowed_field_pattern.match( - parts[0].lower() - ) or parts[1].lower() not in allowed_directions: - raise InvalidInputError( - f"Invalid order_by clause: '{clause.strip()}'" - ) - validated_clauses.append( - f"{parts[0].lower()} {parts[1].lower()}" - ) - else: - raise InvalidInputError( - f"Invalid order_by clause: '{clause.strip()}'" - ) - - validated_order_by = ", ".join(validated_clauses) + validated_order_by = cls._validate_order_by(order_by) query = f"SELECT * FROM {table_name} ORDER BY {validated_order_by}" else: query = f"SELECT * FROM {table_name}" diff --git a/open_notebook/domain/credential.py b/open_notebook/domain/credential.py index 49827976..57462f84 100644 --- a/open_notebook/domain/credential.py +++ b/open_notebook/domain/credential.py @@ -171,11 +171,12 @@ class Credential(ObjectModel): @classmethod async def get_all(cls, order_by=None) -> List["Credential"]: """Override get_all() to handle api_key decryption with per-row error handling.""" - order_clause = f" ORDER BY {order_by}" if order_by else "" - results = await repo_query( - f"SELECT * FROM {cls.table_name}{order_clause}", - {}, - ) + if order_by: + validated_order_by = cls._validate_order_by(order_by) + query = f"SELECT * FROM {cls.table_name} ORDER BY {validated_order_by}" + else: + query = f"SELECT * FROM {cls.table_name}" + results = await repo_query(query, {}) credentials = [] for row in results: try: diff --git a/tests/test_order_by_validation.py b/tests/test_order_by_validation.py new file mode 100644 index 00000000..2581d2fa --- /dev/null +++ b/tests/test_order_by_validation.py @@ -0,0 +1,50 @@ +""" +Regression tests for ObjectModel._validate_order_by() and its use by +Credential.get_all() (open_notebook/domain/credential.py), which builds its +own query instead of delegating to the base get_all() and previously +interpolated order_by into the SurrealQL unvalidated. +""" + +import pytest + +from open_notebook.domain.base import ObjectModel +from open_notebook.domain.credential import Credential +from open_notebook.exceptions import InvalidInputError + + +class TestValidateOrderBy: + @pytest.mark.parametrize( + "clause,expected", + [ + ("provider", "provider"), + ("provider, created", "provider, created"), + ("created DESC", "created desc"), + ("provider asc, created desc", "provider asc, created desc"), + ], + ) + def test_accepts_and_normalizes_valid_clauses(self, clause, expected): + assert ObjectModel._validate_order_by(clause) == expected + + @pytest.mark.parametrize( + "clause", + [ + "field; DROP TABLE credential", + "provider, created; REMOVE TABLE credential", + "provider) FETCH (SELECT * FROM credential", + "created LIMIT 1", + "provider desc extra", + "", + ], + ) + def test_rejects_injection_and_malformed_clauses(self, clause): + with pytest.raises(InvalidInputError): + ObjectModel._validate_order_by(clause) + + +class TestCredentialGetAllRejectsInjection: + @pytest.mark.asyncio + async def test_injection_in_order_by_raises_before_querying(self): + # Must raise on validation - reaching the database with this string + # (the pre-fix behavior) would execute the injected statement. + with pytest.raises(InvalidInputError): + await Credential.get_all(order_by="provider; DROP TABLE credential")