Просмотр исходного кода

refactor: serve.py 改用 starlette + uvicorn,统一支持 HTTP/WebSocket/API 代理

Hermes Agent 1 неделя назад
Родитель
Сommit
cbde341d8f
3 измененных файлов с 91 добавлено и 105 удалено
  1. 3 0
      frontend/requirements.txt
  2. 86 101
      frontend/serve.py
  3. 2 4
      install.sh

+ 3 - 0
frontend/requirements.txt

@@ -1 +1,4 @@
+starlette>=0.27.0
+uvicorn>=0.20.0
+httpx>=0.24.0
 websockets>=10.0

+ 86 - 101
frontend/serve.py

@@ -1,130 +1,115 @@
-"""Simple static file server with API and WebSocket proxy"""
-import asyncio
-import http.server
-import websockets
-import threading
-import urllib.request
-import urllib.error
+"""KVM Frontend - Static files + API proxy + WebSocket VNC proxy"""
 import os
-import json
+import re
 
+from starlette.applications import Starlette
+from starlette.routing import Route, WebSocketRoute, Mount
+from starlette.staticfiles import StaticFiles
+from starlette.responses import FileResponse, RedirectResponse
+from starlette.websockets import WebSocket
+import uvicorn
+
+HOST = "0.0.0.0"
 PORT = 8006
 DIST_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "dist")
 API_BASE = "http://127.0.0.1:8004"
 WS_BASE = "ws://127.0.0.1:8004"
 
 
-class KVMHandler(http.server.SimpleHTTPRequestHandler):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, directory=DIST_DIR, **kwargs)
-
-    def do_GET(self):
-        if self.path.startswith("/api/"):
-            self._proxy("GET")
-        elif self.path == "/" or self.path == "":
-            self.path = "/index.html"
-            super().do_GET()
-        else:
-            file_path = os.path.join(DIST_DIR, self.path.lstrip("/"))
-            if os.path.isfile(file_path):
-                super().do_GET()
-            else:
-                self.path = "/index.html"
-                super().do_GET()
-
-    def do_POST(self):
-        if self.path.startswith("/api/"):
-            self._proxy("POST")
-        else:
-            self.send_error(404)
-
-    def do_PUT(self):
-        if self.path.startswith("/api/"):
-            self._proxy("PUT")
-        else:
-            self.send_error(404)
-
-    def do_DELETE(self):
-        if self.path.startswith("/api/"):
-            self._proxy("DELETE")
-        else:
-            self.send_error(404)
-
-    def _proxy(self, method):
-        content_length = int(self.headers.get("Content-Length", 0))
-        body = self.rfile.read(content_length) if content_length > 0 else None
-
-        url = f"{API_BASE}{self.path}"
-        req = urllib.request.Request(url, data=body, method=method)
-
-        for key in ["Content-Type", "Authorization"]:
-            if key in self.headers:
-                req.add_header(key, self.headers[key])
+# ── SPA fallback ──────────────────────────────────────────────
+async def serve_spa(request):
+    path = request.url.path
+    file_path = os.path.join(DIST_DIR, path.lstrip("/"))
+    if os.path.isfile(file_path):
+        return FileResponse(file_path)
+    return FileResponse(os.path.join(DIST_DIR, "index.html"))
+
+
+# ── API proxy ─────────────────────────────────────────────────
+async def proxy_api(request):
+    from starlette.requests import Request
+    body = await request.body()
+    headers = {}
+    for key in ["content-type", "authorization"]:
+        val = request.headers.get(key)
+        if val:
+            headers[key.replace("_", "-")] = val
 
+    url = f"{API_BASE}{request.url.path}"
+    if request.url.query:
+        url += f"?{request.url.query}"
+
+    async with __import__("httpx").AsyncClient(timeout=30) as client:
         try:
