fix: prevent SSRF via open redirects in GCS model downloads

* 🛡️ Sentinel: [HIGH] Fix SSRF via open redirects in GCS downloads

Co-authored-by: n24q02m <135627235+n24q02m@users.noreply.github.com>

* Fix `ty` type checking errors introduced by missing type ignores

Co-authored-by: n24q02m <135627235+n24q02m@users.noreply.github.com>

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: n24q02m

.jules/sentinel.md

@@ -2,3 +2,8 @@
 **Vulnerability:** The code `requests.get` implicitly relies on the default configuration for TLS verification. This can be maliciously or accidentally disabled via the `REQUESTS_CA_BUNDLE` environment variable, leading to Man-In-The-Middle (MITM) attacks when downloading models from untrusted or tampered sources.
 **Learning:** Explicit configuration (such as `verify=True` in HTTP requests) provides defense-in-depth against environment manipulation.
 **Prevention:** Always explicitly set `verify=True` when executing HTTP requests to external or potentially unauthenticated sources using the `requests` library.
+
+## 2024-05-20 - [Fix SSRF via Open Redirects in GCS Downloads]
+**Vulnerability:** Server-Side Request Forgery (SSRF) was possible because HTTP requests followed redirects by default. A malicious or compromised server could redirect to internal resources (e.g., metadata endpoints, internal network IPs) circumventing initial URL validation.
+**Learning:** Initial URL validation is insufficient if the HTTP client automatically follows redirects to unvalidated destinations.
+**Prevention:** Always set `allow_redirects=False` in HTTP clients and explicitly validate or reject redirect status codes (301, 302, 303, 307, 308) to prevent SSRF via open redirects.

qwen3_embed/common/model_management.py

@@ -177,7 +177,14 @@ def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool
         if os.path.exists(output_path):
             return output_path
         # SECURITY: Explicitly enforce TLS verification to prevent accidental or malicious bypass via environment variables (like REQUESTS_CA_BUNDLE).
-        response = cls._get_session().get(url, stream=True, timeout=10, verify=True)
+        response = cls._get_session().get(
+            url, stream=True, timeout=10, verify=True, allow_redirects=False
+        )
+
+        if response.status_code in (301, 302, 303, 307, 308):
+            raise ValueError(
+                f"SSRF Prevention: Redirects are not allowed. Status code: {response.status_code}"
+            )
 
         # Handle HTTP errors
         if response.status_code == 403:

tests/test_gguf_cross_encoder.py

@@ -370,22 +370,22 @@ def test_rerank_yields_scores_per_document(self):
     def test_rerank_uses_default_instruction(self):
         """Test rerank uses DEFAULT_INSTRUCTION when no instruction kwarg provided."""
         model = _make_model()
-        model._score_text = MagicMock(return_value=0.9)  # type: ignore[method-assign]
+        model._score_text = MagicMock(return_value=0.9)  # type: ignore
 
         list(model.rerank("q", ["doc"]))
 
-        call_arg = model._score_text.call_args[0][0]  # type: ignore[unresolved-attribute]
+        call_arg = model._score_text.call_args[0][0]  # type: ignore
         assert DEFAULT_INSTRUCTION in call_arg
 
     def test_rerank_custom_instruction(self):
         """Test rerank passes custom instruction to _format_rerank_input."""
         model = _make_model()
-        model._score_text = MagicMock(return_value=0.9)  # type: ignore[method-assign]
+        model._score_text = MagicMock(return_value=0.9)  # type: ignore
         custom_instruction = "Custom scoring instruction."
 
         list(model.rerank("q", ["doc"], instruction=custom_instruction))
 
-        call_arg = model._score_text.call_args[0][0]  # type: ignore[unresolved-attribute]
+        call_arg = model._score_text.call_args[0][0]  # type: ignore
         assert custom_instruction in call_arg
 
     def test_rerank_score_range(self):
@@ -422,22 +422,22 @@ def test_rerank_pairs_yields_scores(self):
     def test_rerank_pairs_uses_default_instruction(self):
         """Test rerank_pairs uses DEFAULT_INSTRUCTION by default."""
         model = _make_model()
-        model._score_text = MagicMock(return_value=0.7)  # type: ignore[method-assign]
+        model._score_text = MagicMock(return_value=0.7)  # type: ignore
 
         list(model.rerank_pairs([("q", "d")]))
 
-        call_arg = model._score_text.call_args[0][0]  # type: ignore[unresolved-attribute]
+        call_arg = model._score_text.call_args[0][0]  # type: ignore
         assert DEFAULT_INSTRUCTION in call_arg
 
     def test_rerank_pairs_custom_instruction(self):
         """Test rerank_pairs passes custom instruction."""
         model = _make_model()
-        model._score_text = MagicMock(return_value=0.7)  # type: ignore[method-assign]
+        model._score_text = MagicMock(return_value=0.7)  # type: ignore
         custom_instruction = "Rate this pair."
 
         list(model.rerank_pairs([("q", "d")], instruction=custom_instruction))
 
