Harden web UI security

1 files changed, 40 insertions, 11 deletions

diff --git a/src/ollama_web/routes/sessions.py b/src/ollama_web/routes/sessions.py
index 8640ec3..7b01705 100644
--- a/src/ollama_web/routes/sessions.py
+++ b/src/ollama_web/routes/sessions.py
@@ -2,20 +2,24 @@
 
 from __future__ import annotations
 
+import json
 import logging
-from typing import Any
+from typing import Any, cast
 
+from starlette.datastructures import UploadFile
 from starlette.requests import Request
 from starlette.responses import JSONResponse
 
-from ..sessions import SessionStore, _IMAGE_MIMES, _IMAGE_SUFFIXES
-from ..tools.pdf import extract_pdf_text
+from ..sessions import _IMAGE_MIMES, _IMAGE_SUFFIXES, SessionStore, is_valid_id
+from ..tools.helper.pdf import extract_pdf_text
 
 logger = logging.getLogger("ollama_web.sessions")
 
+_UPLOAD_CHUNK_SIZE = 1024 * 1024
+
 
 def _get_store(request: Request) -> SessionStore:
-    return request.app.state.session_store
+    return cast(SessionStore, request.app.state.session_store)
 
 
 def _is_image_file(name: str, mime: str) -> bool:
@@ -36,7 +40,8 @@ async def list_sessions(request: Request) -> JSONResponse:
 
 async def create_session(request: Request) -> JSONResponse:
     store = _get_store(request)
-    body = await request.json() if await request.body() else {}
+    raw = await request.body()
+    body = json.loads(raw.decode("utf-8")) if raw else {}
     title = body.get("title")
     session = store.create(title=title)
     return JSONResponse(session)
@@ -45,6 +50,8 @@ async def create_session(request: Request) -> JSONResponse:
 async def get_session(request: Request) -> JSONResponse:
     store = _get_store(request)
     session_id = request.path_params["session_id"]
+    if not is_valid_id(session_id):
+        return JSONResponse({"error": "Invalid session id"}, status_code=400)
     session = store.get(session_id)
     if session is None:
         return JSONResponse({"error": "Session not found"}, status_code=404)
@@ -54,6 +61,8 @@ async def get_session(request: Request) -> JSONResponse:
 async def delete_session(request: Request) -> JSONResponse:
     store = _get_store(request)
     session_id = request.path_params["session_id"]
+    if not is_valid_id(session_id):
+        return JSONResponse({"error": "Invalid session id"}, status_code=400)
     ok = store.delete(session_id)
     if not ok:
         return JSONResponse({"error": "Session not found"}, status_code=404)
@@ -63,6 +72,8 @@ async def delete_session(request: Request) -> JSONResponse:
 async def upload_file(request: Request) -> JSONResponse:
     store = _get_store(request)
     session_id = request.path_params["session_id"]
+    if not is_valid_id(session_id):
+        return JSONResponse({"error": "Invalid session id"}, status_code=400)
     session = store.get(session_id)
     if session is None:
         return JSONResponse({"error": "Session not found"}, status_code=404)
@@ -79,22 +90,38 @@ async def upload_file(request: Request) -> JSONResponse:
     for upload in files:
         if upload is None:
             continue
+        if not hasattr(upload, "read"):
+            results.append({"name": "unknown", "error": "Invalid upload"})
+            continue
+        upload_file = cast(UploadFile, upload)
         try:
-            data = await upload.read()
+            chunks: list[bytes] = []
+            total = 0
+            while True:
+                chunk = await upload_file.read(_UPLOAD_CHUNK_SIZE)
+                if not chunk:
+                    break
+                total += len(chunk)
+                if total > max_size:
+                    raise ValueError(
+                        f"File exceeds {request.app.state.settings.max_upload_mb} MB limit"
+                    )
+                chunks.append(chunk)
+            data = b"".join(chunks)
         except Exception as exc:  # noqa: BLE001
-            results.append({"name": getattr(upload, "filename", "unknown"), "error": str(exc)})
+            results.append({"name": getattr(upload_file, "filename", "unknown"), "error": str(exc)})
             continue
         if len(data) > max_size:
             results.append(
                 {
-                    "name": getattr(upload, "filename", "unknown"),
+                    "name": getattr(upload_file, "filename", "unknown"),
                     "error": f"File exceeds {request.app.state.settings.max_upload_mb} MB limit",
                 }
             )
             continue
 
-        name = getattr(upload, "filename", "uploaded-file") or "uploaded-file"
-        content_type = getattr(upload, "content_type", None) or "application/octet-stream"
+        name = getattr(upload_file, "filename", "uploaded-file") or "uploaded-file"
+        content_type = getattr(upload_file, "content_type", None) or "application/octet-stream"
         text = ""
         if _is_image_file(name, content_type):
             # Images are passed to ollama via Message.images (base64), not as text.
@@ -128,7 +155,9 @@ async def delete_file(request: Request) -> JSONResponse:
     store = _get_store(request)
     session_id = request.path_params["session_id"]
     file_id = request.path_params["file_id"]
+    if not is_valid_id(session_id) or not is_valid_id(file_id):
+        return JSONResponse({"error": "Invalid id"}, status_code=400)
     ok = store.remove_file(session_id, file_id)
     if not ok:
         return JSONResponse({"error": "File not found"}, status_code=404)
-    return JSONResponse({"deleted": True})
\ No newline at end of file
+    return JSONResponse({"deleted": True})