-            with urllib.request.urlopen(req, timeout=30) as resp:
-                resp_body = resp.read()
-                self.send_response(resp.status)
-                self.send_header("Content-Type", "application/json")
-                self.send_header("Access-Control-Allow-Origin", "*")
-                self.end_headers()
-                self.wfile.write(resp_body)
-        except urllib.error.HTTPError as e:
-            resp_body = e.read()
-            self.send_response(e.code)
-            self.send_header("Content-Type", "application/json")
-            self.send_header("Access-Control-Allow-Origin", "*")
-            self.end_headers()
-            self.wfile.write(resp_body)
+            resp = await client.request(
+                method=request.method,
+                url=url,
+                content=body,
+                headers=headers,
+            )
+            return Response(
+                resp.content,
+                status_code=resp.status_code,
+                headers={"access-control-allow-origin": "*"},
+            )
         except Exception as e:
-            self.send_response(502)
-            self.send_header("Content-Type", "application/json")
-            self.end_headers()
-            self.wfile.write(json.dumps({"error": str(e)}).encode())
+            from starlette.responses import JSONResponse
+            return JSONResponse({"error": str(e)}, status_code=502)
+    from starlette.responses import Response
+    return Response("", status_code=500)
+
 
-    def log_message(self, format, *args):
-        pass
+# ── WebSocket VNC proxy ────────────────────────────────────────
+async def vnc_websocket(websocket: WebSocket):
+    await websocket.accept()
 
+    # Build backend WebSocket URL
+    path = websocket.url.path  # e.g. /ws/vnc/vm-name?host_id=local
+    backend_url = f"{WS_BASE}{path}"
 
-async def websocket_proxy(websocket, path):
-    """Proxy WebSocket connections to backend"""
-    ws_url = f"{WS_BASE}{path}"
     try:
-        async with websockets.connect(ws_url) as backend_ws:
+        async with __import__("websockets").connect(backend_url) as backend:
             async def forward_to_backend():
-                async for msg in websocket:
-                    await backend_ws.send(msg)
+                async for msg in websocket.iter_text():
+                    await backend.send(msg)
+                async for msg in websocket.iter_bytes():
+                    await backend.send(msg)
 
             async def forward_to_frontend():
-                async for msg in backend_ws:
-                    await websocket.send(msg)
+                async for msg in backend:
+                    if isinstance(msg, bytes):
+                        await websocket.send_bytes(msg)
+                    else:
+                        await websocket.send_text(msg)
 
-            await asyncio.gather(
+            await __import__("asyncio").gather(
                 forward_to_backend(),
-                forward_to_frontend()
+                forward_to_frontend(),
             )
     except Exception as e:
         print(f"WebSocket proxy error: {e}")
+    finally:
+        try:
+            await websocket.close()
+        except Exception:
+            pass
 
 
-async def ws_handler(websocket, path):
-    await websocket_proxy(websocket, path)
-
+from starlette.responses import Response
 
-def run_websocket_server():
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    loop.run_until_complete(websockets.serve(ws_handler, "0.0.0.0", PORT))
-    loop.run_forever()
+routes = [
+    WebSocketRoute("/ws/{path:path}", vnc_websocket, name="vnc"),
+    Route("/api/{path:path}", proxy_api),
+    Route("/api", proxy_api),
+    Mount("/", StaticFiles(directory=DIST_DIR, html=True), name="static"),
+]
 
+app = Starlette(routes=routes)
 
 if __name__ == "__main__":
-    # Start WebSocket server in background thread
-    ws_thread = threading.Thread(target=run_websocket_server, daemon=True)
-    ws_thread.start()
-    print(f"WebSocket proxy listening on ws://0.0.0.0:{PORT}")
-
-    # Start HTTP server
-    server = http.server.HTTPServer(("0.0.0.0", PORT), KVMHandler)
-    print(f"Static file server serving on http://0.0.0.0:{PORT}")
-    server.serve_forever()
+    uvicorn.run(
+        app,
+        host=HOST,
+        port=PORT,
+        log_level="info",
+    )

+ 2 - 4
install.sh

@@ -100,10 +100,8 @@ cd "$FRONTEND_DIR"
 npm install
 npm run build
 
-# 安装前端 Python 依赖 (serve.py 需要 websockets)
-if ! pip show websockets &>/dev/null; then
-    pip install -r requirements.txt
-fi
+# 安装前端 Python 依赖 (starlette + uvicorn)
+pip install -r requirements.txt
 
 echo -e "${YELLOW}[5/5] 验证前端...${NC}"
 if [[ -d "dist" ]]; then