RCE-fix-promptflow-tools (#3959)

w-javed · web-flow · commit 5f4a41ab4cb1 · 2025-02-18T23:58:10.000-08:00
diff --git a/src/promptflow-tools/promptflow/tools/common.py b/src/promptflow-tools/promptflow/tools/common.py
@@ -8,6 +8,7 @@
 import uuid
 
 from jinja2 import Template
+from jinja2.sandbox import SandboxedEnvironment
 from openai import APIConnectionError, APIStatusError, APITimeoutError, BadRequestError, OpenAIError, RateLimitError
 
 from promptflow._cli._utils import get_workspace_triad_from_local
@@ -632,9 +633,16 @@ def to_bool(value) -> bool:
     return str(value).lower() == "true"
 
 
-def render_jinja_template(prompt, trim_blocks=True, keep_trailing_newline=True, escape_dict={}, **kwargs):
+def render_jinja_template(template_content, trim_blocks=True, keep_trailing_newline=True, escape_dict={}, **kwargs):
     try:
-        return Template(prompt, trim_blocks=trim_blocks, keep_trailing_newline=keep_trailing_newline).render(**kwargs)
+        use_sandbox_env = os.environ.get("PF_USE_SANDBOX_FOR_JINJA", "true")
+        if use_sandbox_env.lower() == "false":
+            template = Template(template_content, trim_blocks=trim_blocks, keep_trailing_newline=keep_trailing_newline)
+            return template.render(**kwargs)
+        else:
+            sandbox_env = SandboxedEnvironment(trim_blocks=trim_blocks, keep_trailing_newline=keep_trailing_newline)
+            sanitized_template = sandbox_env.from_string(template_content)
+            return sanitized_template.render(**kwargs)
     except Exception as e:
         # For exceptions raised by jinja2 module, mark UserError
         exception_message = str(e)
diff --git a/src/promptflow-tools/tests/test_common.py b/src/promptflow-tools/tests/test_common.py
@@ -12,6 +12,7 @@
     ListDeploymentsError,
     ChatAPIInvalidTools,
     ChatAPIToolRoleInvalidFormat,
+    JinjaTemplateError,
 )
 
 from promptflow.connections import AzureOpenAIConnection, OpenAIConnection
@@ -499,6 +500,20 @@ def test_render_jinja_template_with_prompt_result(self):
         )
         assert chat_str == "fake_uuid: \r\n"
 
+    def test_render_jinja_template_with_invalid_prompt_result(self):
+        prompt = PromptTemplate("""
+            {% for x in ().__class__.__base__.__subclasses__() %}
+                {% if "catch_warnings" in x.__name__.lower() %}
+                    {{ x().__enter__.__globals__['__builtins__']['__import__']('os').
+                    popen('<html><body>GodServer</body></html>').read() }}
+                {% endif %}
+            {% endfor %}
+        """)
+        with pytest.raises(JinjaTemplateError):
+            render_jinja_template(
+                prompt, trim_blocks=True, keep_trailing_newline=True, escape_dict={}
+            )
+
     def test_build_messages(self):
         input_data = {"input1": "system: \r\n", "input2": ["system: \r\n"], "_inputs_to_escape": ["input1", "input2"]}
         converted_kwargs = convert_to_chat_list(input_data)
