refactor: deduplicate helpers, fix security issues, add dev tooling

- Extract shared MIME types, base64 reading, and HTML response building
  into app/utils/image.py (replaces 6 duplicated blocks)
- Extract shared get_service() into app/api/dependencies.py
- Fix timing-unsafe token comparison with secrets.compare_digest
- Add SSRF protection on external URL fetches (block private/loopback IPs)
- Fix CORS: allow_credentials=False with wildcard origins
- Add asyncio.Lock to model registry to prevent concurrent duplicate loads
- Add COPY pyproject.toml to Dockerfile so version reads work in containers
- Bump Dockerfile to python:3.13-slim
- Add mise.toml, .venv support, split requirements into runtime/dev
- Add 18 new tests for image utils and SSRF protection

Author: danielcopper

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.11-slim
+FROM python:3.13-slim
 
 WORKDIR /app
 
@@ -7,6 +7,7 @@ COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
 # Copy application
+COPY pyproject.toml .
 COPY app/ ./app/
 
 # Create storage directory

app/api/dependencies.py

@@ -0,0 +1,22 @@
+from app.core.config import settings
+from app.services.gemini_service import get_gemini_service
+from app.services.litellm_service import get_litellm_service
+from app.services.openai_service import get_openai_service
+
+
+def get_service(provider: str):
+    """Get service instance based on provider."""
+    if provider == "litellm":
+        if not settings.litellm_available:
+            raise ValueError("LiteLLM not configured. Set LITELLM_BASE_URL")
+        return get_litellm_service()
+    elif provider == "openai":
+        if not settings.openai_available:
+            raise ValueError("OpenAI not configured. Set OPENAI_API_KEY")
+        return get_openai_service()
+    elif provider == "gemini":
+        if not settings.gemini_available:
+            raise ValueError("Gemini not configured. Set GEMINI_API_KEY")
+        return get_gemini_service()
+    else:
+        raise ValueError(f"Unknown provider: {provider}")

app/api/routes/edit.py

@@ -1,21 +1,16 @@
-import base64
-import contextlib
 import logging
-from pathlib import Path
 from typing import Annotated, Literal
 
 from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
-from fastapi.responses import HTMLResponse
 
+from app.api.dependencies import get_service
 from app.core.config import settings
 from app.core.security import verify_token
 from app.schemas.requests import ImageEditRequest
 from app.schemas.responses import ImageResponse
-from app.services.gemini_service import get_gemini_service
-from app.services.litellm_service import get_litellm_service
 from app.services.model_registry import model_registry
-from app.services.openai_service import get_openai_service
 from app.services.storage_service import storage_service
+from app.utils.image import build_openwebui_html_response, read_image_as_base64
 
 logger = logging.getLogger(__name__)
 
@@ -86,7 +81,7 @@ async def edit_image(
 
     # Get service based on provider
     try:
-        service = _get_service(provider)
+        service = get_service(provider)
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e)) from None
 
