Merge branch 'main' into owl-bot-copy

Author: matthew29tang

google/cloud/aiplatform/utils/_ipython_utils.py

@@ -15,6 +15,8 @@
 # limitations under the License.
 #
 
+import html
+import json
 import sys
 import typing
 import urllib
@@ -119,28 +121,31 @@ def display_link(text: str, url: str, icon: Optional[str] = "open_in_new") -> No
 
     button_id = f"view-vertex-resource-{str(uuid4())}"
 
+    # Safe encodings
+    safe_href = html.escape(url, quote=True)
+    safe_text = html.escape(text)
+    safe_icon = html.escape(icon) if icon else ""
+    safe_url_js = json.dumps(url)
+
     # Add the markup for the CSS and link component
-    html = f"""
+    html_out = f"""
         {_get_styles()}
-        <a class="view-vertex-resource" id="{button_id}" href="#view-{button_id}">
-          <span class="material-icons view-vertex-icon">{icon}</span>
-          <span>{text}</span>
+        <a class="view-vertex-resource" id="{button_id}" href="{safe_href}" target="_blank">
+          <span class="material-icons view-vertex-icon">{safe_icon}</span>
+          <span>{safe_text}</span>
         </a>
         """
 
     # Add the click handler for the link
-    html += f"""
+    html_out += f"""
         <script>
           (function () {{
             const link = document.getElementById('{button_id}');
             link.addEventListener('click', (e) => {{
               if (window.google?.colab?.openUrl) {{
-                window.google.colab.openUrl('{url}');
-              }} else {{
-                window.open('{url}', '_blank');
+                window.google.colab.openUrl({safe_url_js});
+                e.preventDefault();
               }}
-              e.stopPropagation();
-              e.preventDefault();
             }});
           }})();
         </script>
@@ -149,7 +154,7 @@ def display_link(text: str, url: str, icon: Optional[str] = "open_in_new") -> No
     from IPython.display import display
     from IPython.display import HTML
 
-    display(HTML(html))
+    display(HTML(html_out))
 
 
 def display_experiment_button(experiment: "experiment_resources.Experiment") -> None:

tests/unit/aiplatform/test_utils.py

@@ -38,6 +38,7 @@
 from google.cloud.aiplatform.compat.types import pipeline_failure_policy
 from google.cloud.aiplatform import datasets
 from google.cloud.aiplatform.utils import (
+    _ipython_utils,
     column_transformations_utils,
     gcs_utils,
     pipeline_utils,
@@ -1137,3 +1138,29 @@ def test_load_yaml_from_invalid_uri(self, uri: str):
         )
         with pytest.raises(ValueError, match=message):
             yaml_utils.load_yaml(uri)
+
+
+class TestIpythonUtils:
+    """Tests for IPython utility functions."""
+
+    def test_display_link_raises_value_error_for_invalid_url(self):
+        with pytest.raises(ValueError, match="Only urls starting with"):
+            _ipython_utils.display_link("bad", "https://example.com")
+
+    def test_display_link_success_and_sanitizes(self):
+        mock_display_module = mock.MagicMock()
+        with mock.patch.dict("sys.modules", {"IPython.display": mock_display_module, "IPython": mock.MagicMock()}):
+            _ipython_utils.display_link(
+                text="<script>alert('xss')</script>",
+                url="https://console.cloud.google.com/test?param=1&another=2",
+                icon="<icon>",
+            )
+            mock_display_module.HTML.assert_called_once()
+            html_arg = mock_display_module.HTML.call_args[0][0]
+
+            assert "&lt;script&gt;alert(&#x27;xss&#x27;)&lt;/script&gt;" in html_arg
+            assert "href=\"https://console.cloud.google.com/test?param=1&amp;another=2\"" in html_arg
+            assert "window.google.colab.openUrl(\"https://console.cloud.google.com/test?param=1&another=2\")" in html_arg
+            assert "&lt;icon&gt;" in html_arg
+
+            mock_display_module.display.assert_called_once_with(mock_display_module.HTML.return_value)

tests/unit/architecture/test_vertexai_import.py

@@ -33,6 +33,7 @@ def test_vertexai_import():
     import google.api_core.operations_v1 as _  # noqa: F811
     import google.api_core.rest_streaming as _  # noqa: F811
     import google.cloud.storage as _  # noqa: F811
+    import html as _  # noqa: F811
 
     try:
         # Needed for Python 3.8