Commit 0e5bf740 authored by Łukasz Wałejko's avatar Łukasz Wałejko
Browse files

Enhance FastAPI uploader service: add error handling for login, improve... 

Enhance FastAPI uploader service: add error handling for login, improve session management with secure cookies, and refine path validation logic.
parent cd41bcaf

trapper/trapper-project/trapper/apps/storage/management/commands/start_uploader_server.py

+26 −16
Original line number Diff line number Diff line
@@ -4,13 +4,14 @@ from django.contrib.sessions.backends.base import SessionBase 
from django.core.management.base import BaseCommand

from django.conf import settings

from django.contrib.auth import authenticate, get_user_model

from django.core.exceptions import ValidationError



import os

import blake3

import uvicorn

import aiofiles

from fastapi import FastAPI, Request, Depends, HTTPException, APIRouter

from fastapi.middleware.cors import CORSMiddleware

from fastapi.responses import JSONResponse

from pydantic import BaseModel



User = get_user_model()
@@ -22,15 +23,6 @@ session_store: SessionBase = import_module(settings.SESSION_ENGINE).SessionStore 
app = FastAPI(title="Chunked Upload Service")

router = APIRouter(prefix="/uploader")



app.add_middleware(

    CORSMiddleware,

    allow_origins=["*"],

    allow_credentials=True,

    allow_methods=["*"],

    allow_headers=["*"],

)





# -------------- AUTH & SESSION --------------
@@ -60,9 +52,13 @@ class LoginData(BaseModel): 


@router.post("/auth/login")

async def login(data: LoginData):

    try:

        user = await sync_to_async(authenticate)(

            username=data.username, password=data.password

        )

    except ValidationError:

        raise HTTPException(401, "Invalid credentials")



    if user is None or not user.is_active:

        raise HTTPException(401, "Invalid credentials")
@@ -72,7 +68,19 @@ async def login(data: LoginData): 
        store.create()

        return {"username": user.username, "sessionid": store.session_key}



    return await sync_to_async(_create_session)(user)

    session = await sync_to_async(_create_session)(user)

    resp = JSONResponse({"username": user.username, "sessionid": session["sessionid"]})

    resp.set_cookie(

        key=settings.SESSION_COOKIE_NAME,

        value=session["sessionid"],

        httponly=True,

        secure=getattr(settings, "SESSION_COOKIE_SECURE", False),

        samesite=getattr(settings, "SESSION_COOKIE_SAMESITE", "Lax"),

        domain=getattr(settings, "SESSION_COOKIE_DOMAIN", None),

        max_age=getattr(settings, "SESSION_COOKIE_AGE", None),

        path=getattr(settings, "SESSION_COOKIE_PATH", "/"),

    )

    return resp





# -------------- UPLOAD ENDPOINTS --------------
@@ -93,12 +101,14 @@ class FinalizeRequest(BaseModel): 


def ensure_user_path(user: User, rel_path: str) -> str:

    base = os.path.join(BASE_DIR, user.username)

    target = os.path.normpath(os.path.join(base, rel_path.lstrip("/")))

    if not target.startswith(base):

    target = os.path.join(base, rel_path.lstrip("/"))

    base_real = os.path.realpath(base)

    target_real = os.path.realpath(target)

    if os.path.commonpath([base_real, target_real]) != base_real:

        raise HTTPException(403, "Invalid path")

    parent = os.path.dirname(target)

    parent = os.path.dirname(target_real)

    os.makedirs(parent, exist_ok=True)

    return target

    return target_real





@router.post("/upload/init")
@@ -169,7 +179,7 @@ async def complete_upload(req: FinalizeRequest, user=Depends(django_session_auth 
            hasher.update(block)



    async def _compute_chunk_hash(f, chunk_index):

        f.seek(chunk_index * CHUNK_SIZE)

        await f.seek(chunk_index * CHUNK_SIZE)

        hasher = blake3.blake3()

        read_chunk_size = 8 * 1024 * 1024

        remaining = CHUNK_SIZE