fix(vertex_ai/gemini): address mypy + gcs credential review feedback

- mypy: cast image_url/file dict to Dict[str, Any] before reading
  mime_type/content_type, since those keys are not declared on the
  ChatCompletionImageUrlObject / ChatCompletionFileObjectFile TypedDicts
  and mypy was inferring `object` for the fallback-or chain, failing
  the lint job.
- security (veria-ai review): in _get_gcs_object_content_type, only
  call VertexBase.get_access_token when the caller explicitly passes
  vertex_project/vertex_credentials. Without explicit credentials we
  now skip the authenticated call entirely and only attempt an
  anonymous metadata lookup, so a Gemini API-key / Google AI Studio
  caller can no longer trick the proxy into querying GCS with the
  server's default Google credentials.
- Add regression test covering the no-credentials path (no ADC token
  requested, no Authorization header on the outbound metadata call).
This commit is contained in:
Kris Xia 2026-05-11 13:56:16 +08:00 • committed by S0ngRu1
parent b9939e8e6d
commit d0061fbea7
2 changed files with 77 additions and 26 deletions

View file

@ -230,6 +230,11 @@ def _get_gcs_object_content_type(
) -> Optional[str]:
"""
Resolve content type from GCS object metadata.
仅当调用方显式传入 Vertex 凭据时才附带 Bearer token,避免在 Gemini
API key(Google AI Studio)路径上自动使用服务端的默认 Google 凭据去
访问 GCS,防止被用作探测私有 GCS 对象的 oracle。
无显式凭据时只做匿名请求,仅对公开可读对象有效。
"""
try:
bucket, object_name = _parse_gs_uri(image_url)
@ -242,14 +247,14 @@ def _get_gcs_object_content_type(
explicit_vertex_auth_provided = (
vertex_project is not None or vertex_credentials is not None
)
try:
access_token, _ = _get_vertex_base().get_access_token(
credentials=vertex_credentials,
project_id=vertex_project,
)
headers["Authorization"] = f"Bearer {access_token}"
except Exception as e:
if explicit_vertex_auth_provided:
if explicit_vertex_auth_provided:
try:
access_token, _ = _get_vertex_base().get_access_token(
credentials=vertex_credentials,
project_id=vertex_project,
)
headers["Authorization"] = f"Bearer {access_token}"
except Exception as e:
raise litellm.BadRequestError(
message=(
"Unable to fetch GCS metadata with provided Vertex credentials/project. "
@ -258,7 +263,6 @@ def _get_gcs_object_content_type(
model=None,
llm_provider="vertex_ai",
)
# 未显式提供 Vertex 凭据时,metadata 对公开对象仍可能可读,继续无 token 尝试。
# 通过 httpx.URL 固定 scheme/host,并对 bucket、object 都做 URL 编码,
# 避免被 CodeQL 误判为可能拼接出任意主机 URL 的 SSRF。
@ -551,12 +555,13 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
model=model,
llm_provider="vertex_ai",
)
image_url_dict = cast(Dict[str, Any], raw_image_url)
format = (
raw_image_url.get("format")
or raw_image_url.get("mime_type")
or raw_image_url.get("content_type")
image_url_dict.get("format")
or image_url_dict.get("mime_type")
or image_url_dict.get("content_type")
)
detail = raw_image_url.get("detail")
detail = image_url_dict.get("detail")
media_resolution_enum = (
_convert_detail_to_media_resolution_enum(detail)
)

View file

@ -1389,9 +1389,14 @@ def test_get_gcs_object_content_type_uses_shared_vertex_base_instance():
mock_http_response.json.return_value = {"contentType": "image/png"}
mock_http_response.raise_for_status.return_value = None
with patch.object(
gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base
), patch("litellm.llms.vertex_ai.gemini.transformation.httpx.get") as mock_http_get:
with (
patch.object(
gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base
),
patch(
"litellm.llms.vertex_ai.gemini.transformation.httpx.get"
) as mock_http_get,
):
mock_http_get.return_value = mock_http_response
content_type = gemini_transformation._get_gcs_object_content_type(
image_url="gs://my-bucket/path/to/image-without-extension",
@ -1428,7 +1433,9 @@ def test_process_gemini_media_rejects_unsupported_metadata_mime_type():
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
return_value="application/octet-stream",
):
with pytest.raises(litellm.BadRequestError, match="File type not supported by gemini"):
with pytest.raises(
litellm.BadRequestError, match="File type not supported by gemini"
):
_process_gemini_media("gs://bucket/image-without-extension")
@ -1437,7 +1444,9 @@ def test_get_gcs_object_content_type_fails_fast_with_explicit_credentials():
mock_vertex_base = MagicMock()
mock_vertex_base.get_access_token.side_effect = Exception("token failure")
with patch.object(gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base):
with patch.object(
gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base
):
with pytest.raises(
litellm.BadRequestError,
match="Unable to fetch GCS metadata with provided Vertex credentials/project",
@ -1449,6 +1458,41 @@ def test_get_gcs_object_content_type_fails_fast_with_explicit_credentials():
)
def test_get_gcs_object_content_type_without_credentials_skips_auth():
"""未显式提供 Vertex 凭据时,不应使用服务端默认凭据去访问 GCS。
这是为了避免 Gemini API key(Google AI Studio)路径被用作探测私有
GCS 对象的 oracle(veria-ai 审阅反馈)。
"""
from litellm.llms.vertex_ai.gemini import transformation as gemini_transformation
mock_vertex_base = MagicMock()
mock_http_response = MagicMock()
mock_http_response.json.return_value = {"contentType": "image/jpeg"}
mock_http_response.raise_for_status.return_value = None
with (
patch.object(
gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base
),
patch(
"litellm.llms.vertex_ai.gemini.transformation.httpx.get"
) as mock_http_get,
):
mock_http_get.return_value = mock_http_response
content_type = gemini_transformation._get_gcs_object_content_type(
image_url="gs://public-bucket/public-object"
)
# 不应触发 get_access_token(避免使用服务端默认凭据)
mock_vertex_base.get_access_token.assert_not_called()
# 匿名请求仍会发出,用于公开可读对象
mock_http_get.assert_called_once()
call_kwargs = mock_http_get.call_args.kwargs
assert "Authorization" not in call_kwargs.get("headers", {})
assert content_type == "image/jpeg"
def test_async_transform_request_body_does_not_block_event_loop():
"""当同步 GCS metadata 查询阻塞时,async_transform_request_body 不应阻塞事件循环。"""
import asyncio
@ -1513,14 +1557,16 @@ def test_async_transform_request_body_does_not_block_event_loop():
await transform_task
return sleep_elapsed
with patch.object(
gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base
), patch.object(
gemini_transformation.httpx, "get", side_effect=slow_http_get
), patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching."
"ContextCachingEndpoints.async_check_and_create_cache",
new=fake_check_and_create_cache,
with (
patch.object(
gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base
),
patch.object(gemini_transformation.httpx, "get", side_effect=slow_http_get),
patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching."
"ContextCachingEndpoints.async_check_and_create_cache",
new=fake_check_and_create_cache,
),
):
sleep_elapsed = asyncio.run(run_scenario())