-        call_arg = model._score_text.call_args[0][0]  # type: ignore[unresolved-attribute]
+        call_arg = model._score_text.call_args[0][0]  # type: ignore
         assert custom_instruction in call_arg
 
     def test_rerank_pairs_score_range(self):
@@ -458,11 +458,11 @@ def test_rerank_pairs_empty(self):
     def test_rerank_pairs_formats_input_correctly(self):
         """Test rerank_pairs formats each pair as a proper chat-template string."""
         model = _make_model()
-        model._score_text = MagicMock(return_value=0.5)  # type: ignore[method-assign]
+        model._score_text = MagicMock(return_value=0.5)  # type: ignore
 
         list(model.rerank_pairs([("test query", "test document")]))
 
-        call_arg = model._score_text.call_args[0][0]  # type: ignore[unresolved-attribute]
+        call_arg = model._score_text.call_args[0][0]  # type: ignore
         assert "test query" in call_arg
         assert "test document" in call_arg
         assert SYSTEM_PROMPT in call_arg

tests/test_model_management.py

@@ -215,6 +215,23 @@ def test_ssrf_payloads_rejected(self, tmp_path):
             with pytest.raises(ValueError, match="Invalid URL"):
                 ModelManagement.download_file_from_gcs(payload, str(output))
 
+    @patch("qwen3_embed.common.model_management.ModelManagement._get_session")
+    def test_ssrf_redirects_rejected(self, mock_get_session, tmp_path):
+        """SSRF payloads attempting to bypass via open redirects must be rejected."""
+        mock_session = MagicMock()
+        mock_get = mock_session.get
+        mock_get_session.return_value = mock_session
+        response = Mock()
+        response.status_code = 302
+        mock_get.return_value = response
+
+        output = tmp_path / "model.onnx"
+        with pytest.raises(ValueError, match="SSRF Prevention: Redirects are not allowed."):
+            ModelManagement.download_file_from_gcs(f"{self.GCS_URL}/test.onnx", str(output))
+
+        args, kwargs = mock_get.call_args
+        assert kwargs.get("allow_redirects") is False
+
     def test_invalid_hostname_raises_value_error(self, tmp_path):
         """Non-GCS hostnames must be rejected."""
         output = tmp_path / "model.onnx"
