upstream-archive/VoiceStudio
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(worker): persist pinned reconnect credentials

Browse Source
This commit is contained in:
debpalash
2026-08-11 23:54:28 +00:00
parent 5818e19137
commit 76d8c11553
5 changed files with 66 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/worker/inbound/keys.py
+19 -1
View File
@@ -111,6 +111,7 @@ class KeyStore:
self._lock = threading.Lock()
self._keys: dict[str, PanelKey] = {}
self._connection_secrets: dict[str, str] = {}
self._connection_fingerprints: dict[str, str] = {}
self._failures: dict[str, _Failures] = {}
self._load()
@@ -145,6 +146,13 @@ class KeyStore:
for endpoint, secret in connections.items()
if endpoint and secret
}
fingerprints = raw.get("connection_fingerprints", {})
if isinstance(fingerprints, dict):
self._connection_fingerprints = {
str(endpoint): str(fingerprint)
for endpoint, fingerprint in fingerprints.items()
if endpoint and fingerprint
}
def _save_locked(self) -> None:
directory = os.path.dirname(os.path.abspath(self._path))
@@ -153,6 +161,7 @@ class KeyStore:
{
"keys": [asdict(k) for k in self._keys.values()],
"connection_secrets": self._connection_secrets,
"connection_fingerprints": self._connection_fingerprints,
},
indent=2,
).encode("utf-8")
@@ -227,19 +236,28 @@ class KeyStore:
# ── Panel-side connection credentials ───────────────────────────────
def remember_connection_secret(self, endpoint: str, secret: str) -> None:
def remember_connection_secret(
self, endpoint: str, secret: str, fingerprint: str = ""
) -> None:
"""Persist a pasted node secret outside the UI-readable settings store."""
with self._lock:
self._connection_secrets[endpoint] = secret
if fingerprint:
self._connection_fingerprints[endpoint] = fingerprint
self._save_locked()
def connection_secret(self, endpoint: str) -> str:
with self._lock:
return self._connection_secrets.get(endpoint, "")
def connection_fingerprint(self, endpoint: str) -> str:
with self._lock:
return self._connection_fingerprints.get(endpoint, "")
def forget_connection_secret(self, endpoint: str) -> None:
with self._lock:
if self._connection_secrets.pop(endpoint, None) is not None:
self._connection_fingerprints.pop(endpoint, None)
self._save_locked()
# ── Authentication ────────────────────────────────────────────────────
backend/worker/inbound/service.py
+9 -4
View File
@@ -351,7 +351,7 @@ class OutboundNodes:
migrated = True
continue
self.credentials.remember_connection_secret(
connection.endpoint, connection.secret
connection.endpoint, connection.secret, connection.fingerprint
)
endpoints.append(connection.endpoint)
migrated = True
@@ -372,7 +372,9 @@ class OutboundNodes:
# replaces it rather than leaving a dead entry that retries forever.
entries = [e for e in entries if _endpoint_of(e) != connection.endpoint]
entries.append(connection.endpoint)
self.credentials.remember_connection_secret(connection.endpoint, connection.secret)
self.credentials.remember_connection_secret(
connection.endpoint, connection.secret, connection.fingerprint
)
self._save(entries)
# Tear down any live session to this machine BEFORE dialling. Without
@@ -451,11 +453,14 @@ class OutboundNodes:
except ValueError as exc:
raise InvalidConnectionString("That saved node address is not valid.") from exc
secret = self.credentials.connection_secret(endpoint)
if not host or not port or not secret:
fingerprint = self.credentials.connection_fingerprint(endpoint)
if not host or not port or not secret or not fingerprint:
raise InvalidConnectionString(
"That saved GPU connection has no protected key. Paste its connection string again."
)
return Connection(host=host, port=port, secret=secret)
return Connection(
host=host, port=port, secret=secret, fingerprint=fingerprint
)
def _endpoint_of(entry: str) -> str:
tests/test_worker_inbound_keys.py
+14 -1
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
import json
import os
import stat
import pytest
@@ -86,6 +87,8 @@ def test_a_locked_out_peer_is_refused_even_with_the_right_key(store):
def test_failed_key_throttle_ignores_ephemeral_source_ports(store):
"""Reconnects from one host must contribute to the same lockout."""
from worker.inbound import keys as keys_module
store.issue("Alice")
for port in range(41000, 41000 + keys_module._MAX_FAILURES):
@@ -95,6 +98,8 @@ def test_failed_key_throttle_ignores_ephemeral_source_ports(store):
def test_failed_key_throttle_normalises_bracketed_ipv6_ports(store):
from worker.inbound import keys as keys_module
store.issue("Alice")
for port in range(41000, 41000 + keys_module._MAX_FAILURES):
@@ -121,6 +126,8 @@ def test_keys_survive_a_restart(store, tmp_path):
def test_pasted_connection_secrets_use_the_protected_key_file(store, tmp_path):
from worker.inbound.keys import KEY_PREFIX, KeyStore
secret = KEY_PREFIX + "s" * 40
store.remember_connection_secret("10.0.0.2:7444", secret)
@@ -133,9 +140,14 @@ def test_pasted_connection_secrets_use_the_protected_key_file(store, tmp_path):
def test_legacy_saved_connection_is_migrated_out_of_settings(store, monkeypatch):
from worker.inbound import service as inbound_service
from worker.inbound.connection_string import format_connection
from worker.inbound.keys import KEY_PREFIX
secret = KEY_PREFIX + "s" * 40
legacy = format_connection(host="10.0.0.2", port=7444, secret=secret)
fingerprint = "a" * 64
legacy = format_connection(
host="10.0.0.2", port=7444, secret=secret, fingerprint=fingerprint
)
settings = {inbound_service._SAVED_KEY: legacy}
monkeypatch.setattr(
inbound_service,
@@ -153,6 +165,7 @@ def test_legacy_saved_connection_is_migrated_out_of_settings(store, monkeypatch)
assert outbound.saved() == ["10.0.0.2:7444"]
assert secret not in settings[inbound_service._SAVED_KEY]
assert store.connection_secret("10.0.0.2:7444") == secret
assert store.connection_fingerprint("10.0.0.2:7444") == fingerprint
def test_a_corrupt_key_file_is_reported_rather_than_read_as_no_keys(tmp_path, caplog):
tests/test_worker_inbound_transport.py
+23 -10
View File
@@ -322,6 +322,8 @@ async def test_attach_ends_when_its_incoming_reader_stops(tmp_path, monkeypatch)
"""A dead reader must not leave a node advertising a healthy session."""
from worker.protocol.gen import worker_v1_pb2 as pb
worker = _worker_modules()
class FakeClient:
def build_register_request(self):
return pb.RegisterRequest()
@@ -348,11 +350,12 @@ async def test_attach_ends_when_its_incoming_reader_stops(tmp_path, monkeypatch)
async def abort(self, code, message):
raise AssertionError(f"unexpected abort: {code}: {message}")
servicer = NodeListener(
keys=KeyStore(str(tmp_path / "keys.json")),
log=ConnectionLog(),
artifacts=ArtifactStore(str(tmp_path / "staged")),
servicer = worker.NodeListener(
keys=worker.KeyStore(str(tmp_path / "keys.json")),
log=worker.ConnectionLog(),
artifacts=worker.ArtifactStore(str(tmp_path / "staged")),
client_factory=lambda _artifacts, _key_id: FakeClient(),
credentials=worker.tls.generate_self_signed(hostnames=["127.0.0.1"]),
)._servicer
monkeypatch.setattr(servicer, "_authenticate", lambda _context: ("panel-a", "A"))
@@ -478,7 +481,7 @@ async def test_an_offset_mismatch_removes_the_partial_input(inbound):
async def test_staged_artifacts_are_isolated_by_panel_key(tmp_path):
from worker.protocol.gen import worker_v1_pb2 as pb
store = ArtifactStore(str(tmp_path / "staged"))
store = _worker_modules().ArtifactStore(str(tmp_path / "staged"))
result = await store.publish(
pb.TaskRef(task_id="t1", attempt_id="a1"),
b"alice audio",
@@ -649,8 +652,14 @@ async def test_result_pull_stops_at_its_runtime_byte_cap(tmp_path):
from worker.inbound.connection_string import Connection
from worker.protocol.gen import worker_v1_pb2 as pb
connection = NodeConnection(
object(), Connection(host="127.0.0.1", port=7444, secret="ovnode_" + "s" * 40)
connection = _worker_modules().NodeConnection(
object(),
Connection(
host="127.0.0.1",
port=7444,
secret="ovnode_" + "s" * 40,
fingerprint="a" * 64,
),
)
class Stub:
@@ -706,7 +715,9 @@ async def test_repasting_a_key_for_a_connected_machine_redials_it(inbound, monke
endpoint = f"127.0.0.1:{inbound.port}"
assert saved == [endpoint], "settings must persist only the non-secret endpoint"
assert inbound.keys.connection_secret(endpoint) == parse_connection(second).secret
parsed_second = inbound.worker.parse_connection(second)
assert inbound.keys.connection_secret(endpoint) == parsed_second.secret
assert inbound.keys.connection_fingerprint(endpoint) == parsed_second.fingerprint
assert outbound._connections[f"127.0.0.1:{inbound.port}"] is not original, (
"the old session must be replaced, not reused"
)
@@ -717,10 +728,11 @@ async def test_repasting_a_key_for_a_connected_machine_redials_it(inbound, monke
async def test_saved_endpoint_reloads_its_key_from_protected_storage(tmp_path, monkeypatch):
from worker.inbound import service as inbound_service
store = KeyStore(str(tmp_path / "keys.json"))
store = _worker_modules().KeyStore(str(tmp_path / "keys.json"))
endpoint = "10.0.0.2:7444"
secret = "ovnode_" + "s" * 40
store.remember_connection_secret(endpoint, secret)
fingerprint = "a" * 64
store.remember_connection_secret(endpoint, secret, fingerprint)
outbound = inbound_service.OutboundNodes(store)
monkeypatch.setattr(outbound, "saved", lambda: [endpoint])
dialled = []
@@ -734,6 +746,7 @@ async def test_saved_endpoint_reloads_its_key_from_protected_storage(tmp_path, m
assert len(dialled) == 1
assert dialled[0].endpoint == endpoint
assert dialled[0].secret == secret
assert dialled[0].fingerprint == fingerprint
@pytest.mark.asyncio
tests/test_worker_transport.py
+1
View File
@@ -664,6 +664,7 @@ async def test_control_stream_setup_failure_clears_open_flag():
with pytest.raises(RuntimeError, match="queue closed"):
await servicer.Control(None, object())
assert session.stream_open is False
@pytest.mark.asyncio
async def test_stream_start_refuses_a_worker_removed_after_registration(tmp_path):
"""Both stream directions must leave a raced session reusable/closed."""
from worker.transport.server import SESSION_METADATA_KEY, _Session