Progress
This commit is contained in:
Jonatan Rek
165
app.py
165
app.py
@@ -1,148 +1,35 @@
|
||||
# main.py
|
||||
import asyncio
|
||||
from fastapi import FastAPI, Request, Response, status
|
||||
from fastapi.responses import JSONResponse
|
||||
import httpx
|
||||
import yaml
|
||||
import docker
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.types import ASGIApp
|
||||
from threading import Thread
|
||||
from time import time, sleep
|
||||
import websockets
|
||||
from starlette.types import Receive, Scope, Send
|
||||
from starlette.websockets import WebSocket
|
||||
from starlette.requests import HTTPConnection
|
||||
|
||||
from config import load_config
|
||||
from docker_manager import DockerManager
|
||||
from middleware import ProxyMiddleware
|
||||
|
||||
|
||||
# --- Configuration Loading ---
|
||||
class Route:
|
||||
def __init__(self, path_prefix: str, target: str, container: str | None = None):
|
||||
self.path_prefix = path_prefix
|
||||
self.target = target
|
||||
self.container = container
|
||||
|
||||
def load_config(path: str):
|
||||
with open(path, 'r') as f:
|
||||
data = yaml.safe_load(f)
|
||||
return [Route(**r) for r in data.get('routes', [])]
|
||||
docker_mgr = DockerManager()
|
||||
routes = load_config("config.yml")
|
||||
|
||||
|
||||
# --- Docker Management ---
|
||||
idle_containers = {}
|
||||
idle_timeout = {}
|
||||
idle_check_interval = 5
|
||||
class _LifespanApp:
|
||||
"""Minimal ASGI app that only handles lifespan events."""
|
||||
|
||||
def idle_watcher():
|
||||
client = docker.from_env()
|
||||
while True:
|
||||
now = time()
|
||||
for container, until in list(idle_containers.items()):
|
||||
if until and now > until:
|
||||
try:
|
||||
c = client.containers.get(container)
|
||||
c.stop()
|
||||
print(f"[watcher] Stopped idle container: {container}")
|
||||
except Exception as e:
|
||||
print(f"[watcher] Failed to stop {container}: {e}")
|
||||
finally:
|
||||
idle_containers.pop(container, None)
|
||||
sleep(idle_check_interval)
|
||||
|
||||
Thread(target=idle_watcher, daemon=True).start()
|
||||
|
||||
|
||||
# --- Proxy Middleware ---
|
||||
class ReverseProxyMiddleware(BaseHTTPMiddleware):
|
||||
def __init__(self, app: ASGIApp, routes: list[Route]):
|
||||
super().__init__(app)
|
||||
self.routes = routes
|
||||
self.docker = docker.from_env()
|
||||
|
||||
async def dispatch(self, request: Request, call_next):
|
||||
path = request.url.path
|
||||
route = next((r for r in self.routes if path.startswith(r.path_prefix)), None)
|
||||
|
||||
if not route:
|
||||
return JSONResponse(status_code=status.HTTP_404_NOT_FOUND, content={"detail": "Not Found"})
|
||||
|
||||
if route.container:
|
||||
try:
|
||||
container = self.docker.containers.get(route.container)
|
||||
if container.status != 'running':
|
||||
container.start()
|
||||
print(f"Started container {route.container}")
|
||||
timeout = idle_timeout.get(route.container)
|
||||
if timeout:
|
||||
idle_containers[route.container] = time() + timeout
|
||||
else:
|
||||
idle_containers[route.container] = None
|
||||
except Exception as e:
|
||||
print(f"[proxy] Failed to ensure container '{route.container}': {e}")
|
||||
return JSONResponse(status_code=status.HTTP_502_BAD_GATEWAY, content={"detail": "Container Error"})
|
||||
|
||||
# WebSocket upgrade detection
|
||||
if request.headers.get("upgrade", "").lower() == "websocket":
|
||||
return await self.handle_websocket_upgrade(request, route)
|
||||
|
||||
new_url = route.target.rstrip('/') + path[len(route.path_prefix):]
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
proxied = await client.request(
|
||||
method=request.method,
|
||||
url=new_url,
|
||||
headers=request.headers.raw,
|
||||
content=await request.body(),
|
||||
timeout=30.0
|
||||
)
|
||||
return Response(content=proxied.content, status_code=proxied.status_code, headers=proxied.headers)
|
||||
except Exception as e:
|
||||
print(f"[proxy] Failed proxying to {new_url}: {e}")
|
||||
return JSONResponse(status_code=status.HTTP_502_BAD_GATEWAY, content={"detail": "Upstream Error"})
|
||||
|
||||
async def handle_websocket_upgrade(self, request: Request, route: Route):
|
||||
scope: Scope = request.scope
|
||||
receive: Receive = request.receive
|
||||
send: Send = request._send # Unsafe, but needed for ASGI hijack
|
||||
ws = WebSocket(scope, receive=receive, send=send)
|
||||
await ws.accept()
|
||||
|
||||
target_ws_url = route.target.rstrip('/') + request.url.path[len(route.path_prefix):]
|
||||
if target_ws_url.startswith("http"):
|
||||
target_ws_url = target_ws_url.replace("http", "ws", 1)
|
||||
|
||||
try:
|
||||
async with websockets.connect(target_ws_url) as backend:
|
||||
async def to_backend():
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] != "lifespan":
|
||||
return
|
||||
task = None
|
||||
while True:
|
||||
event = await receive()
|
||||
if event["type"] == "lifespan.startup":
|
||||
task = asyncio.create_task(docker_mgr.idle_watcher())
|
||||
await send({"type": "lifespan.startup.complete"})
|
||||
elif event["type"] == "lifespan.shutdown":
|
||||
if task:
|
||||
task.cancel()
|
||||
try:
|
||||
while True:
|
||||
data = await ws.receive_text()
|
||||
await backend.send(data)
|
||||
except Exception:
|
||||
await backend.close()
|
||||
|
||||
async def from_backend():
|
||||
try:
|
||||
while True:
|
||||
data = await backend.recv()
|
||||
await ws.send_text(data)
|
||||
except Exception:
|
||||
await ws.close()
|
||||
|
||||
await asyncio.gather(to_backend(), from_backend())
|
||||
except Exception as e:
|
||||
print(f"[ws] Proxy error: {e}")
|
||||
await ws.close(code=1011)
|
||||
return Response(status_code=502, content=b"WebSocket proxy error")
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
await send({"type": "lifespan.shutdown.complete"})
|
||||
return
|
||||
|
||||
|
||||
# --- App Setup ---
|
||||
app = FastAPI()
|
||||
config = load_config("config.yml")
|
||||
app.add_middleware(ReverseProxyMiddleware, routes=config)
|
||||
|
||||
|
||||
# Optional: Health check
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
return {"status": "ok"}
|
||||
app = ProxyMiddleware(_LifespanApp(), routes=routes, docker=docker_mgr)
|
||||
|
||||
Reference in New Issue
Block a user