diff --git a/backend/worker/inbound/connection_log.py b/backend/worker/inbound/connection_log.py index a8f2651c..0e9762d5 100644 --- a/backend/worker/inbound/connection_log.py +++ b/backend/worker/inbound/connection_log.py @@ -23,6 +23,11 @@ from typing import Callable, Optional # nobody rotates. _MAX_EVENTS = 200 +# How long a kicked panel stays out. Long enough that the disconnect is +# visible and the person notices; short enough that it is plainly not a +# revocation, which is a separate and permanent action. +_KICK_COOLDOWN_SECONDS = 60.0 + @dataclass class Session: @@ -54,6 +59,7 @@ class ConnectionLog: self._lock = threading.Lock() self._sessions: dict[str, Session] = {} self._events: deque[Event] = deque(maxlen=_MAX_EVENTS) + self._cooldowns: dict[str, float] = {} # ── Sessions ────────────────────────────────────────────────────────── @@ -108,8 +114,20 @@ class ConnectionLog: if session is None: return False session.disconnect_requested = True + # A panel reconnects on its own, so without this the person is back + # within two seconds and the button appears to do nothing — + # verified on hardware, where the log read disconnected/connected + # in the same breath. The cooldown makes the disconnect visible and + # deliberately does NOT last: revoking the key is how you stop + # somebody for good, and a kick that silently became permanent + # would be a different promise than the button makes. + self._cooldowns[session.key_id] = self._now() + _KICK_COOLDOWN_SECONDS return True + def cooling_down(self, key_id: str) -> bool: + with self._lock: + return self._cooldowns.get(key_id, 0.0) > self._now() + def disconnect_requested(self, session_id: str) -> bool: with self._lock: session = self._sessions.get(session_id) diff --git a/backend/worker/inbound/listener.py b/backend/worker/inbound/listener.py index 429e3e23..25290278 100644 --- a/backend/worker/inbound/listener.py +++ b/backend/worker/inbound/listener.py @@ -81,6 +81,9 @@ class NodeServicer(pb_grpc.NodeServiceServicer): metadata = {k.lower(): v for k, v in (context.invocation_metadata() or ())} secret = metadata.get(KEY_METADATA_KEY, "") key = self._keys.authenticate(secret, peer=peer) + if key is not None and self._log.cooling_down(key.key_id): + self._log.rejected(peer=peer, detail="recently disconnected by the owner") + return None if key is None: self._log.rejected( peer=peer, detail="no key" if not secret else "key not recognised" diff --git a/backend/worker/inbound/service.py b/backend/worker/inbound/service.py index de674ec6..10d83a6b 100644 --- a/backend/worker/inbound/service.py +++ b/backend/worker/inbound/service.py @@ -297,19 +297,31 @@ class OutboundNodes: entries = [e for e in entries if _endpoint_of(e) != connection.endpoint] entries.append(text.strip()) self._save(entries) + + # Tear down any live session to this machine BEFORE dialling. Without + # this, re-pasting for an already-connected machine saved the new key + # and then short-circuited on the existing connection — so a wrong key + # reported success, kept working on the old session, and only failed + # after a restart, by which time nothing pointed at the paste that + # caused it. Verified on hardware. + await self._drop(connection.endpoint) await self._dial(connection, servicer) return connection - async def remove(self, endpoint: str) -> bool: - entries = [e for e in self.saved() if _endpoint_of(e) != endpoint] - self._save(entries) + async def _drop(self, endpoint: str) -> None: connection = self._connections.pop(endpoint, None) task = self._tasks.pop(endpoint, None) if connection is not None: await connection.stop() if task is not None: task.cancel() - return connection is not None + + async def remove(self, endpoint: str) -> bool: + entries = [e for e in self.saved() if _endpoint_of(e) != endpoint] + self._save(entries) + existed = endpoint in self._connections + await self._drop(endpoint) + return existed async def start_all(self, servicer) -> None: for entry in self.saved(): diff --git a/docs/remote-workers.md b/docs/remote-workers.md index 9ec3022a..238bd6a5 100644 --- a/docs/remote-workers.md +++ b/docs/remote-workers.md @@ -97,6 +97,12 @@ working, which is why each person gets their own. panel currently attached, where it connected from, how many jobs it has run, and a **Disconnect** button. +**Disconnect and Remove do different things.** Disconnect ends the session now +and keeps that person out for a minute — use it to get someone off the card +immediately. Their app reconnects by itself after that, because their +connection string is still valid. To stop someone for good, remove their +connection string instead. + > **This mode is not encrypted.** The connection string is a password that > travels in the clear, so anyone who can watch that network can copy it and > use your GPU — and your reference audio and rendered speech cross the network diff --git a/tests/test_worker_inbound_transport.py b/tests/test_worker_inbound_transport.py index 0f1d2f94..d30a99bc 100644 --- a/tests/test_worker_inbound_transport.py +++ b/tests/test_worker_inbound_transport.py @@ -242,6 +242,14 @@ async def test_the_owner_can_see_who_connected_and_kick_them(inbound): # loop that only wakes on outbound traffic would never notice. await _until(lambda: inbound.log.snapshot()["sessions"] == []) + # And it has to STAY landed for a moment. The panel redials on its own, so + # without a cooldown the person is back within two seconds and the button + # appears to do nothing — which is what it did on hardware, where the log + # read disconnected and connected in the same breath. + assert inbound.log.cooling_down(sessions[0]["key_id"]) is True + await asyncio.sleep(2.5) + assert inbound.log.snapshot()["sessions"] == [], "the kicked panel came straight back" + @pytest.mark.asyncio async def test_an_input_pushed_before_the_assignment_is_there_when_the_task_asks( @@ -419,3 +427,35 @@ async def test_a_staged_result_comes_back_whole_when_its_ref_declares_a_size( await inbound.connection.fetch_result(ref, str(destination)) assert destination.read_bytes() == payload + + +@pytest.mark.asyncio +async def test_repasting_a_key_for_a_connected_machine_redials_it(inbound, monkeypatch): + """Re-pasting must replace the live session, not report success against it. + + Found on hardware: `add` saved the new string and then short-circuited + because a connection to that endpoint already existed. A wrong key + therefore overwrote a working one, answered 200 with connected=true from + the stale session, and only failed after a restart — by which point nothing + pointed back at the paste that caused it. + """ + from worker.inbound import service as inbound_service + + await inbound.connect_panel() + outbound = inbound_service.OutboundNodes() + saved = [] + monkeypatch.setattr(outbound, "saved", lambda: list(saved)) + monkeypatch.setattr(outbound, "_save", lambda entries: saved.clear() or saved.extend(entries)) + + first = f"ovnode://{inbound.keys.issue('One').secret}@127.0.0.1:{inbound.port}" + await outbound.add(first, inbound.servicer) + original = outbound._connections[f"127.0.0.1:{inbound.port}"] + + second = f"ovnode://{inbound.keys.issue('Two').secret}@127.0.0.1:{inbound.port}" + await outbound.add(second, inbound.servicer) + + assert saved == [second], "the newly pasted string must be the saved one" + assert outbound._connections[f"127.0.0.1:{inbound.port}"] is not original, ( + "the old session must be replaced, not reused" + ) + await outbound.stop()