101 lines
4.0 KiB
Python
101 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Read-only candidate checks. No login, database operation or task publication."""
|
|
import argparse
|
|
import base64
|
|
import hashlib
|
|
import os
|
|
import socket
|
|
import threading
|
|
import time
|
|
|
|
import uvicorn
|
|
import wsproto # noqa: F401 - required, pinned WebSocket transport
|
|
|
|
|
|
async def probe_app(scope, receive, send):
|
|
if scope["type"] != "websocket":
|
|
raise RuntimeError("The probe must be handled as a WebSocket")
|
|
if (await receive())["type"] != "websocket.connect":
|
|
raise RuntimeError("Missing WebSocket connection event")
|
|
await send({"type": "websocket.accept"})
|
|
await send({"type": "websocket.send", "text": "eusrr-ready"})
|
|
await send({"type": "websocket.close", "code": 1000})
|
|
|
|
|
|
def handshake(port, path, expected_status):
|
|
key = base64.b64encode(os.urandom(16)).decode("ascii")
|
|
with socket.create_connection(("127.0.0.1", port), timeout=5) as client:
|
|
client.settimeout(5)
|
|
client.sendall((
|
|
f"GET {path} HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\n"
|
|
"Connection: Upgrade\r\nUpgrade: websocket\r\n"
|
|
"Sec-WebSocket-Version: 13\r\n"
|
|
f"Sec-WebSocket-Key: {key}\r\n\r\n"
|
|
).encode("ascii"))
|
|
response = b""
|
|
while b"\r\n\r\n" not in response:
|
|
chunk = client.recv(4096)
|
|
if not chunk or len(response) > 16384:
|
|
raise RuntimeError("Incomplete WebSocket handshake")
|
|
response += chunk
|
|
headers, payload = response.split(b"\r\n\r\n", 1)
|
|
if int(headers.split(b" ", 2)[1]) != expected_status:
|
|
raise RuntimeError(f"Unexpected WebSocket HTTP status (expected {expected_status})")
|
|
if expected_status == 101:
|
|
accept = base64.b64encode(hashlib.sha1(
|
|
(key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode("ascii")
|
|
).digest())
|
|
fields = dict(line.split(b":", 1) for line in headers.split(b"\r\n")[1:] if b":" in line)
|
|
fields = {name.lower(): value.strip() for name, value in fields.items()}
|
|
if fields.get(b"sec-websocket-accept") != accept:
|
|
raise RuntimeError("Invalid WebSocket upgrade response")
|
|
expected_frame = b"\x81\x0beusrr-ready"
|
|
while len(payload) < len(expected_frame):
|
|
chunk = client.recv(4096)
|
|
if not chunk:
|
|
raise RuntimeError("Missing WebSocket probe frame")
|
|
payload += chunk
|
|
if not payload.startswith(expected_frame):
|
|
raise RuntimeError("Invalid WebSocket probe frame")
|
|
|
|
|
|
def check_transport():
|
|
# Exercise Uvicorn's actual HTTP-to-WebSocket upgrade, not just an import.
|
|
listener = socket.socket()
|
|
listener.bind(("127.0.0.1", 0))
|
|
listener.listen(4)
|
|
server = uvicorn.Server(uvicorn.Config(
|
|
probe_app, ws="wsproto", lifespan="off", log_level="error", access_log=False,
|
|
))
|
|
thread = threading.Thread(target=server.run, kwargs={"sockets": [listener]}, daemon=True)
|
|
thread.start()
|
|
try:
|
|
deadline = time.monotonic() + 5
|
|
while not server.started and thread.is_alive() and time.monotonic() < deadline:
|
|
time.sleep(0.05)
|
|
if not server.started:
|
|
raise RuntimeError("Uvicorn WebSocket probe did not start")
|
|
handshake(listener.getsockname()[1], "/", 101)
|
|
finally:
|
|
server.should_exit = True
|
|
thread.join(timeout=5)
|
|
listener.close()
|
|
if thread.is_alive():
|
|
raise RuntimeError("Uvicorn WebSocket probe did not stop")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--app-port", type=int, choices=(9001, 9002))
|
|
args = parser.parse_args()
|
|
check_transport()
|
|
if args.app_port:
|
|
# Anonymous access must be rejected. No ticket is minted or consumed;
|
|
# authentication is never bypassed for a deployment health check.
|
|
handshake(args.app_port, "/ws/", 403)
|
|
print("WebSocket transport and requested candidate checks passed.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|