From 6c25a5c73e83aa32c3241ba84a5cbb3ac0e8a89e Mon Sep 17 00:00:00 2001
From: Silvano Cerza <3314350+silvanocerza@users.noreply.github.com>
Date: Fri, 26 Jul 2024 18:17:25 +0200
Subject: [PATCH] Fix issue that could lead to RCE if using unsecure Jinja
 templates in dynamic prompt builders (#8096)

---
 .../components/builders/dynamic_chat_prompt_builder.py   | 9 +++++----
 haystack/components/builders/dynamic_prompt_builder.py   | 9 +++++----
 releasenotes/notes/fix-jinja-env-81c98225b22dc827.yaml   | 2 ++
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/haystack/components/builders/dynamic_chat_prompt_builder.py b/haystack/components/builders/dynamic_chat_prompt_builder.py
index feabfd98af..942a7f2e2e 100644
--- a/haystack/components/builders/dynamic_chat_prompt_builder.py
+++ b/haystack/components/builders/dynamic_chat_prompt_builder.py
@@ -5,7 +5,8 @@
 import warnings
 from typing import Any, Dict, List, Optional, Set
 
-from jinja2 import Template, meta
+from jinja2 import meta
+from jinja2.sandbox import SandboxedEnvironment
 
 from haystack import component, logging
 from haystack.dataclasses.chat_message import ChatMessage, ChatRole
@@ -177,8 +178,8 @@ def _validate_template(self, template_text: str, provided_variables: Set[str]):
         :raises ValueError:
             If all the required template variables are not provided.
         """
-        template = Template(template_text)
-        ast = template.environment.parse(template_text)
+        env = SandboxedEnvironment()
+        ast = env.parse(template_text)
         required_template_variables = meta.find_undeclared_variables(ast)
         filled_template_vars = required_template_variables.intersection(provided_variables)
         if len(filled_template_vars) != len(required_template_variables):
@@ -187,4 +188,4 @@ def _validate_template(self, template_text: str, provided_variables: Set[str]):
                 f"Required variables: {required_template_variables}. Only the following variables were "
                 f"provided: {provided_variables}. Please provide all the required template variables."
             )
-        return template
+        return env.from_string(template_text)
diff --git a/haystack/components/builders/dynamic_prompt_builder.py b/haystack/components/builders/dynamic_prompt_builder.py
index a01b8d17d2..eb0e4adad9 100644
--- a/haystack/components/builders/dynamic_prompt_builder.py
+++ b/haystack/components/builders/dynamic_prompt_builder.py
@@ -5,7 +5,8 @@
 import warnings
 from typing import Any, Dict, List, Optional, Set
 
-from jinja2 import Template, meta
+from jinja2 import meta
+from jinja2.sandbox import SandboxedEnvironment
 
 from haystack import component, logging
 
@@ -156,8 +157,8 @@ def _validate_template(self, template_text: str, provided_variables: Set[str]):
         :raises ValueError:
             If all the required template variables are not provided.
         """
-        template = Template(template_text)
-        ast = template.environment.parse(template_text)
+        env = SandboxedEnvironment()
+        ast = env.parse(template_text)
         required_template_variables = meta.find_undeclared_variables(ast)
         filled_template_vars = required_template_variables.intersection(provided_variables)
         if len(filled_template_vars) != len(required_template_variables):
@@ -166,4 +167,4 @@ def _validate_template(self, template_text: str, provided_variables: Set[str]):
                 f"Required variables: {required_template_variables}. Only the following variables were "
                 f"provided: {provided_variables}. Please provide all the required template variables."
             )
-        return template
+        return env.from_string(template_text)
diff --git a/releasenotes/notes/fix-jinja-env-81c98225b22dc827.yaml b/releasenotes/notes/fix-jinja-env-81c98225b22dc827.yaml
index 328fae8141..6791f249c4 100644
--- a/releasenotes/notes/fix-jinja-env-81c98225b22dc827.yaml
+++ b/releasenotes/notes/fix-jinja-env-81c98225b22dc827.yaml
@@ -8,6 +8,8 @@ security:
 
     - `PromptBuilder`
     - `ChatPromptBuilder`
+    - `DynamicPromptBuilder`
+    - `DynamicChatPromptBuilder`
     - `OutputAdapter`
     - `ConditionalRouter`