contain gemini batch-job results file within results_folder

providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py

@@ -22,6 +22,7 @@
 import os.path
 import time
 from collections.abc import Sequence
+from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 from google.genai.errors import ClientError
@@ -57,6 +58,25 @@
     from airflow.providers.common.compat.sdk import Context
 
 
+def _results_file_path(results_folder: str | None, job) -> str:
+    """
+    Build the local results-file path for a batch job, refusing one that escapes ``results_folder``.
+
+    ``display_name`` and ``name`` come from the batch-job metadata returned by the Gemini
+    API. ``job.name`` is already sanitized against ``/``, but ``display_name`` (used first)
+    is not, so a ``..`` in it would place the file outside ``results_folder`` (CWE-22).
+    """
+    file_name = job.display_name or job.name.replace("/", "-")
+    base = Path(f"{results_folder}").resolve()
+    path_to_file = (base / f"{file_name}.jsonl").resolve()
+    if not path_to_file.is_relative_to(base):
+        raise ValueError(
+            f"Refusing to write batch-job results outside results_folder {results_folder!r}: "
+            f"file name {file_name!r} resolves to {path_to_file!r}."
+        )
+    return str(path_to_file)
+
+
 class GenAIGenerateEmbeddingsOperator(GoogleCloudBaseOperator):
     """
     Uses the Gemini AI Embeddings API to generate embeddings for words, phrases, sentences, and code.
@@ -522,8 +542,7 @@ def _prepare_results_for_xcom(self, job):
             self._validate_results_folder()
             file_content_bytes = self.hook.download_file(file_name=job.dest.file_name)
             file_content = file_content_bytes.decode("utf-8")
-            file_name = job.display_name or job.name.replace("/", "-")
-            path_to_file = os.path.abspath(f"{self.results_folder}/{file_name}.jsonl")
+            path_to_file = _results_file_path(self.results_folder, job)
             with open(path_to_file, "w") as file_with_results:
                 file_with_results.writelines(file_content.splitlines(True))
             results = path_to_file
@@ -972,8 +991,7 @@ def _prepare_results_for_xcom(self, job):
             self._validate_results_folder()
             file_content_bytes = self.hook.download_file(file_name=job.dest.file_name)
             file_content = file_content_bytes.decode("utf-8")
-            file_name = job.display_name or job.name.replace("/", "-")
-            path_to_file = os.path.abspath(f"{self.results_folder}/{file_name}.jsonl")
+            path_to_file = _results_file_path(self.results_folder, job)
             with open(path_to_file, "w") as file_with_results:
                 file_with_results.writelines(file_content.splitlines(True))
             results = path_to_file

providers/google/tests/unit/google/cloud/operators/test_gen_ai.py

@@ -482,6 +482,31 @@ def test_prepare_results_for_xcom_results_folder_not_exists_raises_airflow_excep
 
         mock_hook.return_value.download_file.assert_not_called()
 
+    @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+    def test_prepare_results_for_xcom_rejects_display_name_path_traversal(self, mock_hook, tmp_path):
+        op = GenAIGeminiCreateBatchJobOperator(
+            task_id=TASK_ID,
+            project_id=GCP_PROJECT,
+            location=GCP_LOCATION,
+            model=TEST_GEMINI_MODEL,
+            gcp_conn_id=GCP_CONN_ID,
+            impersonation_chain=IMPERSONATION_CHAIN,
+            input_source=TEST_FILE_NAME,
+            gemini_api_key=TEST_GEMINI_API_KEY,
+            results_folder=str(tmp_path / "results"),
+        )
+        (tmp_path / "results").mkdir()
+        mock_hook.return_value.download_file.return_value = b"data"
+        mock_job = mock.MagicMock()
+        mock_job.dest.inlined_responses = None
+        mock_job.dest.file_name = "results-file"
+        mock_job.display_name = "../evil"
+
+        with pytest.raises(ValueError, match="Refusing to write batch-job results outside"):
+            op._prepare_results_for_xcom(mock_job)
+
+        assert not (tmp_path / "evil.jsonl").exists()
+
     @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
     def test__wait_until_complete_exception_raises_airflow_exception(self, mock_hook):
         op = GenAIGeminiCreateBatchJobOperator(
@@ -882,6 +907,31 @@ def test_prepare_results_for_xcom_results_folder_not_exists_raises_airflow_excep
 
         mock_hook.return_value.download_file.assert_not_called()
 
+    @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+    def test_prepare_results_for_xcom_rejects_display_name_path_traversal(self, mock_hook, tmp_path):
+        op = GenAIGeminiCreateEmbeddingsBatchJobOperator(
+            task_id=TASK_ID,
+            project_id=GCP_PROJECT,
+            location=GCP_LOCATION,
+            input_source=TEST_FILE_NAME,
+            model=EMBEDDING_MODEL,
+            gemini_api_key=TEST_GEMINI_API_KEY,
+            gcp_conn_id=GCP_CONN_ID,
+            impersonation_chain=IMPERSONATION_CHAIN,
+            results_folder=str(tmp_path / "results"),
+        )
+        (tmp_path / "results").mkdir()
+        mock_hook.return_value.download_file.return_value = b"data"
+        mock_job = mock.MagicMock()
+        mock_job.dest.inlined_embed_content_responses = None
+        mock_job.dest.file_name = "results-file"
+        mock_job.display_name = "../evil"
+
+        with pytest.raises(ValueError, match="Refusing to write batch-job results outside"):
+            op._prepare_results_for_xcom(mock_job)
+
+        assert not (tmp_path / "evil.jsonl").exists()
+
     @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
     def test__wait_until_complete_exception_raises_airflow_exception(self, mock_hook):
         op = GenAIGeminiCreateEmbeddingsBatchJobOperator(