Fix XPath injection in search_item_ctrl_f (#1768)

Author: albertvillanova

src/smolagents/vision_web_browser.py

@@ -84,6 +84,24 @@ def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> None:
     return
 
 
+def _escape_xpath_string(s: str) -> str:
+    """
+    Escapes a string for safe use in an XPath expression.
+
+    Args:
+        s (`str`): Arbitrary input string to escape.
+
+    Returns:
+        `str`: Valid XPath expression representing the literal value of `s`.
+    """
+    if "'" not in s:
+        return f"'{s}'"
+    if '"' not in s:
+        return f'"{s}"'
+    parts = s.split("'")
+    return "concat(" + ', "\'", '.join(f"'{p}'" for p in parts) + ")"
+
+
 @tool
 def search_item_ctrl_f(text: str, nth_result: int = 1) -> str:
     """
@@ -92,7 +110,8 @@ def search_item_ctrl_f(text: str, nth_result: int = 1) -> str:
         text: The text to search for
         nth_result: Which occurrence to jump to (default: 1)
     """
-    elements = driver.find_elements(By.XPATH, f"//*[contains(text(), '{text}')]")
+    escaped_text = _escape_xpath_string(text)
+    elements = driver.find_elements(By.XPATH, f"//*[contains(text(), {escaped_text})]")
     if nth_result > len(elements):
         raise Exception(f"Match n°{nth_result} not found (only {len(elements)} matches found)")
     result = f"Found {len(elements)} matches for '{text}'."

tests/test_vision_web_browser.py

@@ -0,0 +1,132 @@
+"""Test XPath injection vulnerability fix in vision_web_browser.py"""
+
+from unittest.mock import Mock, patch
+
+import pytest
+
+from smolagents.vision_web_browser import _escape_xpath_string, search_item_ctrl_f
+
+
+@pytest.fixture
+def mock_driver():
+    """Mock Selenium WebDriver"""
+    driver = Mock()
+    driver.find_elements.return_value = [Mock()]  # Mock found elements
+    driver.execute_script.return_value = None
+    return driver
+
+
+class TestXPathEscaping:
+    """Test XPath string escaping functionality"""
+
+    @pytest.mark.parametrize(
+        "input_text,expected_pattern",
+        [
+            ("normal text", "'normal text'"),
+            ("text with 'quote'", "\"text with 'quote'\""),
+            ('text with "quote"', "'text with \"quote\"'"),
+            ("text with one single'quote", '"text with one single\'quote"'),
+            ('text with one double"quote', "'text with one double\"quote'"),
+            (
+                "text with both 'single' and \"double\" quotes",
+                "concat('text with both ', \"'\", 'single', \"'\", ' and \"double\" quotes')",
+            ),
+            ("", "''"),
+            ("'", '"\'"'),
+            ('"', "'\"'"),
+        ],
+    )
+    def test_escape_xpath_string_basic(self, input_text, expected_pattern):
+        """Test basic XPath escaping cases"""
+        result = _escape_xpath_string(input_text)
+        assert result == expected_pattern
+
+    @pytest.mark.parametrize(
+        "input_text",
+        [
+            "text with both 'single' and \"double\" quotes",
+            'it\'s a "test" case',
+            "'mixed\" quotes'",
+        ],
+    )
+    def test_escape_xpath_string_mixed_quotes(self, input_text):
+        """Test XPath escaping with mixed quotes uses concat()"""
+        result = _escape_xpath_string(input_text)
+        assert result.startswith("concat(")
+        assert result.endswith(")")
+
+    @pytest.mark.parametrize(
+        "malicious_input",
+        [
+            "')] | //script[@src='evil.js'] | foo[contains(text(), '",
+            "') or 1=1 or ('",
+            "')] | //user[contains(@role,'admin')] | foo[contains(text(), '",
+            "') and substring(//user[1]/password,1,1)='a",
+        ],
+    )
+    def test_escape_prevents_injection(self, malicious_input):
+        """Test that malicious XPath injection attempts are safely escaped"""
+        result = _escape_xpath_string(malicious_input)
+        # Should either be wrapped in quotes or use concat()
+        assert (
+            (result.startswith("'") and result.endswith("'"))
+            or (result.startswith('"') and result.endswith('"'))
+            or result.startswith("concat(")
+        )
+
+
+class TestSearchItemCtrlF:
+    """Test the search_item_ctrl_f function with XPath injection protection"""
+
+    @pytest.mark.parametrize(
+        "search_text",
+        [
+            "normal search",
+            "search with 'quotes'",
+            'search with "quotes"',
+            "')] | //script[@src='evil.js'] | foo[contains(text(), '",
+            "') or 1=1 or ('",
+        ],
+    )
+    def test_search_item_prevents_injection(self, search_text, mock_driver):
+        """Test that search_item_ctrl_f prevents XPath injection"""
+        with patch("smolagents.vision_web_browser.driver", mock_driver, create=True):
+            # Call the function
+            result = search_item_ctrl_f(search_text)
+
+            # Verify driver.find_elements was called
+            mock_driver.find_elements.assert_called_once()
+
+            # Get the actual XPath query that was generated
+            call_args = mock_driver.find_elements.call_args
+            xpath_query = call_args[0][1]  # Second positional argument
+
+            # Verify the query doesn't contain unescaped injection
+            if "')] | //" in search_text:
+                # For injection attempts, verify they're properly escaped
+                # The query should either use concat() or be properly quoted
+                is_concat = "concat(" in xpath_query
+                is_properly_quoted = xpath_query.count('"') >= 2 or xpath_query.count("'") >= 2
+                assert is_concat or is_properly_quoted, f"XPath injection not prevented: {xpath_query}"
+
+            # Verify we got a result
+            assert "Found" in result
+
+    def test_search_item_nth_result(self, mock_driver):
+        """Test nth_result parameter works correctly"""
+        mock_driver.find_elements.return_value = [Mock(), Mock(), Mock()]  # 3 elements
+
+        with patch("smolagents.vision_web_browser.driver", mock_driver, create=True):
+            result = search_item_ctrl_f("test", nth_result=2)
+
+            # Should find 3 matches and focus on element 2
+            assert "Found 3 matches" in result
+            assert "Focused on element 2 of 3" in result
+
+    def test_search_item_not_found(self, mock_driver):
+        """Test exception when nth_result exceeds available matches"""
+        mock_driver.find_elements.return_value = [Mock()]  # Only 1 element
+
+        with patch("smolagents.vision_web_browser.driver", mock_driver, create=True):
+            with pytest.raises(Exception, match="Match n°3 not found"):
+                search_item_ctrl_f("test", nth_result=3)