Harden our Jinja sandboxing

@@ -87,7 +87,7 @@     "hf-doc-builder",     "huggingface-hub>=1.31.0,<3.0",     "ipadic>=1.0.0,<2.0",-    "jinja2>=3.1.0",+    "jinja2>=3.1.6",     "kenlm",     "kernels>=0.18.0,<0.19",     "librosa",

@@ -19,7 +19,7 @@     "hf-doc-builder": "hf-doc-builder",     "huggingface-hub": "huggingface-hub>=1.31.0,<3.0",     "ipadic": "ipadic>=1.0.0,<2.0",-    "jinja2": "jinja2>=3.1.0",+    "jinja2": "jinja2>=3.1.6",     "kenlm": "kenlm",     "kernels": "kernels>=0.18.0,<0.19",     "librosa": "librosa",

@@ -12,18 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +import dataclasses import inspect import json import re import types-from collections.abc import Callable+from collections.abc import Callable, Mapping from contextlib import contextmanager-from copy import deepcopy from datetime import datetime from functools import lru_cache, partial from inspect import isfunction from typing import Any, Literal, Union, get_args, get_origin, get_type_hints, no_type_check +import numpy as np from packaging import version  from . import logging@@ -579,9 +580,11 @@ def activate_tracker(self, rendered_blocks: list[int], generation_indices: list[                 self._rendered_blocks = None                 self._generation_indices = None -    if version.parse(jinja2.__version__) < version.parse("3.1.0"):+    if version.parse(jinja2.__version__) < version.parse("3.1.6"):+        # Older versions have known sandbox escapes, and chat templates from the Hub are untrusted code         raise ImportError(-            f"apply_chat_template requires jinja2>=3.1.0 to be installed. Your version is {jinja2.__version__}."+            f"apply_chat_template requires jinja2>=3.1.6 to be installed. Your version is {jinja2.__version__}. "+            "You can upgrade it with `pip install -U jinja2`."         )      def raise_exception(message):@@ -595,12 +598,45 @@ def tojson(x, ensure_ascii=False, indent=None, separators=None, sort_keys=False)     def strftime_now(format):         return datetime.now().strftime(format) +    class TemplateDict:+        # Make a limited subclass so that users can't break out with dictionary attribute access+        # or type-hackery+        def __call__(self, *args, **kwargs):+            return dict(*args, **kwargs)++        def fromkeys(self, *args, **kwargs):+            return dict.fromkeys(*args, **kwargs)++    class TemplateNamespace(jinja2.utils.Namespace):+        # Namespaces reject any dunder attributes to prevent access to CPython internals+        def __init__(*args, **kwargs):+            # Same signature trick as Jinja's Namespace, so that `namespace(self=...)` still works+            self, args = args[0], args[1:]+            attrs = dict(*args, **kwargs)+            for name in attrs:+                TemplateNamespace._check_name(name)+            jinja2.utils.Namespace.__init__(self, attrs)++        def __setitem__(self, name, value):+            TemplateNamespace._check_name(name)+            jinja2.utils.Namespace.__setitem__(self, name, value)++        @staticmethod+        def _check_name(name):+            if isinstance(name, str) and name.startswith("_"):+                raise jinja2.exceptions.SecurityError(+                    f"Chat templates cannot store underscore-prefixed attributes on a namespace, but got {name!r}. "+                    "Please rename the attribute."+                )+     jinja_env = ImmutableSandboxedEnvironment(         trim_blocks=True, lstrip_blocks=True, extensions=[AssistantTracker, jinja2.ext.loopcontrols]     )     jinja_env.filters["tojson"] = tojson     jinja_env.globals["raise_exception"] = raise_exception     jinja_env.globals["strftime_now"] = strftime_now+    jinja_env.globals["dict"] = TemplateDict()+    jinja_env.globals["namespace"] = TemplateNamespace     return jinja_env.from_string(chat_template)  @@ -620,6 +656,42 @@ def _alias_reasoning_fields(chat: ChatType) -> ChatType:     return aliased_chat  +def _plain_template_data(value: Any) -> Any:+    # Chat templates come from the Hub, but callers often put rich objects in their inputs, like the decoded images,+    # arrays and tensors in multimodal content blocks. Those objects have methods that the sandbox doesn't know are+    # dangerous (some of them write files!), so we rebuild everything the template sees out of plain JSON types,+    # and replace anything else with an inert placeholder. Templates only check whether media keys are present,+    # and never read their values, so this doesn't change what they render.+    if value is None or type(value) in (str, int, float, bool):+        # Exact JSON scalars are safe as-is+        return value+    elif isinstance(value, str):+        # Subclasses (enums, numpy scalars, ...) can carry extra methods, so we strip them back to the base type+        return str.__str__(value)+    elif isinstance(value, int):+        return int.__int__(value)+    elif isinstance(value, float):+        return float.__float__(value)+    elif isinstance(value, Mapping):+        return {_plain_template_data(key): _plain_template_data(item) for key, item in value.items()}+    elif isinstance(value, list):+        return [_plain_template_data(item) for item in value]+    elif isinstance(value, tuple):+        return tuple(_plain_template_data(item) for item in value)+    elif isinstance(value, np.generic):+        # `np.bool_(False)` would otherwise become a truthy placeholder string+        return _plain_template_data(value.item())+    elif dataclasses.is_dataclass(value) and not isinstance(value, type):+        # Fields are read one level at a time, so arrays and images inside are not deep-copied+        return {field.name: _plain_template_data(getattr(value, field.name)) for field in dataclasses.fields(value)}+    elif callable(getattr(value, "model_dump", None)):+        # Pydantic models+        return _plain_template_data(value.model_dump())+    else:+        # Anything else becomes a placeholder naming its type, e.g. "<ndarray>"+        return f"<{type(value).__name__}>"++ def render_jinja_template(     conversations: list[ChatType],     tools: list[dict | Callable] | None = None,@@ -659,16 +731,25 @@ def render_jinja_template(             if not isinstance(document, dict):                 raise TypeError("Documents should be a list of dicts with 'title' and 'text' keys!") +    # Everything the template can see goes through `_plain_template_data`. The conversations are converted below.+    tool_schemas = _plain_template_data(tool_schemas)+    documents = _plain_template_data(documents)+    kwargs = _plain_template_data(kwargs)+     rendered = []     all_generation_indices = []     continue_final_message_tag = "CONTINUE_FINAL_MESSAGE_TAG "     for chat in conversations:         if hasattr(chat, "messages"):             # Indicates it's a Conversation object             chat = chat.messages+        # The converted chat is a fresh copy, so the edits below never touch the caller's messages+        chat = _plain_template_data(chat)+        for message in chat:+            if not isinstance(message, dict):+                raise TypeError(f"Chat messages should be dicts, dataclasses or pydantic models, but got {message!r}.")         chat = _alias_reasoning_fields(chat)         if continue_final_message:-            chat = deepcopy(chat)             continue_final_message = continue_final_message if isinstance(continue_final_message, str) else "content"              if (final_message := chat[-1].get(continue_final_message)) is None:

@@ -13,9 +13,13 @@ # limitations under the License.  import copy+import dataclasses import unittest+from enum import Enum from typing import Literal +import numpy as np+ from transformers.testing_utils import require_jinja from transformers.utils import DocstringParsingException, TypeHintParsingException, get_json_schema from transformers.utils.chat_template_utils import render_jinja_template@@ -785,3 +789,98 @@ def test_continue_final_message_on_aliased_reasoning(self):                     continue_final_message=message_field,                 )                 self.assertEqual(rendered[0], expected)+++@require_jinja+class TemplateSandboxTest(unittest.TestCase):+    def render(self, template, chat=None):+        chat = chat if chat is not None else [{"role": "user", "content": "Hi"}]+        rendered, _ = render_jinja_template([chat], chat_template=template)+        return rendered[0]++    def test_dict_global_is_not_a_type(self):+        self.assertEqual(self.render("{{ dict(a=1) }}"), "{'a': 1}")+        self.assertEqual(self.render("{{ dict.fromkeys(['a']) }}"), "{'a': None}")+        # With the real `dict` type, this would build a `types.GenericAlias`+        self.assertEqual(self.render("{{ dict[1] is defined }}"), "False")++    def test_namespace_rejects_underscore_names(self):+        from jinja2.exceptions import SecurityError++        self.assertEqual(+            self.render("{% set ns = namespace(a=1, self=2) %}{% set ns.b = 3 %}{{ ns.a }}{{ ns.self }}{{ ns.b }}"),+            "123",+        )+        for template in (+            "{% set ns = namespace(_a=1) %}",+            "{% set ns = namespace({'__a': 1}) %}",+            "{% set ns = namespace() %}{% set ns.__a = 1 %}",+        ):+            with self.subTest(template=template), self.assertRaises(SecurityError):+                self.render(template)++    def test_rich_objects_become_placeholders(self):+        class Opaque:+            def method(self):+                pass++        chat = [{"role": "user", "content": [{"type": "image", "image": Opaque()}, {"type": "text", "text": "Hi"}]}]+        template = (+            "{{ messages[0].content[0].image }} {{ messages[0].content[0].image.method is defined }} "+            "{{ tools[0].obj }} {{ documents[0].obj }} {{ extra }}"+        )+        rendered, _ = render_jinja_template(+            [chat], chat_template=template, tools=[{"obj": Opaque()}], documents=[{"obj": Opaque()}], extra=Opaque()+        )+        self.assertEqual(rendered[0], "<Opaque> False <Opaque> <Opaque> <Opaque>")+        # The caller's chat keeps the real object, because processors still need it+        self.assertIsInstance(chat[0]["content"][0]["image"], Opaque)++    def test_scalar_subclasses_become_base_types(self):+        class Role(str, Enum):+            USER = "user"++        rendered = self.render(+            "{{ messages[0].role }} {{ messages[0].role.name is defined }}", [{"role": Role.USER, "content": "Hi"}]+        )+        self.assertEqual(rendered, "user False")++    def test_message_objects_render_like_dicts(self):+        @dataclasses.dataclass+        class Message:+            role: str+            content: str++        self.assertEqual(+            self.render("{{ messages[0].role }}: {{ messages[0].content }}", [Message("user", "Hi")]), "user: Hi"+        )++    def test_opaque_messages_raise(self):+        class Message:+            role = "user"+            content = "Hi"++        with self.assertRaises(TypeError):+            self.render("{{ messages[0].content }}", [Message()])++    def test_numpy_scalars_keep_their_value(self):+        template = "{{ 'on' if flag else 'off' }} {{ count + 1 }}"+        rendered, _ = render_jinja_template(+            [[{"role": "user", "content": "Hi"}]], chat_template=template, flag=np.bool_(False), count=np.int64(1)+        )+        # A placeholder string would be truthy and turn the flag on+        self.assertEqual(rendered[0], "off 2")++    def test_continue_final_message_does_not_modify_chat(self):+        chat = [+            {"role": "user", "content": [{"type": "text", "text": "Hi"}]},+            {"role": "assistant", "content": [{"type": "text", "text": "Hello"}]},+        ]+        original_chat = copy.deepcopy(chat)+        rendered, _ = render_jinja_template(+            [chat],+            chat_template="{% for m in messages %}{{ m.content[0].text }}|{% endfor %}",+            continue_final_message=True,+        )+        self.assertEqual(rendered[0], "Hi|Hello")+        self.assertEqual(chat, original_chat)