From 2894b1f06faee65106e1504da61113f1e57840b4 Mon Sep 17 00:00:00 2001
From: "tao.chen" <93983997+taochen-ct@users.noreply.github.com>
Date: Thu, 6 Aug 2026 19:02:11 +0800
Subject: [PATCH] fix: storage_api.py

---
 backend/src/backend/admin.py       |  2 +-
 backend/src/backend/auth.py        | 10 ++++-
 backend/src/backend/main.py        |  8 +---
 backend/src/backend/storage_api.py | 61 ++++++++++--------------------
 common/src/common/config.py        | 11 ++++++
 5 files changed, 43 insertions(+), 49 deletions(-)

diff --git a/backend/src/backend/admin.py b/backend/src/backend/admin.py
index f0ffd5d..f9e314c 100644
--- a/backend/src/backend/admin.py
+++ b/backend/src/backend/admin.py
@@ -239,7 +239,7 @@ async def delete_employee(
         .select_from(WorkspaceMembers)
         .where(WorkspaceMembers.user_id == user_id)
     )
-    if memberships == 0:
+    if int(memberships or 0) == 0:
         user.status = "disabled"
     return {
         "request_id": context.request_id,
diff --git a/backend/src/backend/auth.py b/backend/src/backend/auth.py
index c7bc069..bd13764 100644
--- a/backend/src/backend/auth.py
+++ b/backend/src/backend/auth.py
@@ -24,6 +24,7 @@ from backend.dependencies import database_session
 from common.auth.jwt import JwtError, issue_jwt, verify_jwt_token
 from common.auth.membership import resolve_is_system_admin
 from common.auth.passwords import verify_password
+from common.config import settings
 from common.db.models import Roles, Users, WorkspaceMembers, Workspaces
 from common.ids import new_ulid
 
@@ -41,13 +42,20 @@ COOKIE_SAMESITE = "lax"
 
 def _set_session_cookie(request: Request, response: Response, token: str) -> None:
     forwarded_scheme = request.headers.get("x-forwarded-proto", request.url.scheme)
+    # Operators running behind a TLS-terminating proxy that strips
+    # X-Forwarded-Proto can opt into forcing the Secure flag via
+    # ``settings.cookie_force_secure`` — without that override a plain
+    # HTTP request (no forwarded scheme, no TLS upgrade visible to the
+    # app) would yield an insecure cookie and modern browsers would
+    # silently drop it on the HTTPS round trip.
+    secure = forwarded_scheme == "https" or settings.cookie_force_secure
     response.set_cookie(
         key=COOKIE_NAME,
         value=token,
         max_age=COOKIE_TTL_SECONDS,
         path="/",
         httponly=True,
-        secure=forwarded_scheme == "https",
+        secure=secure,
         samesite=COOKIE_SAMESITE,
     )
 
diff --git a/backend/src/backend/main.py b/backend/src/backend/main.py
index 78a8a8d..35bb910 100644
--- a/backend/src/backend/main.py
+++ b/backend/src/backend/main.py
@@ -1,11 +1,9 @@
 from __future__ import annotations
 
-import asyncio
 from contextlib import asynccontextmanager
 from typing import Any, AsyncIterator
 
 import httpx
-from fastapi.routing import APIRoute
 
 from common.config import settings
 from common.db import create_database_engine, create_session_factory
@@ -27,7 +25,7 @@ from backend.rclone_rc_client import RcloneRCClient
 from backend.schedule_runs import router as schedule_runs_router
 from backend.schedules import router as schedules_router
 from backend.scripts import router as scripts_router
-from backend.storage_api import app as storage_app
+from backend.storage_api import router as storage_api_router
 
 
 @asynccontextmanager
@@ -79,6 +77,4 @@ app.include_router(admin_router)
 app.include_router(platform_router)
 
 # Reuse the proven storage endpoints without running another FastAPI service.
-for route in storage_app.routes:
-    if isinstance(route, APIRoute) and route.path.startswith("/internal/"):
-        app.router.routes.append(route)
+app.include_router(storage_api_router, prefix="/internal")
diff --git a/backend/src/backend/storage_api.py b/backend/src/backend/storage_api.py
index 59c1ab0..622a7e9 100644
--- a/backend/src/backend/storage_api.py
+++ b/backend/src/backend/storage_api.py
@@ -1,16 +1,16 @@
 from __future__ import annotations
 
 import hashlib
-from contextlib import asynccontextmanager
 from datetime import UTC, datetime, timedelta
+from pathlib import PurePosixPath
 from typing import Any, AsyncIterator
 
-from fastapi import Depends, HTTPException, Request, status
+from fastapi import APIRouter, Depends, HTTPException, Request, status
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from common.config import settings
-from common.db import create_database_engine, create_session_factory, session_scope
+from common.db import session_scope
 from common.db.models import (
     StorageObjects,
     UploadSessions,
@@ -19,14 +19,7 @@ from common.db.models import (
     Workspaces,
 )
 from common.ids import new_ulid
-from common.service_app import create_service_app
-from common.storage import (
-    AsyncStorageBackend,
-    PURPOSE_BUCKETS,
-    actual_bucket_name,
-    build_storage_config,
-    create_storage,
-)
+from common.storage import actual_bucket_name
 from common.storage.schemas import (
     CreateUploadRequest,
     DownloadUrlRequest,
@@ -126,27 +119,13 @@ def storage_payload(item: StorageObjects) -> dict[str, Any]:
     }
 
 
-@asynccontextmanager
-async def lifespan(app: Any) -> AsyncIterator[None]:
-    engine = create_database_engine(settings.database_url)
-    app.state.session_factory = create_session_factory(engine)
-    # Buckets are pre-provisioned by the deployment; the storage layer no
-    # longer auto-creates them. ``build_storage_config`` picks s3 vs local
-    # based on ``settings.storage_backend``. The dict is keyed by the
-    # actual bucket name so ``object_stores[upload.bucket_name]``
-    # works without a reverse mapping.
-    app.state.object_stores: dict[str, AsyncStorageBackend] = {
-        actual_bucket_name(purpose): create_storage(build_storage_config(purpose))
-        for purpose in PURPOSE_BUCKETS
-    }
-    app.state.default_bucket = settings.s3_workspace_bucket
-    try:
-        yield
-    finally:
-        await engine.dispose()
-
-
-app = create_service_app(settings.service_name, lifespan=lifespan)
+# Internal routes are mounted from main.py via
+# ``include_router(router, prefix="/internal")``. The engine, session
+# factory, and object_stores live on the main app's lifespan — this
+# module only owns the route definitions and the storage-helper
+# utilities (``storage_payload``, ``resolve_bucket``, ``BUCKET_FOR_USAGE``
+# …) consumed by ``backend.services.storage``.
+router = APIRouter(tags=["internal-storage"])
 
 
 async def database_session(request: Request) -> AsyncIterator[AsyncSession]:
@@ -392,7 +371,7 @@ async def upload_bytes_to_session(
     return item
 
 
-@app.post("/internal/v1/uploads")
+@router.post("/v1/uploads")
 async def create_upload(
     payload: CreateUploadRequest,
     request: Request,
@@ -403,7 +382,7 @@ async def create_upload(
     }
 
 
-@app.put("/internal/v1/uploads/{upload_id}")
+@router.put("/v1/uploads/{upload_id}")
 async def upload_bytes(
     upload_id: str, request: Request, session: AsyncSession = Depends(database_session)
 ) -> dict[str, Any]:
@@ -415,7 +394,7 @@ async def upload_bytes(
     return {"data": storage_payload(item)}
 
 
-@app.post("/internal/v1/uploads/{upload_id}/abort")
+@router.post("/v1/uploads/{upload_id}/abort")
 async def abort_upload(
     upload_id: str, request: Request, session: AsyncSession = Depends(database_session)
 ) -> dict[str, Any]:
@@ -438,7 +417,7 @@ async def abort_upload(
     return {"data": {"upload_id": upload_id, "status": "aborted"}}
 
 
-@app.post("/internal/v1/objects")
+@router.post("/v1/objects")
 async def create_server_object(
     payload: ServerObjectRequest,
     request: Request,
@@ -447,7 +426,7 @@ async def create_server_object(
     return await create_server_object_payload(payload, request, session)
 
 
-@app.post("/internal/v1/objects/{storage_object_id}/download-url")
+@router.post("/v1/objects/{storage_object_id}/download-url")
 async def create_download_url(
     storage_object_id: str,
     payload: DownloadUrlRequest,
@@ -458,7 +437,7 @@ async def create_download_url(
     return await create_download_url_payload(item, payload, request)
 
 
-@app.delete("/internal/v1/objects/{storage_object_id}")
+@router.delete("/v1/objects/{storage_object_id}")
 async def delete_object(
     storage_object_id: str,
     request: Request,
@@ -468,7 +447,7 @@ async def delete_object(
     return await soft_delete_object(storage_object_id, request, session)
 
 
-@app.post("/internal/v1/objects/{storage_object_id}/restore")
+@router.post("/v1/objects/{storage_object_id}/restore")
 async def restore_object(
     storage_object_id: str,
     request: Request,
@@ -517,7 +496,7 @@ async def restore_object(
     }
 
 
-@app.post("/internal/v1/admin/trash/purge")
+@router.post("/v1/admin/trash/purge")
 async def purge_trash_object(
     payload: dict[str, Any],
     request: Request,
@@ -562,6 +541,6 @@ async def purge_trash_object(
     return {"data": {"storage_object_id": storage_object_id, "purged": True}}
 
 
-@app.get("/internal/health/storage")
+@router.get("/health/storage")
 async def internal_health() -> dict[str, str]:
     return {"status": "ready", "service": "storage-api"}
diff --git a/common/src/common/config.py b/common/src/common/config.py
index 4bb9995..aea61cc 100644
--- a/common/src/common/config.py
+++ b/common/src/common/config.py
@@ -42,6 +42,17 @@ class Settings(BaseSettings):
         default=False,
         description="Enable the self-hosted UI's short-lived demo session cookie.",
     )
+    cookie_force_secure: bool = Field(
+        default=False,
+        description=(
+            "Force the ``Secure`` flag on the session cookie even when the "
+            "inbound request scheme is plain HTTP. Enable this when running "
+            "behind a TLS-terminating reverse proxy that strips or rewrites "
+            "``X-Forwarded-Proto`` — without it the cookie is written without "
+            "the Secure flag and modern browsers will refuse to send it back "
+            "over HTTPS."
+        ),
+    )
 
     # ── service identity ──────────────────────────────────────────
     service_name: str = Field(