diff --git a/libs/core/langchain_core/prompts/string.py b/libs/core/langchain_core/prompts/string.py index 2d6f76a693d..f34b2e294b0 100644 --- a/libs/core/langchain_core/prompts/string.py +++ b/libs/core/langchain_core/prompts/string.py @@ -20,65 +20,8 @@ if TYPE_CHECKING: try: from jinja2 import meta - from jinja2.exceptions import SecurityError from jinja2.sandbox import SandboxedEnvironment - class _RestrictedSandboxedEnvironment(SandboxedEnvironment): - """A more restrictive Jinja2 sandbox that blocks all attribute/method access. - - This sandbox only allows simple variable lookups, no attribute or method access. - This prevents template injection attacks via methods like parse_raw(). - """ - - def is_safe_attribute(self, _obj: Any, _attr: str, _value: Any) -> bool: - """Block ALL attribute access for security. - - Only allow accessing variables directly from the context dict, - no attribute access on those objects. - - Args: - _obj: The object being accessed (unused, always blocked). - _attr: The attribute name (unused, always blocked). - _value: The attribute value (unused, always blocked). - - Returns: - False - all attribute access is blocked. - """ - # Block all attribute access - return False - - def is_safe_callable(self, _obj: Any) -> bool: - """Block all method calls for security. - - Args: - _obj: The object being checked (unused, always blocked). - - Returns: - False - all callables are blocked. - """ - return False - - def getattr(self, obj: Any, attribute: str) -> Any: - """Override getattr to block all attribute access. - - Args: - obj: The object. - attribute: The attribute name. - - Returns: - Never returns. - - Raises: - SecurityError: Always, to block attribute access. - """ - msg = ( - f"Access to attributes is not allowed in templates. " - f"Attempted to access '{attribute}' on {type(obj).__name__}. " - f"Use only simple variable names like {{{{variable}}}} " - f"without dots or methods." - ) - raise SecurityError(msg) - _HAS_JINJA2 = True except ImportError: _HAS_JINJA2 = False @@ -121,7 +64,7 @@ def jinja2_formatter(template: str, /, **kwargs: Any) -> str: # Use a restricted sandbox that blocks ALL attribute/method access # Only simple variable lookups like {{variable}} are allowed # Attribute access like {{variable.attr}} or {{variable.method()}} is blocked - return _RestrictedSandboxedEnvironment().from_string(template).render(**kwargs) + return SandboxedEnvironment().from_string(template).render(**kwargs) def validate_jinja2(template: str, input_variables: list[str]) -> None: @@ -156,7 +99,7 @@ def _get_jinja2_variables_from_template(template: str) -> set[str]: "Please install it with `pip install jinja2`." ) raise ImportError(msg) - env = _RestrictedSandboxedEnvironment() + env = SandboxedEnvironment() ast = env.parse(template) return meta.find_undeclared_variables(ast) diff --git a/libs/core/tests/unit_tests/prompts/test_chat.py b/libs/core/tests/unit_tests/prompts/test_chat.py index 6d0c313f74b..0e2954cf5c3 100644 --- a/libs/core/tests/unit_tests/prompts/test_chat.py +++ b/libs/core/tests/unit_tests/prompts/test_chat.py @@ -1636,68 +1636,3 @@ def test_mustache_template_attribute_access_vulnerability() -> None: ) result_dict = prompt_dict.invoke({"person": {"name": "Alice"}}) assert result_dict.messages[0].content == "Alice" # type: ignore[attr-defined] - - -@pytest.mark.requires("jinja2") -def test_jinja2_template_attribute_access_is_blocked() -> None: - """Test that Jinja2 SandboxedEnvironment blocks dangerous attribute access. - - This test verifies that Jinja2's sandbox successfully blocks access to - dangerous dunder attributes like __class__, unlike Mustache. - - GOOD: Jinja2 SandboxedEnvironment raises SecurityError when attempting - to access __class__, __globals__, etc. This is expected behavior. - """ - msg = HumanMessage("howdy") - - # Create a Jinja2 template that attempts to access __class__.__name__ - prompt = ChatPromptTemplate.from_messages( - [("human", "{{question.__class__.__name__}}")], - template_format="jinja2", - ) - - # Jinja2 sandbox should block this with SecurityError - with pytest.raises(Exception, match="attribute") as exc_info: - prompt.invoke( - {"question": msg, "question.__class__.__name__": "safe_placeholder"} - ) - - # Verify it's a SecurityError from Jinja2 blocking __class__ access - error_msg = str(exc_info.value) - assert ( - "SecurityError" in str(type(exc_info.value)) - or "access to attribute '__class__'" in error_msg - ), f"Expected SecurityError blocking __class__, got: {error_msg}" - - -@pytest.mark.requires("jinja2") -def test_jinja2_blocks_all_attribute_access() -> None: - """Test that Jinja2 now blocks ALL attribute/method access for security. - - After the fix, Jinja2 uses _RestrictedSandboxedEnvironment which blocks - ALL attribute access, not just dunder attributes. This prevents the - parse_raw() vulnerability. - """ - msg = HumanMessage("test content") - - # Test 1: Simple variable access should still work - prompt_simple = ChatPromptTemplate.from_messages( - [("human", "Message: {{message}}")], - template_format="jinja2", - ) - result = prompt_simple.invoke({"message": "hello world"}) - assert "hello world" in result.messages[0].content # type: ignore[attr-defined] - - # Test 2: Attribute access should now be blocked (including safe attributes) - prompt_attr = ChatPromptTemplate.from_messages( - [("human", "Content: {{msg.content}}")], - template_format="jinja2", - ) - with pytest.raises(Exception, match="attribute") as exc_info: - prompt_attr.invoke({"msg": msg}) - - error_msg = str(exc_info.value) - assert ( - "SecurityError" in str(type(exc_info.value)) - or "Access to attributes is not allowed" in error_msg - ), f"Expected SecurityError blocking attribute access, got: {error_msg}"