diff --git a/litellm/files/main.py b/litellm/files/main.py index ceccba8d8006..669d50dde414 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -10,6 +10,7 @@ import time import uuid as uuid_module from functools import partial +from types import MappingProxyType from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast import httpx @@ -85,6 +86,16 @@ def _should_sdk_support_streaming( ################################################# +def _add_trusted_model_credentials_to_litellm_params( + litellm_params_dict: Dict[str, Any], kwargs: Dict[str, Any] +) -> None: + trusted_model_credentials = kwargs.get("_litellm_internal_model_credentials") + if isinstance(trusted_model_credentials, type(MappingProxyType({}))): + litellm_params_dict["_litellm_internal_model_credentials"] = ( + trusted_model_credentials + ) + + @client async def acreate_file( file: FileTypes, @@ -373,6 +384,10 @@ def file_retrieve( ) if provider_config is not None: litellm_params_dict = get_litellm_params(**kwargs) + _add_trusted_model_credentials_to_litellm_params( + litellm_params_dict=litellm_params_dict, + kwargs=kwargs, + ) litellm_params_dict["api_key"] = optional_params.api_key litellm_params_dict["api_base"] = optional_params.api_base @@ -497,6 +512,10 @@ def file_delete( pass optional_params = GenericLiteLLMParams(**kwargs) litellm_params_dict = get_litellm_params(**kwargs) + _add_trusted_model_credentials_to_litellm_params( + litellm_params_dict=litellm_params_dict, + kwargs=kwargs, + ) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 # set timeout for 10 minutes by default @@ -846,6 +865,10 @@ def file_content( try: optional_params = GenericLiteLLMParams(**kwargs) litellm_params_dict = get_litellm_params(**kwargs) + _add_trusted_model_credentials_to_litellm_params( + litellm_params_dict=litellm_params_dict, + kwargs=kwargs, + ) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 client = kwargs.get("client") @@ -993,6 +1016,7 @@ def file_content( vertex_location=vertex_ai_location, timeout=timeout, max_retries=optional_params.max_retries, + litellm_params=litellm_params_dict, ) elif custom_llm_provider == "bedrock": response = bedrock_files_instance.file_content( diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 65296bafcf3e..90057984235f 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -6,12 +6,14 @@ from litellm._uuid import uuid from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple -from urllib.parse import quote from litellm._logging import verbose_logger from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE from litellm.integrations.additional_logging_utils import AdditionalLoggingUtils from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase +from litellm.litellm_core_utils.cloud_storage_security import ( + sanitize_cloud_object_component, +) from litellm.proxy._types import CommonProxyErrors from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.gcs_bucket import * @@ -335,7 +337,11 @@ def _get_object_name( _litellm_params = kwargs.get("litellm_params", None) or {} _metadata = _litellm_params.get("metadata", None) or {} if "gcs_log_id" in _metadata: - object_name = _metadata["gcs_log_id"] + safe_log_id = sanitize_cloud_object_component( + _metadata.get("gcs_log_id"), fallback="" + ) + if safe_log_id: + object_name = f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}" return object_name @@ -367,8 +373,7 @@ async def get_request_response_payload( request_date_str=date_str, response_id=request_id, ) - encoded_object_name = quote(object_name, safe="") - response = await self.download_gcs_object(encoded_object_name) + response = await self.download_gcs_object(object_name) if response is not None: loaded_response = json.loads(response) diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_base.py b/litellm/integrations/gcs_bucket/gcs_bucket_base.py index e84b37e689b6..1c5e30777a2d 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_base.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_base.py @@ -11,6 +11,10 @@ from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.litellm_core_utils.cloud_storage_security import ( + encode_gcs_object_name_for_url, + split_configured_cloud_bucket_name, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -133,8 +137,8 @@ def _handle_folders_in_bucket_name( - Returns: bucket_name="my-bucket", object_name="my-folder/dev/my-object" """ - if "/" in bucket_name: - bucket_name, prefix = bucket_name.split("/", 1) + bucket_name, prefix = split_configured_cloud_bucket_name(bucket_name) + if prefix: object_name = f"{prefix}/{object_name}" return bucket_name, object_name return bucket_name, object_name @@ -248,6 +252,7 @@ async def download_gcs_object(self, object_name: str, **kwargs): bucket_name=bucket_name, object_name=object_name, ) + object_name = encode_gcs_object_name_for_url(object_name) url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{object_name}?alt=media" @@ -288,6 +293,7 @@ async def delete_gcs_object(self, object_name: str, **kwargs): bucket_name=bucket_name, object_name=object_name, ) + object_name = encode_gcs_object_name_for_url(object_name) url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{object_name}" @@ -334,10 +340,11 @@ async def _log_json_data_on_gcs( bucket_name=bucket_name, object_name=object_name, ) + encoded_object_name = encode_gcs_object_name_for_url(object_name) response = await self.async_httpx_client.post( headers=headers, - url=f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={object_name}", + url=f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}", data=json_logged_payload, ) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py new file mode 100644 index 000000000000..daa3dc603203 --- /dev/null +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -0,0 +1,175 @@ +import posixpath +import re +from types import MappingProxyType +from typing import Any, Mapping, Optional, Sequence, Tuple, cast +from urllib.parse import quote, unquote + +from litellm._uuid import uuid + +VERTEX_AI_MANAGED_GCS_PREFIX = "litellm-vertex-files/" +BEDROCK_MANAGED_S3_BATCH_PREFIX = "litellm-bedrock-files-" +BEDROCK_MANAGED_S3_UPLOAD_PREFIX = "litellm-bedrock-files/" +BEDROCK_MANAGED_S3_OUTPUT_PREFIX = "litellm-batch-outputs/" +BEDROCK_MANAGED_S3_PREFIXES = ( + BEDROCK_MANAGED_S3_BATCH_PREFIX, + BEDROCK_MANAGED_S3_UPLOAD_PREFIX, + BEDROCK_MANAGED_S3_OUTPUT_PREFIX, +) +_MAPPING_PROXY_TYPE: type = type(MappingProxyType({})) + +_SAFE_OBJECT_COMPONENT_PATTERN = re.compile(r"[^A-Za-z0-9._-]+") + + +def sanitize_cloud_object_component( + value: Optional[str], fallback: str = "file" +) -> str: + if not isinstance(value, str): + return fallback + + component = posixpath.basename(value.replace("\\", "/")).strip() + if component in {"", ".", ".."}: + return fallback + + component = "".join( + "_" if ord(char) < 32 or ord(char) == 127 else char for char in component + ) + component = _SAFE_OBJECT_COMPONENT_PATTERN.sub("_", component) + component = component.strip("._") + if not component: + return fallback + return component[:255] + + +def sanitize_cloud_object_path(value: Optional[str], fallback: str = "file") -> str: + if not isinstance(value, str): + return fallback + + segments = [] + for segment in value.replace("\\", "/").split("/"): + sanitized_segment = sanitize_cloud_object_component(segment, fallback="") + if sanitized_segment: + segments.append(sanitized_segment) + + if not segments: + return fallback + return "/".join(segments) + + +def build_managed_cloud_object_name( + prefix: str, filename: Optional[str], fallback_filename: str = "file" +) -> str: + safe_filename = sanitize_cloud_object_component( + filename, fallback=fallback_filename + ) + return f"{prefix}{uuid.uuid4().hex}-{safe_filename}" + + +def _validate_cloud_object_path(object_name: str) -> None: + if not object_name: + raise ValueError("Cloud storage object name is required") + if object_name.startswith("/"): + raise ValueError("Cloud storage object name must be relative") + if any(ord(char) < 32 or ord(char) == 127 for char in object_name): + raise ValueError("Cloud storage object name contains control characters") + segments = object_name.split("/") + if any(segment in {".", ".."} for segment in segments): + raise ValueError("Cloud storage object name contains an invalid path segment") + if "" in segments[:-1]: + raise ValueError("Cloud storage object name contains an invalid path segment") + + +def split_configured_cloud_bucket_name(bucket_name: str) -> Tuple[str, str]: + if not isinstance(bucket_name, str) or not bucket_name.strip(): + raise ValueError("Cloud storage bucket name is required") + + bucket_name = bucket_name.strip() + if "://" in bucket_name or "?" in bucket_name or "#" in bucket_name: + raise ValueError( + "Cloud storage bucket name must not include a URI scheme or query" + ) + if any(ord(char) < 32 or ord(char) == 127 for char in bucket_name): + raise ValueError("Cloud storage bucket name contains control characters") + + bucket, _, prefix = bucket_name.partition("/") + if not bucket: + raise ValueError("Cloud storage bucket name is required") + if "\\" in bucket: + raise ValueError("Cloud storage bucket name contains an invalid separator") + + prefix = prefix.strip("/") + if prefix: + _validate_cloud_object_path(prefix) + + return bucket, prefix + + +def encode_gcs_object_name_for_url(object_name: str) -> str: + return quote(unquote(object_name), safe="") + + +def encode_s3_object_key_for_url(object_key: str) -> str: + return quote(unquote(object_key), safe="/") + + +def should_allow_legacy_cloud_file_ids( + litellm_params: Optional[Mapping[str, Any]] = None, +) -> bool: + value = None + if isinstance(litellm_params, Mapping): + trusted_model_credentials = litellm_params.get( + "_litellm_internal_model_credentials" + ) + if isinstance(trusted_model_credentials, _MAPPING_PROXY_TYPE): + value = cast(Mapping[str, Any], trusted_model_credentials).get( + "allow_legacy_cloud_file_ids" + ) + + if isinstance(value, bool): + return value + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return False + + +def validate_managed_cloud_file_id( + file_id: str, + scheme: str, + configured_bucket_name: str, + allowed_object_prefixes: Sequence[str], + allow_legacy_cloud_file_ids: bool = False, +) -> Tuple[str, str]: + decoded_file_id = unquote(file_id) + if not decoded_file_id.startswith(scheme): + raise ValueError(f"file_id must be a {scheme} URI") + + full_path = decoded_file_id[len(scheme) :] + if "/" not in full_path: + raise ValueError("file_id must include a cloud storage object name") + + bucket_name, object_name = full_path.split("/", 1) + configured_bucket, configured_prefix = split_configured_cloud_bucket_name( + configured_bucket_name + ) + if bucket_name != configured_bucket: + raise ValueError("file_id bucket does not match the configured storage bucket") + + _validate_cloud_object_path(object_name) + allowed_prefixes = tuple(allowed_object_prefixes) + if configured_prefix: + allowed_prefixes = tuple( + f"{configured_prefix.rstrip('/')}/{prefix}" for prefix in allowed_prefixes + ) + + if object_name.startswith(allowed_prefixes): + return bucket_name, object_name + + if allow_legacy_cloud_file_ids: + if configured_prefix and not object_name.startswith( + f"{configured_prefix.rstrip('/')}/" + ): + raise ValueError( + "file_id object does not match the configured storage prefix" + ) + return bucket_name, object_name + + raise ValueError("file_id must reference a LiteLLM-managed storage object") diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index ffb6436f3875..a89dae523168 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -37,8 +37,6 @@ def validate_no_callback_env_reference( "langfuse_secret_key", "langfuse_host", "langfuse_prompt_version", - "gcs_bucket_name", - "gcs_path_service_account", "langsmith_api_key", "langsmith_project", "langsmith_base_url", @@ -57,6 +55,11 @@ def validate_no_callback_env_reference( "lunary_public_key", ] +_request_blocked_callback_params = { + "gcs_bucket_name", + "gcs_path_service_account", +} + def initialize_standard_callback_dynamic_params( kwargs: Optional[Dict] = None, @@ -64,13 +67,15 @@ def initialize_standard_callback_dynamic_params( """ Initialize the standard callback dynamic params from the kwargs - checks if langfuse_secret_key, gcs_bucket_name in kwargs and sets the corresponding attributes in StandardCallbackDynamicParams + checks supported request callback params in kwargs and sets the corresponding attributes in StandardCallbackDynamicParams """ standard_callback_dynamic_params = StandardCallbackDynamicParams() if kwargs: # 1. Check top-level kwargs for param in _supported_callback_params: + if param in _request_blocked_callback_params: + continue if param in kwargs: _param_value = kwargs.get(param) validate_no_callback_env_reference( @@ -86,6 +91,8 @@ def initialize_standard_callback_dynamic_params( if isinstance(metadata, dict): for param in _supported_callback_params: + if param in _request_blocked_callback_params: + continue if param not in standard_callback_dynamic_params and param in metadata: _param_value = metadata.get(param) validate_no_callback_env_reference( diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index 13bd87a1f019..ecf157e12eef 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -1,10 +1,17 @@ import asyncio import base64 -from typing import Any, Coroutine, Optional, Tuple, Union +import os +from types import MappingProxyType +from typing import Any, Coroutine, Mapping, Optional, Tuple, Union, cast import httpx from litellm import LlmProviders +from litellm.litellm_core_utils.cloud_storage_security import ( + BEDROCK_MANAGED_S3_PREFIXES, + should_allow_legacy_cloud_file_ids, + validate_managed_cloud_file_id, +) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.openai import ( FileContentRequest, @@ -35,7 +42,7 @@ def _extract_s3_uri_from_file_id(self, file_id: str) -> str: The file ID can be in two formats: 1. Base64-encoded unified file ID containing: llm_output_file_id,s3://bucket/path - 2. Direct S3 URI: s3://bucket/path + 2. Direct S3 URI: s3://bucket/litellm-managed-prefix/path Args: file_id: Encoded file ID or direct S3 URI @@ -58,14 +65,19 @@ def _extract_s3_uri_from_file_id(self, file_id: str) -> str: except Exception: pass - # If not base64 encoded or doesn't contain llm_output_file_id, assume it's already an S3 URI + # If not base64 encoded or doesn't contain llm_output_file_id, accept only + # explicit S3 URIs. Bucket and key validation happens before any S3 call. if file_id.startswith("s3://"): return file_id - # If it doesn't start with s3://, assume it's a direct S3 URI and add the prefix - return f"s3://{file_id}" + raise ValueError("file_id must be a managed LiteLLM S3 file id") - def _parse_s3_uri(self, s3_uri: str) -> Tuple[str, str]: + def _parse_s3_uri( + self, + s3_uri: str, + configured_bucket_name: str, + allow_legacy_cloud_file_ids: bool = False, + ) -> Tuple[str, str]: """ Parse S3 URI to extract bucket name and object key. @@ -75,21 +87,34 @@ def _parse_s3_uri(self, s3_uri: str) -> Tuple[str, str]: Returns: Tuple of (bucket_name, object_key) """ - if not s3_uri.startswith("s3://"): + return validate_managed_cloud_file_id( + file_id=s3_uri, + scheme="s3://", + configured_bucket_name=configured_bucket_name, + allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES, + allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids, + ) + + def _get_configured_s3_bucket_name(self, litellm_params: dict) -> str: + trusted_model_credentials = litellm_params.get( + "_litellm_internal_model_credentials" + ) + bucket_name = None + if isinstance(trusted_model_credentials, type(MappingProxyType({}))): + trusted_model_credentials_mapping = cast( + Mapping[str, Any], trusted_model_credentials + ) + candidate_bucket_name = trusted_model_credentials_mapping.get( + "s3_bucket_name" + ) + if isinstance(candidate_bucket_name, str): + bucket_name = candidate_bucket_name + bucket_name = bucket_name or os.getenv("AWS_S3_BUCKET_NAME") + if not bucket_name: raise ValueError( - f"Invalid S3 URI format: {s3_uri}. Expected format: s3://bucket-name/path/to/file" + "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval." ) - - # Remove 's3://' prefix - path = s3_uri[5:] - - if "/" in path: - bucket_name, object_key = path.split("/", 1) - else: - bucket_name = path - object_key = "" - - return bucket_name, object_key + return bucket_name async def afile_content( self, @@ -119,7 +144,14 @@ async def afile_content( # Extract S3 URI from file ID s3_uri = self._extract_s3_uri_from_file_id(file_id) - bucket_name, object_key = self._parse_s3_uri(s3_uri) + configured_bucket_name = self._get_configured_s3_bucket_name(optional_params) + bucket_name, object_key = self._parse_s3_uri( + s3_uri=s3_uri, + configured_bucket_name=configured_bucket_name, + allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids( + optional_params + ), + ) # Get AWS credentials aws_region_name = self._get_aws_region_name( diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 3007b54808ce..6669363093b2 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -2,6 +2,7 @@ import os import time from typing import Any, Dict, List, Optional, Tuple, Union +from urllib.parse import unquote import httpx from httpx import Headers, Response @@ -10,6 +11,14 @@ from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.files.utils import FilesAPIUtils +from litellm.litellm_core_utils.cloud_storage_security import ( + BEDROCK_MANAGED_S3_BATCH_PREFIX, + BEDROCK_MANAGED_S3_UPLOAD_PREFIX, + build_managed_cloud_object_name, + encode_s3_object_key_for_url, + sanitize_cloud_object_component, + split_configured_cloud_bucket_name, +) from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( @@ -116,10 +125,13 @@ def _get_s3_object_name_from_batch_jsonl( if _model.startswith("bedrock/"): _model = _model[8:] - # Replace colons with hyphens for Bedrock S3 URI compliance - _model = _model.replace(":", "-") + safe_model = sanitize_cloud_object_component( + _model.replace(":", "-"), fallback="model" + ) - object_name = f"litellm-bedrock-files-{_model}-{uuid.uuid4()}.jsonl" + object_name = ( + f"{BEDROCK_MANAGED_S3_BATCH_PREFIX}{safe_model}-{uuid.uuid4()}.jsonl" + ) return object_name def get_object_name( @@ -146,12 +158,13 @@ def get_object_name( if len(openai_jsonl_content) > 0: return self._get_s3_object_name_from_batch_jsonl(openai_jsonl_content) - ## 2. If not jsonl, return the filename + ## 2. If not jsonl, store under a server-generated managed object name filename = extracted_file_data.get("filename") - if filename: - return filename - ## 3. If no file name, return timestamp - return str(int(time.time())) + return build_managed_cloud_object_name( + prefix=BEDROCK_MANAGED_S3_UPLOAD_PREFIX, + filename=filename, + fallback_filename="file", + ) def get_complete_file_url( self, @@ -172,6 +185,7 @@ def get_complete_file_url( raise ValueError( "S3 bucket_name is required. Set 's3_bucket_name' in litellm_params or AWS_S3_BUCKET_NAME env var" ) + bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) s3_region_name = litellm_params.get("s3_region_name") or optional_params.get( "s3_region_name" @@ -188,14 +202,17 @@ def get_complete_file_url( raise ValueError("purpose is required") extracted_file_data = extract_file_data(file_data) object_name = self.get_object_name(extracted_file_data, purpose) + if object_prefix: + object_name = f"{object_prefix}/{object_name}" + encoded_object_name = encode_s3_object_key_for_url(object_name) # S3 endpoint URL format s3_endpoint_url = ( optional_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com" - ) + ).rstrip("/") - return f"{s3_endpoint_url}/{bucket_name}/{object_name}" + return f"{s3_endpoint_url}/{bucket_name}/{encoded_object_name}" def get_supported_openai_params( self, model: str @@ -532,10 +549,12 @@ def _convert_https_url_to_s3_uri(self, https_url: str) -> tuple[str, str]: if match1: # Pattern: https://s3.region.amazonaws.com/bucket/key region, bucket, key = match1.groups() + key = unquote(key) s3_uri = f"s3://{bucket}/{key}" elif match2: # Pattern: https://bucket.s3.region.amazonaws.com/key bucket, region, key = match2.groups() + key = unquote(key) s3_uri = f"s3://{bucket}/{key}" else: # Fallback: try to extract bucket and key from URL path @@ -545,6 +564,7 @@ def _convert_https_url_to_s3_uri(self, https_url: str) -> tuple[str, str]: path_parts = parsed.path.lstrip("/").split("/", 1) if len(path_parts) >= 2: bucket, key = path_parts[0], path_parts[1] + key = unquote(key) s3_uri = f"s3://{bucket}/{key}" else: raise ValueError(f"Unable to parse S3 URL: {https_url}") @@ -722,7 +742,12 @@ def _get_s3_object_name( # Remove bedrock/ prefix if present if _model.startswith("bedrock/"): _model = _model[8:] - object_name = f"litellm-bedrock-files-{_model}-{uuid.uuid4()}.jsonl" + safe_model = sanitize_cloud_object_component( + _model.replace(":", "-"), fallback="model" + ) + object_name = ( + f"{BEDROCK_MANAGED_S3_BATCH_PREFIX}{safe_model}-{uuid.uuid4()}.jsonl" + ) return object_name def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str: diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index bd4b2ac8bbca..c31bfde69e75 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -1,6 +1,6 @@ import asyncio import time -import urllib.parse +from urllib.parse import unquote from typing import Any, Coroutine, Optional, Tuple, Union import httpx @@ -10,6 +10,11 @@ GCSBucketBase, GCSLoggingConfig, ) +from litellm.litellm_core_utils.cloud_storage_security import ( + VERTEX_AI_MANAGED_GCS_PREFIX, + should_allow_legacy_cloud_file_ids, + validate_managed_cloud_file_id, +) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.openai import ( CreateFileRequest, @@ -114,34 +119,31 @@ def create_file( ) ) - def _extract_bucket_and_object_from_file_id(self, file_id: str) -> Tuple[str, str]: + def _extract_bucket_and_object_from_file_id( + self, + file_id: str, + configured_bucket_name: str, + litellm_params: Optional[dict] = None, + ) -> Tuple[str, str]: """ - Extract bucket name and object path from URL-encoded file_id. + Validate and extract bucket name and object path from file_id. - Expected format: gs%3A%2F%2Fbucket-name%2Fpath%2Fto%2Ffile - Which decodes to: gs://bucket-name/path/to/file + Expected format: gs://bucket-name/litellm-vertex-files/path/to/file Returns: - tuple: (bucket_name, url_encoded_object_path) + tuple: (bucket_name, object_path) - bucket_name: "bucket-name" - - url_encoded_object_path: "path%2Fto%2Ffile" + - object_path: "litellm-vertex-files/path/to/file" """ - decoded_path = urllib.parse.unquote(file_id) - - if decoded_path.startswith("gs://"): - full_path = decoded_path[5:] # Remove 'gs://' prefix - else: - full_path = decoded_path - - if "/" in full_path: - bucket_name, object_path = full_path.split("/", 1) - else: - bucket_name = full_path - object_path = "" - - encoded_object_path = urllib.parse.quote(object_path, safe="") - - return bucket_name, encoded_object_path + return validate_managed_cloud_file_id( + file_id=file_id, + scheme="gs://", + configured_bucket_name=configured_bucket_name, + allowed_object_prefixes=(VERTEX_AI_MANAGED_GCS_PREFIX,), + allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids( + litellm_params + ), + ) async def afile_content( self, @@ -151,6 +153,7 @@ async def afile_content( vertex_location: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], + litellm_params: Optional[dict] = None, ) -> HttpxBinaryResponseContent: """ Download file content from GCS bucket for VertexAI files. @@ -170,23 +173,30 @@ async def afile_content( if not file_id: raise ValueError("file_id is required in file_content_request") - bucket_name, encoded_object_path = self._extract_bucket_and_object_from_file_id( - file_id + gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config( + kwargs={} + ) + bucket_name, object_path = self._extract_bucket_and_object_from_file_id( + file_id=file_id, + configured_bucket_name=gcs_logging_config["bucket_name"], + litellm_params=litellm_params, ) download_kwargs = { - "standard_callback_dynamic_params": {"gcs_bucket_name": bucket_name} + "standard_callback_dynamic_params": { + "gcs_bucket_name": bucket_name, + "gcs_path_service_account": gcs_logging_config["path_service_account"], + } } file_content = await self.download_gcs_object( - object_name=encoded_object_path, **download_kwargs + object_name=object_path, **download_kwargs ) + decoded_file_id = unquote(file_id) if file_content is None: - decoded_path = urllib.parse.unquote(file_id) - raise ValueError(f"Failed to download file from GCS: {decoded_path}") + raise ValueError(f"Failed to download file from GCS: {decoded_file_id}") - decoded_path = urllib.parse.unquote(file_id) mock_response = httpx.Response( status_code=200, content=file_content, @@ -194,7 +204,7 @@ async def afile_content( "content-type": "application/octet-stream", "content-length": str(len(file_content)), }, - request=httpx.Request(method="GET", url=decoded_path), + request=httpx.Request(method="GET", url=decoded_file_id), ) # Apply transformation to convert Vertex AI batch outputs to OpenAI format @@ -225,6 +235,7 @@ def file_content( vertex_location: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], + litellm_params: Optional[dict] = None, ) -> Union[ HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent] ]: @@ -253,6 +264,7 @@ def file_content( vertex_location=vertex_location, timeout=timeout, max_retries=max_retries, + litellm_params=litellm_params, ) else: return asyncio.run( @@ -263,5 +275,6 @@ def file_content( vertex_location=vertex_location, timeout=timeout, max_retries=max_retries, + litellm_params=litellm_params, ) ) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 01b2125d0303..f30518bc7ca0 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -12,6 +12,15 @@ import litellm from litellm._uuid import uuid from litellm.files.utils import FilesAPIUtils +from litellm.litellm_core_utils.cloud_storage_security import ( + VERTEX_AI_MANAGED_GCS_PREFIX, + build_managed_cloud_object_name, + encode_gcs_object_name_for_url, + sanitize_cloud_object_path, + should_allow_legacy_cloud_file_ids, + split_configured_cloud_bucket_name, + validate_managed_cloud_file_id, +) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -248,7 +257,8 @@ def _get_gcs_object_name_from_batch_jsonl( _model = openai_jsonl_content[0].get("body", {}).get("model", "") if "publishers/google/models" not in _model: _model = f"publishers/google/models/{_model}" - object_name = f"litellm-vertex-files/{_model}/{uuid.uuid4()}" + safe_model_path = sanitize_cloud_object_path(_model, fallback="model") + object_name = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" return object_name def get_object_name( @@ -275,12 +285,19 @@ def get_object_name( if len(openai_jsonl_content) > 0: return self._get_gcs_object_name_from_batch_jsonl(openai_jsonl_content) - ## 2. If not jsonl, return the filename + ## 2. If not jsonl, store under a server-generated managed object name filename = extracted_file_data.get("filename") - if filename: - return filename - ## 3. If no file name, return timestamp - return str(int(time.time())) + return build_managed_cloud_object_name( + prefix=f"{VERTEX_AI_MANAGED_GCS_PREFIX}uploads/", + filename=filename, + fallback_filename="file", + ) + + def _get_configured_bucket_name(self, litellm_params: Dict) -> str: + bucket_name = litellm_params.get("bucket_name") or os.getenv("GCS_BUCKET_NAME") + if not bucket_name: + raise ValueError("GCS bucket_name is required") + return bucket_name def get_complete_file_url( self, @@ -294,13 +311,8 @@ def get_complete_file_url( """ Get the complete url for the request """ - bucket_name = ( - litellm_params.get("bucket_name") - or litellm_params.get("litellm_metadata", {}).pop("gcs_bucket_name", None) - or os.getenv("GCS_BUCKET_NAME") - ) - if not bucket_name: - raise ValueError("GCS bucket_name is required") + bucket_name = self._get_configured_bucket_name(litellm_params) + bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) file_data = data.get("file") purpose = data.get("purpose") if file_data is None: @@ -309,9 +321,10 @@ def get_complete_file_url( raise ValueError("purpose is required") extracted_file_data = extract_file_data(file_data) object_name = self.get_object_name(extracted_file_data, purpose) - endpoint = ( - f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={object_name}" - ) + if object_prefix: + object_name = f"{object_prefix}/{object_name}" + encoded_object_name = encode_gcs_object_name_for_url(object_name) + endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}" api_base = api_base or "https://storage.googleapis.com" if not api_base: raise ValueError("api_base is required") @@ -450,27 +463,23 @@ def get_error_class( status_code=status_code, message=error_message, headers=headers ) - def _parse_gcs_uri(self, file_id: str) -> Tuple[str, str]: + def _parse_gcs_uri( + self, file_id: str, litellm_params: Optional[Dict] = None + ) -> Tuple[str, str]: """ - Parse a GCS URI (gs://bucket/path/to/object) into (bucket, url-encoded-object-path). - Handles both raw and URL-encoded input. + Validate a managed GCS file_id and return (bucket, url-encoded-object-path). """ - import urllib.parse - - decoded = urllib.parse.unquote(file_id) - if decoded.startswith("gs://"): - full_path = decoded[5:] - else: - full_path = decoded - - if "/" in full_path: - bucket_name, object_path = full_path.split("/", 1) - else: - bucket_name = full_path - object_path = "" - - encoded_object = urllib.parse.quote(object_path, safe="") - return bucket_name, encoded_object + configured_bucket_name = self._get_configured_bucket_name(litellm_params or {}) + bucket_name, object_path = validate_managed_cloud_file_id( + file_id=file_id, + scheme="gs://", + configured_bucket_name=configured_bucket_name, + allowed_object_prefixes=(VERTEX_AI_MANAGED_GCS_PREFIX,), + allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids( + litellm_params + ), + ) + return bucket_name, encode_gcs_object_name_for_url(object_path) def transform_retrieve_file_request( self, @@ -478,7 +487,7 @@ def transform_retrieve_file_request( optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - bucket, encoded_object = self._parse_gcs_uri(file_id) + bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params) url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}" return url, {} @@ -510,7 +519,7 @@ def transform_delete_file_request( optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - bucket, encoded_object = self._parse_gcs_uri(file_id) + bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params) url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}" return url, {} @@ -554,7 +563,7 @@ def transform_file_content_request( litellm_params: dict, ) -> tuple[str, dict]: file_id = file_content_request.get("file_id", "") - bucket, encoded_object = self._parse_gcs_uri(file_id) + bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params) url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}?alt=media" return url, {} @@ -842,7 +851,8 @@ def _get_gcs_object_name( _model = openai_jsonl_content[0].get("body", {}).get("model", "") if "publishers/google/models" not in _model: _model = f"publishers/google/models/{_model}" - object_name = f"litellm-vertex-files/{_model}/{uuid.uuid4()}" + safe_model_path = sanitize_cloud_object_path(_model, fallback="model") + object_name = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" return object_name def _map_openai_to_vertex_params( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 25ad1b25aad6..30c78ed5ba77 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -2,6 +2,7 @@ import mimetypes import re from dataclasses import dataclass, field +from types import MappingProxyType from typing import TYPE_CHECKING, List, Literal, Optional, Union from litellm.types.utils import SpecialEnums @@ -298,6 +299,7 @@ def prepare_data_with_credentials( data: dict, credentials: dict, file_id: Optional[str] = None, + include_internal_credentials: bool = False, ) -> None: """ Update data dictionary with model credentials (in-place). @@ -306,8 +308,14 @@ def prepare_data_with_credentials( data: Data dictionary to update credentials: Credentials from router file_id: Optional original file_id to set (for decoded file IDs) + include_internal_credentials: Preserve an immutable server-side snapshot + for code paths that must distinguish proxy config from request params. """ data.update(credentials) + if include_internal_credentials: + data["_litellm_internal_model_credentials"] = MappingProxyType( + dict(credentials) + ) data.pop("custom_llm_provider", None) if file_id is not None: diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index c9be707af0c8..378cbbda89c0 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -774,6 +774,7 @@ async def get_file_content( # noqa: PLR0915 data=data, credentials=credentials, # type: ignore file_id=original_file_id, # Use decoded file ID if from encoded ID + include_internal_credentials=True, ) response = await litellm.afile_content( custom_llm_provider=credentials["custom_llm_provider"], # type: ignore @@ -949,6 +950,7 @@ async def get_file( data=data, credentials=credentials, # type: ignore file_id=original_file_id, + include_internal_credentials=True, ) response = await litellm.afile_retrieve(**data) # type: ignore @@ -1149,6 +1151,7 @@ async def delete_file( data=data, credentials=credentials, # type: ignore file_id=original_file_id, + include_internal_credentials=True, ) response = await litellm.afile_delete( diff --git a/tests/litellm/proxy/test_model_based_routing_files_batches.py b/tests/litellm/proxy/test_model_based_routing_files_batches.py index 94e2c4603bc1..f9bc07212a26 100644 --- a/tests/litellm/proxy/test_model_based_routing_files_batches.py +++ b/tests/litellm/proxy/test_model_based_routing_files_batches.py @@ -6,10 +6,13 @@ and file proxy endpoints. """ +from types import MappingProxyType + from litellm.proxy.openai_files_endpoints.common_utils import ( decode_model_from_file_id, encode_file_id_with_model, get_original_file_id, + prepare_data_with_credentials, ) @@ -50,6 +53,42 @@ def test_gcs_uri_gets_file_prefix(self): assert result.startswith("file-") +class TestPrepareDataWithCredentials: + def test_preserves_trusted_internal_credentials_snapshot(self): + data = {"file_id": "file-abc"} + credentials = { + "custom_llm_provider": "bedrock", + "s3_bucket_name": "safe-bucket", + } + + prepare_data_with_credentials( + data=data, + credentials=credentials, + include_internal_credentials=True, + ) + + assert data["s3_bucket_name"] == "safe-bucket" + assert "custom_llm_provider" not in data + assert isinstance( + data["_litellm_internal_model_credentials"], type(MappingProxyType({})) + ) + assert ( + data["_litellm_internal_model_credentials"]["s3_bucket_name"] + == "safe-bucket" + ) + + def test_does_not_add_internal_credentials_by_default(self): + data = {"file_id": "file-abc"} + credentials = { + "custom_llm_provider": "bedrock", + "s3_bucket_name": "safe-bucket", + } + + prepare_data_with_credentials(data=data, credentials=credentials) + + assert "_litellm_internal_model_credentials" not in data + + class TestRoundTrip: """Tests for encode -> decode round-trip integrity.""" diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py index d7b2a842bc08..a4e16500aee5 100644 --- a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py +++ b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py @@ -1,5 +1,9 @@ import os -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase @@ -86,3 +90,41 @@ def test_construct_request_headers_without_project_id(self): project_id=None, # Should be None when no env var is set custom_llm_provider="vertex_ai", ) + + @pytest.mark.asyncio + async def test_log_json_data_on_gcs_url_encodes_object_name(self): + handler = GCSBucketBase(bucket_name="test-bucket") + handler.async_httpx_client = AsyncMock() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"name": "logs/object"} + handler.async_httpx_client.post.return_value = mock_response + + await handler._log_json_data_on_gcs( + headers={"Authorization": "Bearer token"}, + bucket_name="test-bucket", + object_name="logs/object?uploadType=media&name=evil", + logging_payload={"ok": True}, + ) + + post_url = handler.async_httpx_client.post.call_args.kwargs["url"] + assert "name=logs%2Fobject%3FuploadType%3Dmedia%26name%3Devil" in post_url + assert "name=logs/object?" not in post_url + + def test_gcs_log_id_is_only_used_as_sanitized_hint(self): + logger = GCSBucketLogger.__new__(GCSBucketLogger) + + object_name = logger._get_object_name( + kwargs={ + "litellm_params": { + "metadata": {"gcs_log_id": "../../target?uploadType=media"} + } + }, + logging_payload={"id": "payload"}, + response_obj={"id": "response-id"}, + ) + + assert "/custom-" in object_name + assert object_name.endswith("-target_uploadType_media") + assert ".." not in object_name + assert "?" not in object_name diff --git a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py index 55f3c2ba3aa7..f63216b96d43 100644 --- a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py +++ b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py @@ -64,7 +64,7 @@ def test_env_reference_in_metadata_raises_with_guidance(): assert "metadata" in message -def test_env_reference_in_litellm_params_metadata_raises(): +def test_gcs_bucket_name_in_litellm_params_metadata_is_ignored(): kwargs = { "litellm_params": { "metadata": { @@ -73,10 +73,21 @@ def test_env_reference_in_litellm_params_metadata_raises(): } } - with pytest.raises(ValueError) as exc_info: - initialize_standard_callback_dynamic_params(kwargs) + params = initialize_standard_callback_dynamic_params(kwargs) + + assert params.get("gcs_bucket_name") is None + + +def test_gcs_callback_params_are_not_extracted_from_request_kwargs(): + kwargs = { + "gcs_bucket_name": "server-bucket", + "gcs_path_service_account": "/path/to/service-account.json", + } + + params = initialize_standard_callback_dynamic_params(kwargs) - assert "gcs_bucket_name" in str(exc_info.value) + assert params.get("gcs_bucket_name") is None + assert params.get("gcs_path_service_account") is None def test_non_string_values_are_not_flagged(): diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py new file mode 100644 index 000000000000..7f91b49a6f5d --- /dev/null +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py @@ -0,0 +1,206 @@ +import base64 +import os +from types import MappingProxyType +from unittest.mock import MagicMock, patch + +import pytest + +import litellm.files.main as files_main +from litellm.llms.bedrock.files.handler import BedrockFilesHandler +from litellm.types.utils import SpecialEnums + + +def _encode_unified_file_id(s3_uri: str) -> str: + unified_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", + "unified-id", + "", + s3_uri, + "model-id", + ) + return base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") + + +class TestBedrockFilesHandler: + def setup_method(self): + self.handler = BedrockFilesHandler() + + def test_should_parse_direct_managed_s3_uri(self): + bucket, key = self.handler._parse_s3_uri( + s3_uri="s3://safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", + configured_bucket_name="safe-bucket", + ) + + assert bucket == "safe-bucket" + assert key == "litellm-bedrock-files-model-id-abc.jsonl" + + def test_should_parse_managed_batch_output_uri(self): + bucket, key = self.handler._parse_s3_uri( + s3_uri="s3://safe-bucket/litellm-batch-outputs/job/", + configured_bucket_name="safe-bucket", + ) + + assert bucket == "safe-bucket" + assert key == "litellm-batch-outputs/job/" + + def test_should_reject_arbitrary_bucket(self): + with pytest.raises(ValueError, match="configured storage bucket"): + self.handler._parse_s3_uri( + s3_uri="s3://other-bucket/litellm-bedrock-files-model-id-abc.jsonl", + configured_bucket_name="safe-bucket", + ) + + def test_should_reject_unmanaged_same_bucket_key(self): + with pytest.raises(ValueError, match="LiteLLM-managed"): + self.handler._parse_s3_uri( + s3_uri="s3://safe-bucket/private/output.jsonl", + configured_bucket_name="safe-bucket", + ) + + def test_should_allow_legacy_same_bucket_key_when_server_flag_enabled(self): + bucket, key = self.handler._parse_s3_uri( + s3_uri="s3://safe-bucket/private/output.jsonl", + configured_bucket_name="safe-bucket", + allow_legacy_cloud_file_ids=True, + ) + + assert bucket == "safe-bucket" + assert key == "private/output.jsonl" + + def test_should_keep_configured_prefix_for_legacy_keys(self): + bucket, key = self.handler._parse_s3_uri( + s3_uri="s3://safe-bucket/team-a/private/output.jsonl", + configured_bucket_name="safe-bucket/team-a", + allow_legacy_cloud_file_ids=True, + ) + + assert bucket == "safe-bucket" + assert key == "team-a/private/output.jsonl" + + def test_should_reject_legacy_key_outside_configured_prefix(self): + with pytest.raises(ValueError, match="configured storage prefix"): + self.handler._parse_s3_uri( + s3_uri="s3://safe-bucket/team-b/private/output.jsonl", + configured_bucket_name="safe-bucket/team-a", + allow_legacy_cloud_file_ids=True, + ) + + def test_should_reject_dot_segment_key(self): + with pytest.raises(ValueError, match="invalid path segment"): + self.handler._parse_s3_uri( + s3_uri="s3://safe-bucket/litellm-bedrock-files/../secret.jsonl", + configured_bucket_name="safe-bucket", + ) + + def test_should_reject_empty_middle_path_segment(self): + with pytest.raises(ValueError, match="invalid path segment"): + self.handler._parse_s3_uri( + s3_uri="s3://safe-bucket/litellm-bedrock-files//secret.jsonl", + configured_bucket_name="safe-bucket", + ) + + def test_should_extract_unified_managed_s3_uri(self): + file_id = _encode_unified_file_id( + "s3://safe-bucket/litellm-batch-outputs/job/output.jsonl" + ) + + assert ( + self.handler._extract_s3_uri_from_file_id(file_id) + == "s3://safe-bucket/litellm-batch-outputs/job/output.jsonl" + ) + + def test_should_reject_file_id_without_s3_scheme(self): + with pytest.raises(ValueError, match="managed LiteLLM S3 file id"): + self.handler._extract_s3_uri_from_file_id("safe-bucket/private.jsonl") + + def test_should_reject_unified_unmanaged_s3_uri(self): + file_id = _encode_unified_file_id("s3://safe-bucket/private/output.jsonl") + s3_uri = self.handler._extract_s3_uri_from_file_id(file_id) + + with pytest.raises(ValueError, match="LiteLLM-managed"): + self.handler._parse_s3_uri( + s3_uri=s3_uri, + configured_bucket_name="safe-bucket", + ) + + def test_should_not_trust_request_s3_bucket_name_for_expected_bucket(self): + with patch.dict(os.environ, {"AWS_S3_BUCKET_NAME": "safe-bucket"}): + assert ( + self.handler._get_configured_s3_bucket_name( + {"s3_bucket_name": "attacker-bucket"} + ) + == "safe-bucket" + ) + + def test_should_trust_proxy_config_s3_bucket_name_for_expected_bucket(self): + trusted_credentials = MappingProxyType({"s3_bucket_name": "safe-bucket"}) + + with patch.dict(os.environ, {}, clear=True): + assert ( + self.handler._get_configured_s3_bucket_name( + { + "s3_bucket_name": "attacker-bucket", + "_litellm_internal_model_credentials": trusted_credentials, + } + ) + == "safe-bucket" + ) + + def test_should_not_trust_user_supplied_internal_credentials_dict(self): + with patch.dict(os.environ, {}, clear=True): + with pytest.raises(ValueError, match="S3 bucket_name is required"): + self.handler._get_configured_s3_bucket_name( + { + "_litellm_internal_model_credentials": { + "s3_bucket_name": "attacker-bucket" + } + } + ) + + def test_should_require_server_s3_bucket_name(self): + with patch.dict(os.environ, {}, clear=True): + with pytest.raises(ValueError, match="S3 bucket_name is required"): + self.handler._get_configured_s3_bucket_name( + {"s3_bucket_name": "attacker-bucket"} + ) + + +def test_should_forward_trusted_model_credentials_to_bedrock_provider_config(): + trusted_credentials = MappingProxyType({"s3_bucket_name": "safe-bucket"}) + mock_response = MagicMock() + + with patch.object( + files_main.base_llm_http_handler, + "retrieve_file_content", + return_value=mock_response, + ) as mock_retrieve_file_content: + response = files_main.file_content( + file_id="s3://safe-bucket/litellm-bedrock-files/file.jsonl", + custom_llm_provider="bedrock", + _litellm_internal_model_credentials=trusted_credentials, + ) + + assert response is mock_response + litellm_params = mock_retrieve_file_content.call_args.kwargs["litellm_params"] + assert litellm_params["_litellm_internal_model_credentials"] is trusted_credentials + assert "s3_bucket_name" not in litellm_params + + +def test_should_forward_trusted_model_credentials_to_retrieve_provider_config(): + trusted_credentials = MappingProxyType({"allow_legacy_cloud_file_ids": True}) + mock_response = MagicMock() + + with patch.object( + files_main.base_llm_http_handler, + "retrieve_file", + return_value=mock_response, + ) as mock_retrieve_file: + response = files_main.file_retrieve( + file_id="gs://safe-bucket/private/file.jsonl", + custom_llm_provider="vertex_ai", + _litellm_internal_model_credentials=trusted_credentials, + ) + + assert response is mock_response + litellm_params = mock_retrieve_file.call_args.kwargs["litellm_params"] + assert litellm_params["_litellm_internal_model_credentials"] is trusted_credentials diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index d9a2ddefd344..5245612e9d3a 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -4,9 +4,7 @@ import json import os -from typing import Any, Dict, List - -import pytest +from urllib.parse import unquote, urlparse from litellm.llms.bedrock.files.transformation import BedrockJsonlFilesTransformation @@ -43,19 +41,6 @@ def test_transform_openai_jsonl_content_to_bedrock_jsonl_content(self): ) ) - # Print the transformation results for validation - print("\n=== INPUT (OpenAI format) ===") - for i, content in enumerate(openai_jsonl_content): - print(f"Record {i+1}:") - print(json.dumps(content, indent=2)) - print() - - print("\n=== OUTPUT (Bedrock format) ===") - for i, content in enumerate(bedrock_jsonl_content): - print(f"Record {i+1}:") - print(json.dumps(content, indent=2)) - print() - # Basic validation assert len(bedrock_jsonl_content) == len( openai_jsonl_content @@ -88,17 +73,6 @@ def test_transform_openai_jsonl_content_to_bedrock_jsonl_content(self): "max_tokens" in model_input ), f"Record {i+1} should have max_tokens" - # Write expected output to file for reference - expected_output_path = os.path.join( - os.path.dirname(__file__), "expected_bedrock_batch_completions.jsonl" - ) - - with open(expected_output_path, "w") as f: - for record in bedrock_jsonl_content: - f.write(json.dumps(record) + "\n") - - print(f"\n=== Expected output written to: {expected_output_path} ===") - def test_nova_text_only_uses_converse_format(self): """ Test that Nova models produce Converse API format in batch modelInput. @@ -327,6 +301,44 @@ def test_get_complete_file_url_respects_s3_region_name(self): ), f"us-west-2 must not appear when s3_region_name is set, got: {url}" assert "litellm-batch-352026" in url + def test_get_complete_file_url_sanitizes_untrusted_filename(self): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + create_file_data = { + "file": ("../../owned.jsonl?acl=public", b"hello", "application/jsonl"), + "purpose": "assistants", + } + + url = config.get_complete_file_url( + api_base=None, + api_key=None, + model="amazon.nova-pro-v1:0", + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={"s3_bucket_name": "safe-bucket"}, + data=create_file_data, + ) + + parsed_url = urlparse(url) + object_key = unquote(parsed_url.path).split("/safe-bucket/", 1)[1] + assert object_key.startswith("litellm-bedrock-files/") + assert object_key.endswith("-owned.jsonl_acl_public") + assert ".." not in object_key + assert parsed_url.query == "" + + def test_batch_object_name_sanitizes_model_path(self): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + object_name = config._get_s3_object_name_from_batch_jsonl( + [{"body": {"model": "bedrock/../../secret:model"}}] + ) + + assert object_name.startswith("litellm-bedrock-files-") + assert object_name.endswith(".jsonl") + assert "/" not in object_name + assert ".." not in object_name + def test_transform_create_file_request_injects_s3_region_for_signing(self): """ When s3_region_name is provided, transform_create_file_request must pass diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py index ea056a2fc80c..453a0c14bf91 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -3,6 +3,7 @@ """ import asyncio +from types import MappingProxyType import pytest from unittest.mock import AsyncMock, patch @@ -12,6 +13,14 @@ from litellm.types.llms.openai import FileContentRequest, HttpxBinaryResponseContent +def _mock_gcs_logging_config(bucket_name: str = "test-bucket"): + return { + "bucket_name": bucket_name, + "path_service_account": None, + "vertex_instance": None, + } + + class TestVertexAIFilesHandler: """Test Vertex AI files handler""" @@ -22,57 +31,84 @@ def setup_method(self): def test_extract_bucket_and_object_from_file_id_standard_path(self): """Test extraction of bucket and object from URL-encoded file_id with standard path""" # Sample file_id with nested folder structure - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-folder" "%2Fsub-folder%2Ftest-file.txt" + file_id = ( + "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files" + "%2Ftest-folder%2Fsub-folder%2Ftest-file.txt" + ) - bucket_name, encoded_object_path = ( - self.handler._extract_bucket_and_object_from_file_id(file_id) + bucket_name, object_path = self.handler._extract_bucket_and_object_from_file_id( + file_id=file_id, + configured_bucket_name="test-bucket", ) # Verify bucket name extraction assert bucket_name == "test-bucket" - # Verify object path encoding - expected_encoded_object = "test-folder%2Fsub-folder%2Ftest-file.txt" - assert encoded_object_path == expected_encoded_object + expected_object = "litellm-vertex-files/test-folder/sub-folder/test-file.txt" + assert object_path == expected_object - def test_extract_bucket_and_object_from_file_id_bucket_only(self): + def test_extract_bucket_and_object_from_file_id_rejects_bucket_only(self): """Test extraction when only bucket name is provided""" file_id = "gs%3A%2F%2Ftest-bucket" - bucket_name, encoded_object_path = ( - self.handler._extract_bucket_and_object_from_file_id(file_id) - ) - - assert bucket_name == "test-bucket" - assert encoded_object_path == "" + with pytest.raises(ValueError, match="object name"): + self.handler._extract_bucket_and_object_from_file_id( + file_id=file_id, + configured_bucket_name="test-bucket", + ) - def test_extract_bucket_and_object_from_file_id_simple_path(self): + def test_extract_bucket_and_object_from_file_id_rejects_unmanaged_path(self): """Test extraction with simple path""" file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - bucket_name, encoded_object_path = ( - self.handler._extract_bucket_and_object_from_file_id(file_id) + with pytest.raises(ValueError, match="LiteLLM-managed"): + self.handler._extract_bucket_and_object_from_file_id( + file_id=file_id, + configured_bucket_name="test-bucket", + ) + + def test_extract_bucket_and_object_from_file_id_allows_trusted_legacy_flag(self): + file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" + trusted_credentials = MappingProxyType({"allow_legacy_cloud_file_ids": True}) + + bucket_name, object_path = self.handler._extract_bucket_and_object_from_file_id( + file_id=file_id, + configured_bucket_name="test-bucket", + litellm_params={ + "_litellm_internal_model_credentials": trusted_credentials, + }, ) assert bucket_name == "test-bucket" - assert encoded_object_path == "test-file.txt" + assert object_path == "test-file.txt" - def test_extract_bucket_and_object_from_file_id_no_gs_prefix(self): + def test_extract_bucket_and_object_from_file_id_rejects_no_gs_prefix(self): """Test extraction when gs:// prefix is missing""" - file_id = "test-bucket%2Ftest-file.txt" + file_id = "test-bucket%2Flitellm-vertex-files%2Ftest-file.txt" - bucket_name, encoded_object_path = ( - self.handler._extract_bucket_and_object_from_file_id(file_id) - ) + with pytest.raises(ValueError, match="gs://"): + self.handler._extract_bucket_and_object_from_file_id( + file_id=file_id, + configured_bucket_name="test-bucket", + ) - assert bucket_name == "test-bucket" - assert encoded_object_path == "test-file.txt" + def test_extract_bucket_and_object_from_file_id_rejects_wrong_bucket(self): + file_id = "gs%3A%2F%2Fother-bucket%2Flitellm-vertex-files%2Ftest-file.txt" + + with pytest.raises(ValueError, match="configured storage bucket"): + self.handler._extract_bucket_and_object_from_file_id( + file_id=file_id, + configured_bucket_name="test-bucket", + ) @pytest.mark.asyncio async def test_afile_content_success(self): """Test successful async file content retrieval""" # Setup test data - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" + file_id = ( + "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files" + "%2Fuploads%2Fabc-test-file.txt" + ) expected_content = b"test file content" file_content_request = FileContentRequest( @@ -80,9 +116,17 @@ async def test_afile_content_success(self): ) # Mock the download_gcs_object method - with patch.object( - self.handler, "download_gcs_object", new_callable=AsyncMock - ) as mock_download: + with ( + patch.object( + self.handler, "download_gcs_object", new_callable=AsyncMock + ) as mock_download, + patch.object( + self.handler, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_mock_gcs_logging_config(), + ), + ): mock_download.return_value = expected_content # Call the method @@ -104,7 +148,10 @@ async def test_afile_content_success(self): # Verify the download was called with correct parameters mock_download.assert_called_once() call_args = mock_download.call_args - assert call_args.kwargs["object_name"] == "test-file.txt" + assert ( + call_args.kwargs["object_name"] + == "litellm-vertex-files/uploads/abc-test-file.txt" + ) assert "standard_callback_dynamic_params" in call_args.kwargs assert ( call_args.kwargs["standard_callback_dynamic_params"]["gcs_bucket_name"] @@ -132,22 +179,33 @@ async def test_afile_content_missing_file_id(self): @pytest.mark.asyncio async def test_afile_content_download_failure(self): """Test async file content retrieval when download fails""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" + file_id = ( + "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files" + "%2Fuploads%2Fabc-test-file.txt" + ) file_content_request = FileContentRequest( file_id=file_id, extra_headers=None, extra_body=None ) # Mock download to return None (failure) - with patch.object( - self.handler, "download_gcs_object", new_callable=AsyncMock - ) as mock_download: + with ( + patch.object( + self.handler, "download_gcs_object", new_callable=AsyncMock + ) as mock_download, + patch.object( + self.handler, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_mock_gcs_logging_config(), + ), + ): mock_download.return_value = None # Should raise ValueError for failed download with pytest.raises( ValueError, - match="Failed to download file from GCS: gs://test-bucket/test-file.txt", + match="Failed to download file from GCS: gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt", ): await self.handler.afile_content( file_content_request=file_content_request, diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index b8482ce6aae5..7c063c726076 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -5,6 +5,8 @@ import json import urllib.parse +from types import MappingProxyType +from urllib.parse import parse_qs, urlparse import httpx import pytest @@ -29,13 +31,20 @@ class TestParseGcsUri: """Tests for the _parse_gcs_uri helper used by retrieve / content / delete.""" def test_should_parse_standard_gs_uri(self, config): - bucket, encoded = config._parse_gcs_uri("gs://my-bucket/path/to/object.jsonl") + file_id = "gs://my-bucket/litellm-vertex-files/path/to/object.jsonl" + bucket, encoded = config._parse_gcs_uri( + file_id, litellm_params={"bucket_name": "my-bucket"} + ) assert bucket == "my-bucket" - assert encoded == urllib.parse.quote("path/to/object.jsonl", safe="") + assert encoded == urllib.parse.quote( + "litellm-vertex-files/path/to/object.jsonl", safe="" + ) def test_should_parse_uri_with_nested_publisher_path(self, config): uri = "gs://litellm-local/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" - bucket, encoded = config._parse_gcs_uri(uri) + bucket, encoded = config._parse_gcs_uri( + uri, litellm_params={"bucket_name": "litellm-local"} + ) assert bucket == "litellm-local" expected_path = ( "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" @@ -43,30 +52,141 @@ def test_should_parse_uri_with_nested_publisher_path(self, config): assert encoded == urllib.parse.quote(expected_path, safe="") def test_should_handle_url_encoded_input(self, config): - encoded_uri = urllib.parse.quote("gs://my-bucket/some/path", safe="") - bucket, encoded = config._parse_gcs_uri(encoded_uri) + encoded_uri = urllib.parse.quote( + "gs://my-bucket/litellm-vertex-files/some/path", safe="" + ) + bucket, encoded = config._parse_gcs_uri( + encoded_uri, litellm_params={"bucket_name": "my-bucket"} + ) assert bucket == "my-bucket" - assert encoded == urllib.parse.quote("some/path", safe="") + assert encoded == urllib.parse.quote("litellm-vertex-files/some/path", safe="") + + def test_should_reject_bucket_only(self, config): + with pytest.raises(ValueError, match="object name"): + config._parse_gcs_uri( + "gs://my-bucket", litellm_params={"bucket_name": "my-bucket"} + ) + + def test_should_reject_no_gs_prefix(self, config): + with pytest.raises(ValueError, match="gs://"): + config._parse_gcs_uri( + "my-bucket/litellm-vertex-files/object.txt", + litellm_params={"bucket_name": "my-bucket"}, + ) + + def test_should_reject_unmanaged_object_path(self, config): + with pytest.raises(ValueError, match="LiteLLM-managed"): + config._parse_gcs_uri( + "gs://my-bucket/private/object.txt", + litellm_params={"bucket_name": "my-bucket"}, + ) + + def test_should_reject_request_supplied_legacy_flag(self, config): + with pytest.raises(ValueError, match="LiteLLM-managed"): + config._parse_gcs_uri( + "gs://my-bucket/private/object.txt", + litellm_params={ + "bucket_name": "my-bucket", + "allow_legacy_cloud_file_ids": True, + }, + ) + + def test_should_allow_legacy_object_path_with_trusted_server_flag(self, config): + trusted_credentials = MappingProxyType({"allow_legacy_cloud_file_ids": True}) + bucket, encoded = config._parse_gcs_uri( + "gs://my-bucket/private/object.txt", + litellm_params={ + "bucket_name": "my-bucket", + "_litellm_internal_model_credentials": trusted_credentials, + }, + ) - def test_should_handle_bucket_only(self, config): - bucket, encoded = config._parse_gcs_uri("gs://my-bucket") assert bucket == "my-bucket" - assert encoded == "" + assert encoded == urllib.parse.quote("private/object.txt", safe="") + + def test_should_reject_user_supplied_legacy_flag_snapshot(self, config): + with pytest.raises(ValueError, match="LiteLLM-managed"): + config._parse_gcs_uri( + "gs://my-bucket/private/object.txt", + litellm_params={ + "bucket_name": "my-bucket", + "_litellm_internal_model_credentials": { + "allow_legacy_cloud_file_ids": True + }, + }, + ) + + def test_should_keep_configured_prefix_for_legacy_object_path(self, config): + trusted_credentials = MappingProxyType({"allow_legacy_cloud_file_ids": True}) + bucket, encoded = config._parse_gcs_uri( + "gs://my-bucket/team-a/private/object.txt", + litellm_params={ + "bucket_name": "my-bucket/team-a", + "_litellm_internal_model_credentials": trusted_credentials, + }, + ) - def test_should_handle_no_gs_prefix(self, config): - bucket, encoded = config._parse_gcs_uri("my-bucket/object.txt") assert bucket == "my-bucket" - assert encoded == "object.txt" + assert encoded == urllib.parse.quote("team-a/private/object.txt", safe="") + + def test_should_reject_legacy_object_outside_configured_prefix(self, config): + trusted_credentials = MappingProxyType({"allow_legacy_cloud_file_ids": True}) + with pytest.raises(ValueError, match="configured storage prefix"): + config._parse_gcs_uri( + "gs://my-bucket/team-b/private/object.txt", + litellm_params={ + "bucket_name": "my-bucket/team-a", + "_litellm_internal_model_credentials": trusted_credentials, + }, + ) + + def test_should_reject_unconfigured_bucket(self, config): + with pytest.raises(ValueError, match="configured storage bucket"): + config._parse_gcs_uri( + "gs://other-bucket/litellm-vertex-files/object.txt", + litellm_params={"bucket_name": "my-bucket"}, + ) + + +class TestCreateFileUrl: + def test_should_ignore_request_metadata_bucket_and_sanitize_filename(self, config): + url = config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={ + "bucket_name": "safe-bucket", + "litellm_metadata": {"gcs_bucket_name": "attacker-bucket"}, + }, + data={ + "file": ("../../owned.jsonl?alt=media", b"{}", "application/jsonl"), + "purpose": "assistants", + }, + ) + + parsed_url = urlparse(url) + object_name = parse_qs(parsed_url.query)["name"][0] + assert "/b/safe-bucket/" in parsed_url.path + assert "attacker-bucket" not in url + assert object_name.startswith("litellm-vertex-files/uploads/") + assert object_name.endswith("-owned.jsonl_alt_media") + assert ".." not in object_name + assert "?" not in object_name class TestTransformRetrieveFile: def test_should_build_correct_gcs_metadata_url(self, config): - file_id = "gs://my-bucket/path/to/file.jsonl" + file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" url, params = config.transform_retrieve_file_request( - file_id=file_id, optional_params={}, litellm_params={} + file_id=file_id, + optional_params={}, + litellm_params={"bucket_name": "my-bucket"}, + ) + expected_encoded = urllib.parse.quote( + "litellm-vertex-files/path/to/file.jsonl", safe="" ) - expected_encoded = urllib.parse.quote("path/to/file.jsonl", safe="") assert ( url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{expected_encoded}" @@ -119,13 +239,13 @@ def test_should_default_purpose_to_batch_when_metadata_missing(self, config): class TestTransformFileContent: def test_should_build_gcs_media_download_url(self, config): - file_id = "gs://my-bucket/path/to/file.jsonl" + file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" url, params = config.transform_file_content_request( file_content_request={"file_id": file_id}, optional_params={}, - litellm_params={}, + litellm_params={"bucket_name": "my-bucket"}, ) - encoded = urllib.parse.quote("path/to/file.jsonl", safe="") + encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") assert ( url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}?alt=media" @@ -254,11 +374,13 @@ def test_should_skip_batch_output_transformation_when_opt_out_flag_set( class TestTransformDeleteFile: def test_should_build_correct_gcs_delete_url(self, config): - file_id = "gs://my-bucket/path/to/file.jsonl" + file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" url, params = config.transform_delete_file_request( - file_id=file_id, optional_params={}, litellm_params={} + file_id=file_id, + optional_params={}, + litellm_params={"bucket_name": "my-bucket"}, ) - encoded = urllib.parse.quote("path/to/file.jsonl", safe="") + encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") assert ( url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}" )