diff --git a/docker-compose.yml b/docker-compose.yml index 5b0c690..1b1c07e 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -45,5 +45,8 @@ services: depends_on: - c2 ports: - - "6080:6080" # noVNC: http://localhost:6080/vnc.html + # Ephemeral host port so replicas don't collide under --scale. Find a worker's + # noVNC with `docker compose port worker 6080` (or `docker ps`), then open + # http://localhost:/vnc.html to watch / solve a challenge. + - "6080" restart: unless-stopped diff --git a/worker/worker.py b/worker/worker.py index 9fc0ab2..d6e6a5a 100644 --- a/worker/worker.py +++ b/worker/worker.py @@ -209,6 +209,11 @@ class LocalForwardingProxy: async def _relay( client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter, up_reader: asyncio.StreamReader, up_writer: asyncio.StreamWriter) -> None: + # Pipe both directions, but tear the whole tunnel down as soon as EITHER side + # closes (mirrors the .NET WhenAny). Waiting for both — as a plain gather does — + # leaks a task holding two sockets on every half-closed connection, which piles + # up fast across a long multi-worker run. Closing both writers when the first + # pipe finishes unblocks the other's pending read so both tasks settle. async def pipe(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: try: while data := await reader.read(65536): @@ -216,10 +221,18 @@ class LocalForwardingProxy: await writer.drain() except Exception: pass - await asyncio.gather( - pipe(client_reader, up_writer), - pipe(up_reader, client_writer), - ) + + a = asyncio.create_task(pipe(client_reader, up_writer)) + b = asyncio.create_task(pipe(up_reader, client_writer)) + try: + await asyncio.wait({a, b}, return_when=asyncio.FIRST_COMPLETED) + finally: + for w in (client_writer, up_writer): + try: + w.close() + except Exception: + pass + await asyncio.gather(a, b, return_exceptions=True) def looks_like_challenge(body: str) -> bool: