From a6b9958c9cc4604d09936c8de9608cdaf767bf4e Mon Sep 17 00:00:00 2001 From: bibi samina Date: Wed, 19 Aug 2026 16:03:51 +0530 Subject: [PATCH 1/3] contain gemini batch-job results file within results_folder --- .../google/cloud/operators/gen_ai.py | 25 ++++++++-- .../google/cloud/operators/test_gen_ai.py | 50 +++++++++++++++++++ 2 files changed, 71 insertions(+), 4 deletions(-) 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..f999eb46509e4 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 @@ -57,6 +57,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("/", "-") + path_to_file = os.path.abspath(f"{results_folder}/{file_name}.jsonl") + base = os.path.abspath(f"{results_folder}") + if base != path_to_file and os.path.commonpath([base, path_to_file]) != 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 path_to_file + + class GenAIGenerateEmbeddingsOperator(GoogleCloudBaseOperator): """ Uses the Gemini AI Embeddings API to generate embeddings for words, phrases, sentences, and code. @@ -522,8 +541,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 +990,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( From e0f86e31ed8375496ea45648eb476fe5087e1192 Mon Sep 17 00:00:00 2001 From: bibi samina Date: Thu, 20 Aug 2026 13:21:48 +0530 Subject: [PATCH 2/3] Resolve symlinks in gemini batch-job results containment check --- .../airflow/providers/google/cloud/operators/gen_ai.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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 f999eb46509e4..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 @@ -66,14 +67,14 @@ def _results_file_path(results_folder: str | None, job) -> str: is not, so a ``..`` in it would place the file outside ``results_folder`` (CWE-22). """ file_name = job.display_name or job.name.replace("/", "-") - path_to_file = os.path.abspath(f"{results_folder}/{file_name}.jsonl") - base = os.path.abspath(f"{results_folder}") - if base != path_to_file and os.path.commonpath([base, path_to_file]) != base: + 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 path_to_file + return str(path_to_file) class GenAIGenerateEmbeddingsOperator(GoogleCloudBaseOperator): From 6d19f3f6038079cd156f7f1dbd0b610098aba58a Mon Sep 17 00:00:00 2001 From: bibi samina Date: Fri, 21 Aug 2026 12:27:45 +0530 Subject: [PATCH 3/3] Add symlink-escape regression test for gemini batch-job results path --- .../google/cloud/operators/test_gen_ai.py | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) 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 a63bfb2d76f9b..54668677a3f78 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 @@ -507,6 +507,36 @@ def test_prepare_results_for_xcom_rejects_display_name_path_traversal(self, mock assert not (tmp_path / "evil.jsonl").exists() + @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook")) + def test_prepare_results_for_xcom_rejects_display_name_symlink_escape(self, mock_hook, tmp_path): + results_folder = tmp_path / "results" + results_folder.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + (results_folder / "link").symlink_to(outside, target_is_directory=True) + + 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(results_folder), + ) + 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 = "link/evil" + + with pytest.raises(ValueError, match="Refusing to write batch-job results outside"): + op._prepare_results_for_xcom(mock_job) + + assert not (outside / "evil.jsonl").exists() + @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook")) def test__wait_until_complete_exception_raises_airflow_exception(self, mock_hook): op = GenAIGeminiCreateBatchJobOperator( @@ -932,6 +962,36 @@ def test_prepare_results_for_xcom_rejects_display_name_path_traversal(self, mock assert not (tmp_path / "evil.jsonl").exists() + @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook")) + def test_prepare_results_for_xcom_rejects_display_name_symlink_escape(self, mock_hook, tmp_path): + results_folder = tmp_path / "results" + results_folder.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + (results_folder / "link").symlink_to(outside, target_is_directory=True) + + 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(results_folder), + ) + 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 = "link/evil" + + with pytest.raises(ValueError, match="Refusing to write batch-job results outside"): + op._prepare_results_for_xcom(mock_job) + + assert not (outside / "evil.jsonl").exists() + @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook")) def test__wait_until_complete_exception_raises_airflow_exception(self, mock_hook): op = GenAIGeminiCreateEmbeddingsBatchJobOperator(