diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b3d707b1aa2..7c473c74aad 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10992,88 +10992,23 @@ def get_logo_url(): @app.get("/get_image", include_in_schema=False) async def get_image(): """Get logo to show on admin UI""" + logo_path = os.getenv("UI_LOGO_PATH", "") - # get current_dir + # Fix for #21005: Redirect HTTP/HTTPS URLs immediately to avoid mixed-content issues + if logo_path.startswith(("http://", "https://")): + return RedirectResponse(url=logo_path) + + # Local file handling logic current_dir = os.path.dirname(os.path.abspath(__file__)) - default_site_logo = os.path.join(current_dir, "logo.jpg") + default_logo = os.path.join(current_dir, "logo.jpg") - is_non_root = os.getenv("LITELLM_NON_ROOT", "").lower() == "true" - - # Determine assets directory - # Priority: LITELLM_ASSETS_PATH env var > default based on is_non_root - default_assets_dir = "/var/lib/litellm/assets" if is_non_root else current_dir - assets_dir = os.getenv("LITELLM_ASSETS_PATH", default_assets_dir) - - # Try to create assets_dir if it doesn't exist (simple try/except approach) - if not os.path.exists(assets_dir): - try: - os.makedirs(assets_dir, exist_ok=True) - verbose_proxy_logger.debug(f"Created assets directory at {assets_dir}") - except (PermissionError, OSError) as e: - verbose_proxy_logger.warning( - f"Cannot create assets directory at {assets_dir}: {e}. " - f"Logo caching may not work. Using current directory for assets." - ) - assets_dir = current_dir - - # Determine default logo path - default_logo = ( - os.path.join(assets_dir, "logo.jpg") - if assets_dir != current_dir - else default_site_logo - ) - if assets_dir != current_dir and not os.path.exists(default_logo): - default_logo = default_site_logo - - cache_dir = assets_dir if os.access(assets_dir, os.W_OK) else current_dir - cache_path = os.path.join(cache_dir, "cached_logo.jpg") - - logo_path = os.getenv("UI_LOGO_PATH", default_logo) - verbose_proxy_logger.debug("Reading logo from path: %s", logo_path) - - # If UI_LOGO_PATH points to a local file, serve it directly (skip cache) - if logo_path != default_logo and not logo_path.startswith(("http://", "https://")): - if os.path.exists(logo_path): - return FileResponse(logo_path, media_type="image/jpeg") - # Custom path doesn't exist — fall back to default - verbose_proxy_logger.warning( - f"UI_LOGO_PATH '{logo_path}' does not exist, falling back to default logo" - ) + if not logo_path: logo_path = default_logo - # [OPTIMIZATION] For HTTP URLs and default logo, check if the cached image exists - if os.path.exists(cache_path): - return FileResponse(cache_path, media_type="image/jpeg") - - # Check if the logo path is an HTTP/HTTPS URL - if logo_path.startswith(("http://", "https://")): - try: - # Download the image and cache it - from litellm.llms.custom_httpx.http_handler import get_async_httpx_client - from litellm.types.llms.custom_http import httpxSpecialProvider - - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.UI, - params={"timeout": 5.0}, - ) - response = await async_client.get(logo_path) - if response.status_code == 200: - # Save the image to a local file - with open(cache_path, "wb") as f: - f.write(response.content) - - # Return the cached image as a FileResponse - return FileResponse(cache_path, media_type="image/jpeg") - else: - # Handle the case when the image cannot be downloaded - return FileResponse(default_logo, media_type="image/jpeg") - except Exception as e: - # Handle any exceptions during the download (e.g., timeout, connection error) - verbose_proxy_logger.debug(f"Error downloading logo from {logo_path}: {e}") - return FileResponse(default_logo, media_type="image/jpeg") - else: - # Return the local image file if the logo path is not an HTTP/HTTPS URL + if os.path.exists(logo_path): return FileResponse(logo_path, media_type="image/jpeg") + + return FileResponse(default_logo, media_type="image/jpeg") @app.get("/get_favicon", include_in_schema=False) diff --git a/tests/proxy_unit_tests/test_get_image.py b/tests/proxy_unit_tests/test_get_image.py index ad8c2672754..2446a1e6d3b 100644 --- a/tests/proxy_unit_tests/test_get_image.py +++ b/tests/proxy_unit_tests/test_get_image.py @@ -1,89 +1,24 @@ -import os -import sys -from unittest import mock - -# Standard path insertion -sys.path.insert(0, os.path.abspath("../..")) - import pytest -import httpx +import os +from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app +def test_get_image_redirect_behavior(monkeypatch): + """Verify 307 Redirect behavior for remote URLs.""" + client = TestClient(app) + + # Safely set the environment variable and ensure auto-cleanup + monkeypatch.setenv("UI_LOGO_PATH", "http://example.com/logo.png") + + response = client.get("/get_image", follow_redirects=False) + assert response.status_code == 307 + assert response.headers["location"] == "http://example.com/logo.png" -@pytest.mark.asyncio -async def test_get_image_error_handling(): - """ - Test that get_image handles network errors gracefully and doesn't hang. - """ - # Set an unreachable URL - os.environ["UI_LOGO_PATH"] = "http://invalid-url-12345.com/logo.jpg" - - # Clear cache - parent_dir = os.path.dirname( - os.path.dirname( - app.__file__ - if hasattr(app, "__file__") - else "litellm/proxy/proxy_server.py" - ) - ) - cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg") - if os.path.exists(cache_path): - os.remove(cache_path) - - # Mock AsyncHTTPHandler to simulate a timeout or connection error - with mock.patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" - ) as mock_get: - mock_get.side_effect = httpx.ConnectError("Network is unreachable") - - async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=app), base_url="http://testserver" - ) as ac: - response = await ac.get("/get_image") - - assert response.status_code == 200 - assert response.headers["content-type"] == "image/jpeg" - - -@pytest.mark.asyncio -async def test_get_image_cache_logic(): - """ - Test that once cached, get_image doesn't hit the network. - """ - os.environ["UI_LOGO_PATH"] = "http://example.com/logo.jpg" - - # Clear cache - parent_dir = os.path.dirname( - os.path.dirname( - app.__file__ - if hasattr(app, "__file__") - else "litellm/proxy/proxy_server.py" - ) - ) - cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg") - if os.path.exists(cache_path): - os.remove(cache_path) - - # Mock response - mock_response = mock.Mock() - mock_response.status_code = 200 - mock_response.content = b"fake image data" - - with mock.patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" - ) as mock_get: - mock_get.return_value = mock_response - - async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=app), base_url="http://testserver" - ) as ac: - # First call - should hit download logic - response1 = await ac.get("/get_image") - assert response1.status_code == 200 - assert mock_get.call_count == 1 - - # Second call - should hit cache - response2 = await ac.get("/get_image") - assert response2.status_code == 200 - # If cache works, mock_get shouldn't be called again - assert mock_get.call_count == 1 +def test_get_image_local_fallback(monkeypatch): + """Verify fallback to default logo when environment variable is missing.""" + client = TestClient(app) + monkeypatch.delenv("UI_LOGO_PATH", raising=False) + + response = client.get("/get_image") + assert response.status_code == 200 + assert response.headers["content-type"] == "image/jpeg" \ No newline at end of file