refactor: remove Redis dependency and related code

2026-05-05 15:01:14 +02:00
parent a523c2faa3
commit d977a953cc
21 changed files with 504 additions and 158 deletions
@@ -1,8 +1,13 @@
+import asyncio
 from pathlib import Path

 from core.config import settings
-from core.security import require_authenticated_user
-from fastapi import APIRouter, Depends, HTTPException, Request, status
+from core.security import (
+    InvalidTokenError,
+    authenticate_token,
+    require_authenticated_user,
+)
+from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket, status
 from integrations.docker import get_all_services_in_dir
 from integrations.docker.discovery import service_scan_root
 from models.example import AddonExampleCreateInput
@@ -33,6 +38,7 @@ from services.mitmproxy.options import (
    validate_options,
 )
 from services.mitmproxy.validator import ProxyCodeValidationError, validate_proxy_code
+from starlette.websockets import WebSocketDisconnect

 SUPPORTED_MODES = {
    "regular",
@@ -60,8 +66,8 @@ SUPPORTED_PROTOCOLS = {
 router = APIRouter(
    prefix="/proxy",
    tags=["Proxy"],
-    dependencies=[Depends(require_authenticated_user)],
 )
+auth_dependencies = [Depends(require_authenticated_user)]


 async def _validate_code_or_raise(code: str) -> None:
@@ -117,12 +123,71 @@ async def get_proxy_or_404(proxy_port: int):
    return proxy


-@router.get("")
+async def authenticate_proxy_socket(websocket: WebSocket) -> None:
+    token = websocket.query_params.get("token")
+    if token is None:
+        authorization = websocket.headers.get("authorization", "")
+        scheme, _, value = authorization.partition(" ")
+        if scheme.lower() == "bearer" and value:
+            token = value
+
+    if not token:
+        raise InvalidTokenError
+
+    await authenticate_token(token)
+
+
+def proxy_log_path(proxy) -> Path:
+    if proxy.log_path:
+        return Path(proxy.log_path)
+    return manager.log_path_for(proxy)
+
+
+def read_log_snapshot(path: Path, limit: int = 8000) -> str:
+    if not path.is_file():
+        return ""
+    return path.read_bytes()[-limit:].decode("utf-8", errors="replace")
+
+
+async def watch_socket_disconnect(websocket: WebSocket) -> None:
+    while True:
+        await websocket.receive_text()
+
+
+async def stream_proxy_logs(websocket: WebSocket, path: Path) -> None:
+    offset = path.stat().st_size if path.is_file() else 0
+
+    while True:
+        if not path.exists():
+            await asyncio.sleep(0.4)
+            continue
+
+        try:
+            size = path.stat().st_size
+        except OSError:
+            await asyncio.sleep(0.4)
+            continue
+
+        if size < offset:
+            offset = size
+
+        if size > offset:
+            with path.open("rb") as handle:
+                handle.seek(offset)
+                chunk = handle.read(size - offset).decode("utf-8", errors="replace")
+            offset = size
+            if chunk:
+                await websocket.send_json({"type": "append", "log": chunk})
+
+        await asyncio.sleep(0.4)
+
+
+@router.get("", dependencies=auth_dependencies)
 async def list_proxies_handler():
    return await list_proxy_documents()


-@router.get("/services")
+@router.get("/services", dependencies=auth_dependencies)
 async def list_services(request: Request, scan_root: str | None = None):
    _ = settings
    root = Path(scan_root).expanduser() if scan_root else service_scan_root()
@@ -131,17 +196,21 @@ async def list_services(request: Request, scan_root: str | None = None):
    return {"scan_root": str(root), "services": services}


-@router.get("/mitmproxy/options")
+@router.get("/mitmproxy/options", dependencies=auth_dependencies)
 async def list_mitmproxy_options():
    return public_options(await load_options())


-@router.get("/mitmproxy/examples")
+@router.get("/mitmproxy/examples", dependencies=auth_dependencies)
 async def list_mitmproxy_examples():
    return [*BUILTIN_EXAMPLES, *await list_custom_examples()]


-@router.post("/mitmproxy/examples", status_code=status.HTTP_201_CREATED)
+@router.post(
+    "/mitmproxy/examples",
+    status_code=status.HTTP_201_CREATED,
+    dependencies=auth_dependencies,
+)
 async def create_mitmproxy_example(payload: AddonExampleCreateInput):
    await _validate_code_or_raise(payload.code)
    try:
@@ -156,6 +225,7 @@ async def create_mitmproxy_example(payload: AddonExampleCreateInput):
@router.delete(
    "/mitmproxy/examples/{example_id}",
    status_code=status.HTTP_204_NO_CONTENT,
+    dependencies=auth_dependencies,
 )
 async def delete_mitmproxy_example(example_id: str):
    if any(example.id == example_id for example in BUILTIN_EXAMPLES):
@@ -172,12 +242,12 @@ async def delete_mitmproxy_example(example_id: str):
        )


-@router.get("/{proxy_port}")
+@router.get("/{proxy_port}", dependencies=auth_dependencies)
 async def get_proxy(proxy_port: int):
    return await get_proxy_or_404(proxy_port)


-@router.post("", status_code=status.HTTP_201_CREATED)
+@router.post("", status_code=status.HTTP_201_CREATED, dependencies=auth_dependencies)
 async def create_proxy_handler(payload: ProxyCreateInput):
    existing_proxy = await get_proxy_by_port(payload.proxy_port)
    if existing_proxy is not None:
@@ -195,7 +265,7 @@ async def create_proxy_handler(payload: ProxyCreateInput):
    return await create_proxy(payload)


-@router.put("/{proxy_port}")
+@router.put("/{proxy_port}", dependencies=auth_dependencies)
 async def edit_proxy(proxy_port: int, payload: ProxyUpdateInput):
    if payload.code is not None:
        await _validate_code_or_raise(payload.code)
@@ -217,7 +287,7 @@ async def edit_proxy(proxy_port: int, payload: ProxyUpdateInput):
    return updated_proxy


-@router.post("/{proxy_port}/start")
+@router.post("/{proxy_port}/start", dependencies=auth_dependencies)
 async def start_proxy(proxy_port: int):
    proxy = await get_proxy_or_404(proxy_port)

@@ -233,14 +303,14 @@ async def start_proxy(proxy_port: int):
    return await update_proxy_runtime(proxy_port, runtime)


-@router.post("/{proxy_port}/stop")
+@router.post("/{proxy_port}/stop", dependencies=auth_dependencies)
 async def stop_proxy(proxy_port: int):
    proxy = await get_proxy_or_404(proxy_port)
    runtime = await manager.stop(proxy)
    return await update_proxy_runtime(proxy_port, runtime)


-@router.post("/{proxy_port}/restart")
+@router.post("/{proxy_port}/restart", dependencies=auth_dependencies)
 async def restart_proxy(proxy_port: int):
    proxy = await get_proxy_or_404(proxy_port)
    stopped = await manager.stop(proxy)
@@ -256,7 +326,7 @@ async def restart_proxy(proxy_port: int):
    return await update_proxy_runtime(proxy_port, runtime)


-@router.get("/{proxy_port}/status")
+@router.get("/{proxy_port}/status", dependencies=auth_dependencies)
 async def proxy_status(proxy_port: int):
    proxy = await get_proxy_or_404(proxy_port)
    runtime_status = await manager.status(proxy)
@@ -268,29 +338,81 @@ async def proxy_status(proxy_port: int):
    return {"proxy_port": proxy.proxy_port, "status": runtime_status, "pid": proxy.pid}


-@router.get("/{proxy_port}/health")
+@router.get("/{proxy_port}/health", dependencies=auth_dependencies)
 async def proxy_health(proxy_port: int):
    proxy = await get_proxy_or_404(proxy_port)
    return await check_proxy(proxy)


-@router.get("/{proxy_port}/logs")
+@router.get("/{proxy_port}/logs", dependencies=auth_dependencies)
 async def proxy_logs(proxy_port: int):
    proxy = await get_proxy_or_404(proxy_port)
-    if proxy.log_path is None:
-        return {"proxy_port": proxy.proxy_port, "log": ""}
-
-    path = Path(proxy.log_path)
-    if not path.is_file():
-        return {"proxy_port": proxy.proxy_port, "log": ""}
+    path = proxy_log_path(proxy)

    return {
        "proxy_port": proxy.proxy_port,
-        "log": path.read_text(encoding="utf-8", errors="replace")[-8000:],
+        "log": read_log_snapshot(path),
    }


-@router.delete("/{proxy_port}", status_code=status.HTTP_204_NO_CONTENT)
+@router.websocket("/{proxy_port}/logs/ws")
+async def proxy_logs_socket(websocket: WebSocket, proxy_port: int):
+    disconnect_task: asyncio.Task[None] | None = None
+    stream_task: asyncio.Task[None] | None = None
+
+    try:
+        try:
+            await authenticate_proxy_socket(websocket)
+        except InvalidTokenError:
+            await websocket.close(code=4401)
+            return
+
+        proxy = await get_proxy_by_port(proxy_port)
+        if proxy is None:
+            await websocket.close(code=4404)
+            return
+
+        path = proxy_log_path(proxy)
+        await websocket.accept()
+        await websocket.send_json({"type": "snapshot", "log": read_log_snapshot(path)})
+
+        disconnect_task = asyncio.create_task(watch_socket_disconnect(websocket))
+        stream_task = asyncio.create_task(stream_proxy_logs(websocket, path))
+
+        done, pending = await asyncio.wait(
+            {disconnect_task, stream_task},
+            return_when=asyncio.FIRST_EXCEPTION,
+        )
+
+        for task in pending:
+            task.cancel()
+        if pending:
+            await asyncio.gather(*pending, return_exceptions=True)
+
+        for task in done:
+            try:
+                task.result()
+            except asyncio.CancelledError:
+                return
+            except WebSocketDisconnect:
+                return
+            except RuntimeError:
+                await websocket.close(code=1011, reason="log stream failed")
+                return
+    except asyncio.CancelledError:
+        return
+    finally:
+        for task in (disconnect_task, stream_task):
+            if task is None or task.done():
+                continue
+            task.cancel()
+
+
+@router.delete(
+    "/{proxy_port}",
+    status_code=status.HTTP_204_NO_CONTENT,
+    dependencies=auth_dependencies,
+)
 async def remove_proxy(proxy_port: int):
    proxy = await get_proxy_by_port(proxy_port)
    if proxy is None:
@@ -12,7 +12,6 @@ class Settings(BaseSettings):

    MONGODB_URI: str
    MONGODB_DB_NAME: str
-    REDIS_URL: str
    JWT_SECRET_KEY: str
    JWT_ALGORITHM: str = "HS256"
    JWT_EXPIRE_MINUTES: int = Field(default=720, gt=0)
@@ -1,22 +0,0 @@
-import redis.asyncio as aioredis
-from core.config import settings
-
-_redis_pool: aioredis.Redis | None = None
-
-
-async def init_redis() -> None:
-    global _redis_pool
-    _redis_pool = aioredis.from_url(
-        settings.REDIS_URL, encoding="utf-8", decode_responses=True
-    )
-
-
-async def close_redis() -> None:
-    if _redis_pool:
-        await _redis_pool.aclose()
-
-
-async def get_redis() -> aioredis.Redis:
-    if _redis_pool is None:
-        raise RuntimeError("Redis pool has not been initialized")
-    return _redis_pool
@@ -7,7 +7,6 @@ from api.editor.router import router as editor_router
 from api.proxy.router import router as proxy_router
 from core.config import settings
 from core.mongodb import close_mongodb, init_mongodb
-from core.redis import close_redis, init_redis
 from core.security import hash_password
 from fastapi import FastAPI
 from integrations.docker import get_all_services_in_dir
@@ -21,7 +20,6 @@ from services.mitmproxy.recovery import reconcile_runtime
@asynccontextmanager
 async def lifespan(_: FastAPI):
    await init_mongodb()
-    await init_redis()
    await ensure_auth_indexes()
    await ensure_single_account(
        username=settings.AUTH_USERNAME,
@@ -56,7 +54,6 @@ async def lifespan(_: FastAPI):
    try:
        yield
    finally:
-        await close_redis()
        await close_mongodb()