BIG BIG BIG BIG BIG UPDATES

Co-authored-by: Copilot <[email protected]>

backend/app/services/mitmproxy/health.py

@@ -0,0 +1,185 @@
+import asyncio
+import ssl
+from dataclasses import dataclass
+
+from models.proxy import HealthCheckDocument, ProxyDocument
+
+
+@dataclass(slots=True)
+class Target:
+    name: str
+    host: str
+    port: int
+    protocol: str
+
+
+async def check_proxy(proxy: ProxyDocument) -> dict[str, object]:
+    if not proxy.healthcheck.enabled:
+        return {"enabled": False, "overall": "disabled", "checks": []}
+
+    upstream_port = proxy.service.upstream_port or proxy.service.port
+    checks = await asyncio.gather(
+        check_target(
+            Target(
+                name="upstream",
+                host=proxy.service.host,
+                port=upstream_port,
+                protocol=health_protocol(proxy),
+            ),
+            proxy.healthcheck,
+        ),
+        check_target(
+            Target(
+                name="proxy",
+                host=listen_host_for_check(proxy.listen_host),
+                port=proxy.proxy_port,
+                protocol=health_protocol(proxy),
+            ),
+            proxy.healthcheck,
+        ),
+    )
+    overall = "healthy" if all(check["healthy"] for check in checks) else "unhealthy"
+    return {"enabled": True, "overall": overall, "checks": checks}
+
+
+def health_protocol(proxy: ProxyDocument) -> str:
+    return proxy.healthcheck.protocol or proxy.protocol
+
+
+def listen_host_for_check(host: str) -> str:
+    if host in {"", "0.0.0.0", "::"}:
+        return "127.0.0.1"
+    return host
+
+
+async def check_target(
+    target: Target,
+    config: HealthCheckDocument,
+) -> dict[str, object]:
+    try:
+        match target.protocol:
+            case "http":
+                return await check_http(target, config, use_tls=False)
+            case "https":
+                return await check_http(target, config, use_tls=True)
+            case "tcp":
+                return await check_tcp(target, config, use_tls=False)
+            case "tls":
+                return await check_tcp(target, config, use_tls=True)
+            case "udp" | "dtls" | "quic" | "dns" | "http3":
+                return await check_udp(target, config)
+            case _:
+                return result(target, False, f"Unsupported protocol: {target.protocol}")
+    except TimeoutError:
+        return result(target, False, "Timed out")
+    except OSError as error:
+        return result(target, False, str(error))
+
+
+async def check_http(
+    target: Target,
+    config: HealthCheckDocument,
+    *,
+    use_tls: bool,
+) -> dict[str, object]:
+    reader, writer = await open_connection(target, config, use_tls=use_tls)
+    try:
+        path = config.path if config.path.startswith("/") else f"/{config.path}"
+        request = (
+            f"{config.method} {path} HTTP/1.1\r\n"
+            f"Host: {target.host}\r\n"
+            "Connection: close\r\n\r\n"
+        )
+        writer.write(request.encode("ascii"))
+        await writer.drain()
+        line = await asyncio.wait_for(reader.readline(), config.timeout_seconds)
+    finally:
+        writer.close()
+        await writer.wait_closed()
+
+    parts = line.decode("iso-8859-1", errors="replace").split()
+    status = int(parts[1]) if len(parts) >= 2 and parts[1].isdigit() else None
+    expected = config.expected_status
+    healthy = status == expected if expected is not None else status is not None
+    return result(target, healthy, None, {"status": status})
+
+
+async def check_tcp(
+    target: Target,
+    config: HealthCheckDocument,
+    *,
+    use_tls: bool,
+) -> dict[str, object]:
+    _reader, writer = await open_connection(target, config, use_tls=use_tls)
+    writer.close()
+    await writer.wait_closed()
+    return result(target, True)
+
+
+async def check_udp(target: Target, config: HealthCheckDocument) -> dict[str, object]:
+    loop = asyncio.get_running_loop()
+    transport, protocol = await loop.create_datagram_endpoint(
+        lambda: DatagramProbe(config.expect_response),
+        remote_addr=(target.host, target.port),
+    )
+    try:
+        payload = config.udp_payload.encode()
+        transport.sendto(payload)
+        if not config.expect_response:
+            return result(target, True, None, {"mode": "send-only"})
+        data = await asyncio.wait_for(protocol.response, config.timeout_seconds)
+        return result(target, True, None, {"bytes": len(data)})
+    finally:
+        transport.close()
+
+
+async def open_connection(
+    target: Target,
+    config: HealthCheckDocument,
+    *,
+    use_tls: bool,
+):
+    if use_tls and config.verify_tls:
+        context = ssl.create_default_context()
+    elif use_tls:
+        context = ssl.create_default_context()
+        context.check_hostname = False
+        context.verify_mode = ssl.CERT_NONE
+    else:
+        context = None
+    return await asyncio.wait_for(
+        asyncio.open_connection(target.host, target.port, ssl=context),
+        config.timeout_seconds,
+    )
+
+
+class DatagramProbe(asyncio.DatagramProtocol):
+    def __init__(self, expect_response: bool):
+        self.response = asyncio.get_running_loop().create_future()
+        if not expect_response:
+            self.response.set_result(b"")
+
+    def datagram_received(self, data, _addr):
+        if not self.response.done():
+            self.response.set_result(data)
+
+    def error_received(self, exc):
+        if not self.response.done():
+            self.response.set_exception(exc)
+
+
+def result(
+    target: Target,
+    healthy: bool,
+    error: str | None = None,
+    details: dict[str, object] | None = None,
+) -> dict[str, object]:
+    return {
+        "target": target.name,
+        "host": target.host,
+        "port": target.port,
+        "protocol": target.protocol,
+        "healthy": healthy,
+        "error": error,
+        "details": details or {},
+    }