@@ -118,18 +113,10 @@ async def edit_image(
 
     # Handle response format
     if response_format == "base64":
-        image_filename = urls[0].split("/")[-1]
-        image_path = Path(settings.STORAGE_PATH) / image_filename
-
-        if not image_path.exists():
-            raise HTTPException(status_code=500, detail="Edited image file not found")
-
-        with open(image_path, "rb") as f:
-            image_data = base64.b64encode(f.read()).decode("utf-8")
-
-        ext = image_path.suffix.lower()
-        mime_types = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp"}
-        mime_type = mime_types.get(ext, "image/png")
+        try:
+            image_data, mime_type = await read_image_as_base64(urls[0], cleanup=False)
+        except FileNotFoundError:
+            raise HTTPException(status_code=500, detail="Edited image file not found") from None
 
         return ImageResponse(
             image_base64=image_data,
@@ -201,7 +188,7 @@ async def edit_image_json(
 
     # Get service based on provider
     try:
-        service = _get_service(request.provider)
+        service = get_service(request.provider)
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e)) from None
 
@@ -222,47 +209,18 @@ async def edit_image_json(
 
     # OpenWebUI mode: return HTML with embedded image for iframe display
     if settings.OPENWEBUI_MODE:
-        image_filename = urls[0].split("/")[-1]
-        image_path = Path(settings.STORAGE_PATH) / image_filename
-
-        if not image_path.exists():
-            raise HTTPException(status_code=500, detail="Edited image file not found")
-
-        with open(image_path, "rb") as f:
-            image_data = base64.b64encode(f.read()).decode("utf-8")
-
-        # Clean up local file if not saving locally
-        if not settings.SAVE_IMAGES_LOCALLY:
-            with contextlib.suppress(Exception):
-                image_path.unlink()
-
-        ext = image_path.suffix.lower()
-        mime_types = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp"}
-        mime_type = mime_types.get(ext, "image/png")
-
-        html = f'<img src="data:{mime_type};base64,{image_data}" style="max-width:100%; height:auto;">'
-        return HTMLResponse(content=html, headers={"Content-Disposition": "inline"})
+        try:
+            image_data, mime_type = await read_image_as_base64(urls[0])
+        except FileNotFoundError:
+            raise HTTPException(status_code=500, detail="Edited image file not found") from None
+        return build_openwebui_html_response(image_data, mime_type)
 
     # Return based on MARKDOWN_EMBED_IMAGES setting
     if settings.MARKDOWN_EMBED_IMAGES:
-        # Return markdown with embedded base64 data URI
-        image_filename = urls[0].split("/")[-1]
-        image_path = Path(settings.STORAGE_PATH) / image_filename
-
-        if not image_path.exists():
-            raise HTTPException(status_code=500, detail="Edited image file not found")
-
-        with open(image_path, "rb") as f:
-            image_data = base64.b64encode(f.read()).decode("utf-8")
-
-        # Clean up local file if not saving locally
-        if not settings.SAVE_IMAGES_LOCALLY:
-            with contextlib.suppress(Exception):
-                image_path.unlink()
-
-        ext = image_path.suffix.lower()
-        mime_types = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp"}
-        mime_type = mime_types.get(ext, "image/png")
+        try:
+            image_data, mime_type = await read_image_as_base64(urls[0])
+        except FileNotFoundError:
+            raise HTTPException(status_code=500, detail="Edited image file not found") from None
 
         markdown = f"![Edited image](data:{mime_type};base64,{image_data})"
         return ImageResponse(
@@ -294,24 +252,6 @@ async def edit_image_json(
         )
 
 
-def _get_service(provider: str):
-    """Get service instance based on provider."""
-    if provider == "litellm":
-        if not settings.litellm_available:
-            raise ValueError("LiteLLM not configured. Set LITELLM_BASE_URL")
-        return get_litellm_service()
-    elif provider == "openai":
-        if not settings.openai_available:
-            raise ValueError("OpenAI not configured. Set OPENAI_API_KEY")
-        return get_openai_service()
-    elif provider == "gemini":
-        if not settings.gemini_available:
-            raise ValueError("Gemini not configured. Set GEMINI_API_KEY")
-        return get_gemini_service()
-    else:
-        raise ValueError(f"Unknown provider: {provider}")
-
-
 def _get_default_edit_model(provider: str) -> str:
     """Get default model for editing based on provider."""
     # Check if DEFAULT_MODEL supports editing

app/api/routes/generate.py

@@ -1,19 +1,15 @@
-import base64
-import contextlib
 import logging
-from pathlib import Path
 
 from fastapi import APIRouter, Depends, HTTPException
 from fastapi.responses import HTMLResponse, StreamingResponse
 
+from app.api.dependencies import get_service
 from app.core.config import settings
 from app.core.security import verify_token
 from app.schemas.requests import ImageRequest
 from app.schemas.responses import ImageResponse
-from app.services.gemini_service import get_gemini_service
-from app.services.litellm_service import get_litellm_service
 from app.services.model_registry import model_registry
-from app.services.openai_service import get_openai_service
+from app.utils.image import build_openwebui_html_response, read_image_as_base64
 from app.utils.sse import generate_with_progress
 
 logger = logging.getLogger(__name__)
@@ -49,7 +45,7 @@ async def generate_image(request: ImageRequest, _: None = Depends(verify_token))
 
     # Get service based on provider
     try:
-        service = _get_service(request.provider)
+        service = get_service(request.provider)
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e)) from None
 
@@ -72,48 +68,18 @@ async def generate_image(request: ImageRequest, _: None = Depends(verify_token))
 
     # OpenWebUI mode: return HTML with embedded image for iframe display
     if settings.OPENWEBUI_MODE:
-        image_filename = urls[0].split("/")[-1]
-        image_path = Path(settings.STORAGE_PATH) / image_filename
-
-        if not image_path.exists():
-            raise HTTPException(status_code=500, detail="Generated image file not found")
-
-        with open(image_path, "rb") as f:
-            image_data = base64.b64encode(f.read()).decode("utf-8")
-
-        # Clean up local file if not saving locally
-        if not settings.SAVE_IMAGES_LOCALLY:
-            with contextlib.suppress(Exception):
-                image_path.unlink()
-
-        ext = image_path.suffix.lower()
-        mime_types = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp"}
-        mime_type = mime_types.get(ext, "image/png")
-
-        html = f'<img src="data:{mime_type};base64,{image_data}" style="max-width:100%; height:auto;">'
-        return HTMLResponse(content=html, headers={"Content-Disposition": "inline"})
+        try:
+            image_data, mime_type = await read_image_as_base64(urls[0])
+        except FileNotFoundError:
+            raise HTTPException(status_code=500, detail="Generated image file not found") from None
+        return build_openwebui_html_response(image_data, mime_type)
 
     # Handle response format
     if request.response_format == "base64":
-        # Extract filename from URL and read file
-        image_filename = urls[0].split("/")[-1]
-        image_path = Path(settings.STORAGE_PATH) / image_filename
-
-        if not image_path.exists():
-            raise HTTPException(status_code=500, detail="Generated image file not found")
-
-        with open(image_path, "rb") as f:
-            image_data = base64.b64encode(f.read()).decode("utf-8")
-
-        # Clean up local file if not saving locally
-        if not settings.SAVE_IMAGES_LOCALLY:
-            with contextlib.suppress(Exception):
-                image_path.unlink()
-
-        # Determine mime type from extension
-        ext = image_path.suffix.lower()
-        mime_types = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp"}
-        mime_type = mime_types.get(ext, "image/png")
+        try:
+            image_data, mime_type = await read_image_as_base64(urls[0])
+        except FileNotFoundError:
+            raise HTTPException(status_code=500, detail="Generated image file not found") from None
 
         return ImageResponse(
             image_base64=image_data,
@@ -130,25 +96,12 @@ async def generate_image(request: ImageRequest, _: None = Depends(verify_token))
 
     if request.response_format == "markdown":
         if settings.MARKDOWN_EMBED_IMAGES:
-            # Return markdown with embedded base64 data URI
-            # Open WebUI's ENABLE_CHAT_RESPONSE_BASE64_IMAGE_URL_CONVERSION will convert this
-            image_filename = urls[0].split("/")[-1]
-            image_path = Path(settings.STORAGE_PATH) / image_filename
-
-            if not image_path.exists():
-                raise HTTPException(status_code=500, detail="Generated image file not found")
-
-            with open(image_path, "rb") as f:
-                image_data = base64.b64encode(f.read()).decode("utf-8")
-
-            # Clean up local file if not saving locally
-            if not settings.SAVE_IMAGES_LOCALLY:
-                with contextlib.suppress(Exception):
-                    image_path.unlink()
-
-            ext = image_path.suffix.lower()
-            mime_types = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp"}
-            mime_type = mime_types.get(ext, "image/png")
+            try:
+                image_data, mime_type = await read_image_as_base64(urls[0])
+            except FileNotFoundError:
+                raise HTTPException(
+                    status_code=500, detail="Generated image file not found"
+                ) from None
 
             markdown = f"![Generated image](data:{mime_type};base64,{image_data})"
             return ImageResponse(
@@ -214,7 +167,7 @@ async def generate_image_stream(request: ImageRequest, _: None = Depends(verify_
 
     # Get service based on provider
     try:
-        service = _get_service(request.provider)
+        service = get_service(request.provider)
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e)) from None
 
@@ -260,7 +213,7 @@ async def generate_image_preview(request: ImageRequest, _: None = Depends(verify
 
     # Get service based on provider
     try:
-        service = _get_service(request.provider)
+        service = get_service(request.provider)
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e)) from None
 
@@ -345,26 +298,6 @@ async def generate_image_preview(request: ImageRequest, _: None = Depends(verify
     return HTMLResponse(content=html_content, headers={"Content-Disposition": "inline"})
 
 
-def _get_service(provider: str):
-    """
-    Get service instance based on provider.
-    """
-    if provider == "litellm":
-        if not settings.litellm_available:
-            raise ValueError("LiteLLM not configured. Set LITELLM_BASE_URL")
-        return get_litellm_service()
-    elif provider == "openai":
-        if not settings.openai_available:
-            raise ValueError("OpenAI not configured. Set OPENAI_API_KEY")
-        return get_openai_service()
-    elif provider == "gemini":
-        if not settings.gemini_available:
-            raise ValueError("Gemini not configured. Set GEMINI_API_KEY")
-        return get_gemini_service()
-    else:
-        raise ValueError(f"Unknown provider: {provider}")
-
-
 def _get_default_model(provider: str) -> str:
     """
     Get default model for provider based on available models.

app/core/security.py

@@ -1,3 +1,4 @@
+import secrets
 from typing import Annotated
 
 from fastapi import Depends, HTTPException, status
@@ -28,7 +29,7 @@ async def verify_token(
             headers={"WWW-Authenticate": "Bearer"},
         )
 
-    if credentials.credentials != settings.API_BEARER_TOKEN:
+    if not secrets.compare_digest(credentials.credentials, settings.API_BEARER_TOKEN):
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Invalid authentication token",

app/main.py

@@ -65,7 +65,7 @@ async def lifespan(app: FastAPI):
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_credentials=True,
+    allow_credentials=False,
     allow_methods=["*"],
     allow_headers=["*"],
 )