mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
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:
parent
b9939e8e6d
commit
d0061fbea7
2 changed files with 77 additions and 26 deletions
|
|
@ -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)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue