feat: add proxy template origin field and implement socket utilities for log streaming

backend/app/api/proxy/socket_utils.py

@@ -0,0 +1,127 @@
+import asyncio
+from pathlib import Path
+
+from fastapi import WebSocket
+
+from app.core.security import InvalidTokenError, authenticate_token
+from app.models.health import HealthSnapshot
+from app.services.mitmproxy.manager import manager
+
+
+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 serialize_health_snapshot(
+    snapshot: HealthSnapshot | None,
+) -> dict[str, object] | None:
+    if snapshot is None:
+        return None
+    return snapshot.model_dump(mode="json")
+
+
+async def read_log_snapshot(path: Path, limit: int = 8000) -> str:
+    snapshot, _offset, _marker = await snapshot_log_state(path, limit=limit)
+    return snapshot
+
+
+async def snapshot_log_state(
+    path: Path,
+    limit: int = 8000,
+    marker_width: int = 64,
+) -> tuple[str, int, bytes]:
+    if not await asyncio.to_thread(path.is_file):
+        return "", 0, b""
+
+    data = await asyncio.to_thread(path.read_bytes)
+    size = len(data)
+    marker = data[max(0, size - marker_width) : size]
+    snapshot = data[-limit:].decode("utf-8", errors="replace")
+    return snapshot, size, marker
+
+
+def _read_log_marker_bytes(path: Path, offset: int, width: int) -> bytes:
+    if offset <= 0 or not path.is_file():
+        return b""
+
+    start = max(0, offset - width)
+    with path.open("rb") as handle:
+        handle.seek(start)
+        return handle.read(offset - start)
+
+
+async def read_log_marker(path: Path, offset: int, width: int = 64) -> bytes:
+    return await asyncio.to_thread(_read_log_marker_bytes, path, offset, width)
+
+
+async def watch_socket_disconnect(websocket: WebSocket) -> None:
+    while True:
+        await websocket.receive_text()
+
+
+async def stream_proxy_logs(
+    websocket: WebSocket,
+    path: Path,
+    *,
+    offset: int,
+    marker: bytes,
+) -> None:
+    while True:
+        if not await asyncio.to_thread(path.exists):
+            await asyncio.sleep(0.4)
+            continue
+
+        try:
+            size = await asyncio.to_thread(lambda: path.stat().st_size)
+        except OSError:
+            await asyncio.sleep(0.4)
+            continue
+
+        if marker and size >= offset and await read_log_marker(path, offset) != marker:
+            snapshot, offset, marker = await snapshot_log_state(path)
+            await websocket.send_json({"type": "snapshot", "log": snapshot})
+            await asyncio.sleep(0.4)
+            continue
+
+        if size < offset:
+            snapshot, offset, marker = await snapshot_log_state(path)
+            await websocket.send_json({"type": "snapshot", "log": snapshot})
+            await asyncio.sleep(0.4)
+            continue
+
+        if size > offset:
+
+            def read_log_chunk(
+                chunk_offset: int = offset,
+                chunk_size: int = size,
+            ) -> str:
+                with path.open("rb") as handle:
+                    handle.seek(chunk_offset)
+                    return handle.read(chunk_size - chunk_offset).decode(
+                        "utf-8",
+                        errors="replace",
+                    )
+
+            chunk = await asyncio.to_thread(read_log_chunk)
+            offset = size
+            marker = await read_log_marker(path, offset)
+            if chunk:
+                await websocket.send_json({"type": "append", "log": chunk})
+
+        await asyncio.sleep(0.4)