diff --git a/google/cloud/aiplatform/utils/_ipython_utils.py b/google/cloud/aiplatform/utils/_ipython_utils.py
index afa6f3f3d7..a4319f6cca 100644
--- a/google/cloud/aiplatform/utils/_ipython_utils.py
+++ b/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()}
-
- {icon}
- {text}
+
+ {safe_icon}
+ {safe_text}
"""
# Add the click handler for the link
- html += f"""
+ html_out += f"""
@@ -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:
diff --git a/tests/unit/aiplatform/test_utils.py b/tests/unit/aiplatform/test_utils.py
index e40aae37b8..0116fd1815 100644
--- a/tests/unit/aiplatform/test_utils.py
+++ b/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="",
+ url="https://console.cloud.google.com/test?param=1&another=2",
+ icon="",
+ )
+ mock_display_module.HTML.assert_called_once()
+ html_arg = mock_display_module.HTML.call_args[0][0]
+
+ assert "<script>alert('xss')</script>" in html_arg
+ assert "href=\"https://console.cloud.google.com/test?param=1&another=2\"" in html_arg
+ assert "window.google.colab.openUrl(\"https://console.cloud.google.com/test?param=1&another=2\")" in html_arg
+ assert "<icon>" in html_arg
+
+ mock_display_module.display.assert_called_once_with(mock_display_module.HTML.return_value)
diff --git a/tests/unit/architecture/test_vertexai_import.py b/tests/unit/architecture/test_vertexai_import.py
index 2170c7fa9e..227ede0eea 100644
--- a/tests/unit/architecture/test_vertexai_import.py
+++ b/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