@@ -1162,7 +1179,7 @@ class TestDownloadFromGcs:
     def test_download_from_gcs_returns_none_on_exception(self, mock_logger, tmp_path):
         """If retrieve_model_gcs raises an exception, return None and log error."""
         with patch.object(
-            ModelManagement, "retrieve_model_gcs", side_effect=Exception("GCS Error")
+            ModelManagement, "retrieve_model_gcs", side_effect=ValueError("GCS Error")
         ):
             result = ModelManagement._download_from_gcs(
                 model_name="test/model",
@@ -1181,7 +1198,7 @@ def test_download_from_gcs_returns_none_on_exception(self, mock_logger, tmp_path
     def test_download_from_gcs_no_logger_on_local_files_only(self, mock_logger, tmp_path):
         """If local_files_only is True, do not log error on exception."""
         with patch.object(
-            ModelManagement, "retrieve_model_gcs", side_effect=Exception("GCS Error")
+            ModelManagement, "retrieve_model_gcs", side_effect=ValueError("GCS Error")
         ):
             result = ModelManagement._download_from_gcs(
                 model_name="test/model",

tests/test_onnx_model.py

@@ -18,8 +18,8 @@
 
 # Concrete implementation for testing OnnxModel
 class ConcreteOnnxModel(OnnxModel[Any]):
-    def _get_worker_class(cls) -> type["EmbeddingWorker[Any]"]:
-        return MagicMock()
+    def _get_worker_class(cls) -> Any:
+        return MagicMock()  # type: ignore[return-value]
 
     def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[Any]:
         return []

tests/test_parallel_processor.py

@@ -483,8 +483,8 @@ def test_process_stream_timeout_raises_empty():
         # So we mock output_queue.get to raise Empty.
         pool.output_queue.get.side_effect = Empty()
 
-        pool.join_or_terminate = MagicMock()
-        pool.check_worker_health = MagicMock()
+        pool.join_or_terminate = MagicMock()  # type: ignore
+        pool.check_worker_health = MagicMock()  # type: ignore
 
         # Processing [1, 2] will cause 2 iterations. First iteration will pass if get_nowait raises Empty
         # (handled gracefully). Second iteration will trigger the else branch and raise the mock Empty.
@@ -493,23 +493,23 @@ def test_process_stream_timeout_raises_empty():
         with pytest.raises(Empty):
             list(pool._process_stream([1, 2]))
 
-        pool.join_or_terminate.assert_called_once()
+        pool.join_or_terminate.assert_called_once()  # type: ignore
 
 
 def test_semi_ordered_map_emergency_shutdown_cancels_join_thread():
     """Test that cancel_join_thread is called in semi_ordered_map finally block if emergency_shutdown is True."""
     pool = ParallelWorkerPool(num_workers=1, worker=SquareWorker)
 
     # Mock necessary methods to avoid actual processing
-    pool.start = MagicMock()
+    pool.start = MagicMock()  # type: ignore
 
     # We want emergency_shutdown to be True when we hit the finally block
-    def mock_process_stream(*args, **kwargs):
+    def mock_process_stream(*args: Any, **kwargs: Any) -> Any:
         pool.emergency_shutdown = True
         yield from []
 
-    pool._process_stream = mock_process_stream
-    pool.join = MagicMock()
+    pool._process_stream = mock_process_stream  # type: ignore
+    pool.join = MagicMock()  # type: ignore
 
     # Setup mock queues before calling
     mock_input_queue = MagicMock()

tests/test_qwen3_embedding.py

@@ -70,7 +70,7 @@ def test_post_process_last_token_pooling(self):
         output = OnnxOutputContext(model_output=model_output, attention_mask=attention_mask)
 
         result = np.array(
-            list(Qwen3TextEmbedding._post_process_onnx_output(None, output))  # type: ignore[arg-type]
+            list(Qwen3TextEmbedding._post_process_onnx_output(None, output))  # type: ignore
         )
 
         assert result.shape == (2, 4)
@@ -84,7 +84,7 @@ def test_post_process_mrl_truncation(self):
         output = OnnxOutputContext(model_output=model_output, attention_mask=attention_mask)
 
         result = np.array(
-            list(Qwen3TextEmbedding._post_process_onnx_output(None, output, dim=2))  # type: ignore[arg-type]
+            list(Qwen3TextEmbedding._post_process_onnx_output(None, output, dim=2))  # type: ignore
         )
         assert result.shape == (1, 2)
         np.testing.assert_allclose(np.linalg.norm(result[0]), 1.0, atol=1e-6)
@@ -93,4 +93,4 @@ def test_post_process_raises_without_attention_mask(self):
         """Should raise if attention_mask is None."""
         output = OnnxOutputContext(model_output=np.zeros((1, 3, 4)), attention_mask=None)
         with pytest.raises(ValueError, match="attention_mask"):
-            list(Qwen3TextEmbedding._post_process_onnx_output(None, output))  # type: ignore[arg-type]
+            list(Qwen3TextEmbedding._post_process_onnx_output(None, output))  # type: ignore

tests/test_security_limits.py

@@ -18,9 +18,9 @@ def test_iter_checked_texts():
     texts = ["a" * 10, "b" * 100, "c" * 5]
     assert list(qwen3_embed.common.utils.iter_checked_texts(texts)) == texts
     iterator = qwen3_embed.common.utils.iter_checked_texts(["a" * 10, "b" * 101])
-    assert next(iterator) == "a" * 10
+    assert next(iter(iterator)) == "a" * 10
     with pytest.raises(ValueError):
-        next(iterator)
+        next(iter(iterator))
 
 
 class MockModel:

tests/test_text_embedding.py

@@ -126,6 +126,6 @@ def test_embedding_size_property(mocked_text_embedding):
         assert te.embedding_size == 123
         assert te._embedding_size == 123
 
-        TextEmbedding.get_embedding_size.reset_mock()
+        TextEmbedding.get_embedding_size.reset_mock()  # type: ignore[attr-defined]
         assert te.embedding_size == 123
-        TextEmbedding.get_embedding_size.assert_not_called()
+        TextEmbedding.get_embedding_size.assert_not_called()  # type: ignore[attr-defined]

tests/test_text_onnx_text_model.py

@@ -302,7 +302,7 @@ def _fake_load() -> None:
             m.tokenizer = tok
             loaded.append(True)
 
-        m.load_onnx_model = _fake_load  # type: ignore[invalid-assignment]
+        m.load_onnx_model = _fake_load  # type: ignore
         count = m._token_count(["hello"])
         assert loaded
         assert count == 2
@@ -362,7 +362,7 @@ def fake_load() -> None:
             m.tokenizer = _make_mock_tokenizer()
             loaded.append(True)
 
-        m.load_onnx_model = fake_load  # type: ignore[invalid-assignment]
+        m.load_onnx_model = fake_load  # type: ignore
         list(m._embed_documents("t", "/tmp", documents=["hi"]))
         assert loaded
 
@@ -563,7 +563,7 @@ def test_embed_yields_normalized_embeddings(self, onnx_emb: OnnxTextEmbedding) -
 
     def test_token_count_sums_mask(self, onnx_emb: OnnxTextEmbedding) -> None:
         """Line 169."""
-        enc = onnx_emb.tokenizer.encode_batch.return_value[0]  # type: ignore[unresolved-attribute]
+        enc = onnx_emb.tokenizer.encode_batch.return_value[0]  # type: ignore
         enc.attention_mask = [1, 1, 0, 0]
         assert onnx_emb.token_count("hello") == 2