fix: friendly 502 handling, 2 uvicorn workers, file bundle cache (0.20.18)

backend/app/services/rdp_bundle_delivery.py

@@ -1,26 +1,19 @@
-"""Short-lived in-memory RDP agent bundles for WinRM clients to download."""
+"""Short-lived RDP agent bundles for WinRM clients to download."""
 
 from __future__ import annotations
 
 import io
+import re
 import secrets
 import time
 import zipfile
 from pathlib import Path
-from threading import Lock
 
-TTL_SECONDS = 600
 _UTF8_BOM = b"\xef\xbb\xbf"
 _BOM_SUFFIXES = {".ps1", ".txt"}
-
-_lock = Lock()
-_entries: dict[str, tuple[bytes, float]] = {}
-
-
-def _prune_expired(now: float) -> None:
-    expired = [token for token, (_, expires_at) in _entries.items() if expires_at <= now]
-    for token in expired:
-        _entries.pop(token, None)
+_TOKEN_RE = re.compile(r"^[A-Za-z0-9_-]{16,128}$")
+TTL_SECONDS = 600
+BUNDLE_DIR = Path("/tmp/sac-rdp-bundles")
 
 
 def _bundle_file_bytes(path: Path) -> bytes:
@@ -43,27 +36,34 @@ def build_rdp_bundle_zip(repo_dir: Path, filenames: tuple[str, ...]) -> bytes:
     return buffer.getvalue()
 
 
+def _prune_expired_bundles(now: float | None = None) -> None:
+    if not BUNDLE_DIR.is_dir():
+        return
+    cutoff = (now or time.time()) - TTL_SECONDS
+    for path in BUNDLE_DIR.glob("*.zip"):
+        try:
+            if path.stat().st_mtime < cutoff:
+                path.unlink(missing_ok=True)
+        except OSError:
+            continue
+
+
 def register_rdp_bundle_zip(zip_bytes: bytes) -> str:
     token = secrets.token_urlsafe(32)
-    expires_at = time.time() + TTL_SECONDS
-    with _lock:
-        _prune_expired(time.time())
-        _entries[token] = (zip_bytes, expires_at)
+    BUNDLE_DIR.mkdir(parents=True, exist_ok=True)
+    _prune_expired_bundles()
+    (BUNDLE_DIR / f"{token}.zip").write_bytes(zip_bytes)
     return token
 
 
 def get_rdp_bundle_zip(token: str) -> bytes | None:
     text = (token or "").strip()
-    if not text:
+    if not _TOKEN_RE.fullmatch(text):
         return None
-    now = time.time()
-    with _lock:
-        _prune_expired(now)
-        entry = _entries.get(text)
-        if entry is None:
-            return None
-        data, expires_at = entry
-        if expires_at <= now:
-            _entries.pop(text, None)
-            return None
-        return data
+    path = BUNDLE_DIR / f"{text}.zip"
+    if not path.is_file():
+        return None
+    if time.time() - path.stat().st_mtime > TTL_SECONDS:
+        path.unlink(missing_ok=True)
+        return None
+    return path.read_bytes()