diff --git a/sdk/python/cocoonsandbox/sandbox.py b/sdk/python/cocoonsandbox/sandbox.py index 7d57855a..0f357330 100644 --- a/sdk/python/cocoonsandbox/sandbox.py +++ b/sdk/python/cocoonsandbox/sandbox.py @@ -381,22 +381,27 @@ def pump_out(): if not chunk: break local.sendall(chunk) + # a close cannot wake the read below, and the socket outlives it with contextlib.suppress(OSError): - local.close() + local.shutdown(socket.SHUT_WR) pump = threading.Thread(target=pump_out, daemon=True) pump.start() try: - while True: - chunk = local.recv(FS_CHUNK) - if not chunk: - break - guest.send(chunk) + with contextlib.suppress(OSError): + while True: + chunk = local.recv(FS_CHUNK) + if not chunk: + break + guest.send(chunk) # half-close, not close: the guest's reply is still in flight. - guest.close_write() + with contextlib.suppress(SandboxError, OSError): + guest.close_write() pump.join() finally: guest.close() + with contextlib.suppress(OSError): + local.close() def _call(self, op: str, expect: str, **fields) -> dict: with self._dial() as conn: diff --git a/sdk/python/tests/test_proxy.py b/sdk/python/tests/test_proxy.py index e4dc99fb..629a8b82 100644 --- a/sdk/python/tests/test_proxy.py +++ b/sdk/python/tests/test_proxy.py @@ -31,6 +31,28 @@ def close(self) -> None: self._closed = True +class ClosedPortConn: + """A guest port whose upstream refuses, reported while the client is idle.""" + + def __init__(self): + self._dialled = False + + def send(self, data: bytes) -> None: + pass + + def recv(self) -> bytes: + if not self._dialled: + self._dialled = True + time.sleep(0.3) + return b"" + + def close_write(self) -> None: + pass + + def close(self) -> None: + pass + + def test_proxy_port_accepts_and_pipes(monkeypatch): sb = Sandbox(client=Client("127.0.0.1:1"), id="sb_1", token="tok", owner="127.0.0.1:1") monkeypatch.setattr(sb, "dial_port", lambda port: FakePortConn()) @@ -50,3 +72,17 @@ def test_proxy_port_accepts_and_pipes(monkeypatch): c.close() finally: listener.close() + + +def test_proxy_port_ends_the_local_connection_when_the_guest_port_is_closed(monkeypatch): + sb = Sandbox(client=Client("127.0.0.1:1"), id="sb_1", token="tok", owner="127.0.0.1:1") + monkeypatch.setattr(sb, "dial_port", lambda port: ClosedPortConn()) + + listener = sb.proxy_port("127.0.0.1:0", 8080) + try: + c = socket.create_connection(listener.getsockname(), timeout=5) + c.settimeout(5) + assert c.recv(256) == b"" + c.close() + finally: + listener.close()