fix: prevent ReDoS in code block extraction

Merge #6118

## Summary
- Replace regular expression-based code block extraction with a simple and safe string-find based search. This avoids exponential backtracking (ReDoS) when processing long or repeating inputs with missing trailing delimiters.
- Add unit tests to verify standard behavior and test against ReDoS vulnerability.

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 933834549

Author: 2 people

src/google/adk/code_executors/code_execution_utils.py

@@ -20,7 +20,6 @@
 import binascii
 import copy
 import dataclasses
-import re
 from typing import List
 from typing import Optional
 
@@ -145,31 +144,41 @@ def extract_code_and_truncate_content(
     first_text_part = copy.deepcopy(text_parts[0])
     response_text = '\n'.join([p.text for p in text_parts])
 
-    # Find the first code block.
-    leading_delimiter_pattern = '|'.join(d[0] for d in code_block_delimiters)
-    trailing_delimiter_pattern = '|'.join(d[1] for d in code_block_delimiters)
-    pattern = re.compile(
-        (
-            rf'(?P<prefix>.*?)({leading_delimiter_pattern})(?P<code>.*?)({trailing_delimiter_pattern})(?P<suffix>.*?)$'
-        ).encode(),
-        re.DOTALL,
-    )
-    pattern_match = pattern.search(response_text.encode())
-    if pattern_match is None:
+    # Find the first code block using simple string search
+    best_start = -1
+    best_end = -1
+    best_lead_len = 0
+
+    for lead, trail in code_block_delimiters:
+      start_idx = response_text.find(lead)
+      if start_idx == -1:
+        continue
+      code_start = start_idx + len(lead)
+      end_idx = response_text.find(trail, code_start)
+      if end_idx == -1:
+        continue
+      # Pick the earliest occurring code block.
+      if best_start == -1 or start_idx < best_start:
+        best_start = start_idx
+        best_end = end_idx
+        best_lead_len = len(lead)
+
+    if best_start == -1:
       return
 
-    code_str = pattern_match.group('code').decode()
+    code_str = response_text[best_start + best_lead_len : best_end]
     if not code_str:
       return
 
     content.parts = []
-    if pattern_match.group('prefix'):
-      first_text_part.text = pattern_match.group('prefix').decode()
+    prefix_text = response_text[:best_start]
+    if prefix_text:
+      first_text_part.text = prefix_text
       content.parts.append(first_text_part)
     content.parts.append(
         CodeExecutionUtils.build_executable_code_part(code_str)
     )
-    return pattern_match.group('code').decode()
+    return code_str
 
   @staticmethod
   def build_executable_code_part(code: str) -> types.Part:

tests/unittests/code_executors/test_code_execution_utils.py

@@ -0,0 +1,179 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import signal
+
+from google.adk.code_executors import code_execution_utils
+from google.genai import types
+
+
+def test_extract_code_and_truncate_content_basic():
+  """Tests basic code extraction and content truncation."""
+  content = types.Content(
+      role="model",
+      parts=[
+          types.Part(
+              text=(
+                  "Here is some code:\n```python\nx = 1\n```\nAnd some text"
+                  " after."
+              )
+          )
+      ],
+  )
+  delimiters = [("```python\n", "\n```")]
+  code = (
+      code_execution_utils.CodeExecutionUtils.extract_code_and_truncate_content(
+          content, delimiters
+      )
+  )
+  assert code == "x = 1"
+  assert len(content.parts) == 2
+  assert content.parts[0].text == "Here is some code:\n"
+  assert content.parts[1].executable_code.code == "x = 1"
+
+
+def test_extract_code_and_truncate_content_multiple_blocks():
+  """Tests that the first code block is extracted when multiple exist."""
+  content = types.Content(
+      role="model",
+      parts=[
+          types.Part(
+              text=(
+                  "First:\n"
+                  "```python\n"
+                  "x = 1\n"
+                  "```\n"
+                  "Second:\n"
+                  "```python\n"
+                  "y = 2\n"
+                  "```"
+              )
+          )
+      ],
+  )
+  delimiters = [("```python\n", "\n```")]
+  code = (
+      code_execution_utils.CodeExecutionUtils.extract_code_and_truncate_content(
+          content, delimiters
+      )
+  )
+  assert code == "x = 1"
+  assert len(content.parts) == 2
+  assert content.parts[0].text == "First:\n"
+  assert content.parts[1].executable_code.code == "x = 1"
+
+
+def test_extract_code_and_truncate_content_no_delimiter():
+  """Tests when no delimiters are found in the content."""
+  content = types.Content(
+      role="model",
+      parts=[types.Part(text="Just plain text without code.")],
+  )
+  delimiters = [("```python\n", "\n```")]
+  code = (
+      code_execution_utils.CodeExecutionUtils.extract_code_and_truncate_content(
+          content, delimiters
+      )
+  )
+  assert code is None
+  # Content should be unmodified.
+  assert len(content.parts) == 1
+  assert content.parts[0].text == "Just plain text without code."
+
+
+def test_extract_code_and_truncate_content_redos_vulnerability():
+  """Tests that a string that would cause ReDoS behaves reasonably."""
+  # Construct a long string that contains repeating patterns without matching delimiters.
+  # The old regex pattern would backtrack exponentially.
+  ticks = "`" * 3
+  long_invalid_payload = ticks + "python\n" + "x = 1\n" * 5000 + "not_matching"
+  content = types.Content(
+      role="model",
+      parts=[types.Part(text=long_invalid_payload)],
+  )
+  delimiters = [(ticks + "python\n", "\n" + ticks)]
+
+  def handler(_signum, _frame):
+    raise TimeoutError("Test timed out (possible ReDoS regression)")
+
+  signal.signal(signal.SIGALRM, handler)
+  signal.alarm(2)
+  try:
+    # If ReDoS vulnerability exists, this call will hang or take a very long time.
+    code = code_execution_utils.CodeExecutionUtils.extract_code_and_truncate_content(
+        content, delimiters
+    )
+  finally:
+    signal.alarm(0)
+  assert code is None
+
+
+def test_extract_code_and_truncate_content_multiple_delimiter_pairs():
+  """Tests code extraction when multiple different delimiter pairs are provided."""
+  ticks = "`" * 3
+  # Case 1: First delimiter pair matches first
+  content = types.Content(
+      role="model",
+      parts=[
+          types.Part(
+              text="Here is tool code:\n"
+              + ticks
+              + "tool_code\nx = 1\n"
+              + ticks
+              + "\nAnd python code:\n"
+              + ticks
+              + "python\ny = 2\n"
+              + ticks
+          )
+      ],
+  )
+  delimiters = [
+      (ticks + "tool_code\n", "\n" + ticks),
+      (ticks + "python\n", "\n" + ticks),
+  ]
+  code = (
+      code_execution_utils.CodeExecutionUtils.extract_code_and_truncate_content(
+          content, delimiters
+      )
+  )
+  assert code == "x = 1"
+  assert len(content.parts) == 2
+  assert content.parts[0].text == "Here is tool code:\n"
+  assert content.parts[1].executable_code.code == "x = 1"
+
+  # Case 2: Second delimiter pair matches first
+  content = types.Content(
+      role="model",
+      parts=[
+          types.Part(
+              text="Here is python code:\n"
+              + ticks
+              + "python\ny = 2\n"
+              + ticks
+              + "\nAnd tool code:\n"
+              + ticks
+              + "tool_code\nx = 1\n"
+              + ticks
+          )
+      ],
+  )
+  code = (
+      code_execution_utils.CodeExecutionUtils.extract_code_and_truncate_content(
+          content, delimiters
+      )
+  )
+  assert code == "y = 2"
+  assert len(content.parts) == 2
+  assert content.parts[0].text == "Here is python code:\n"
+  assert content.parts[1].executable_code.code == "y = 2"