diff --git a/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py b/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py index 65d8a9433742c..f5471ffa1b168 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py +++ b/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 diff --git a/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py b/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py index 1071f2de82920..a63bfb2d76f9b 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py +++ b/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(