diff --git a/README.md b/README.md index 53eb614..2b86efd 100644 --- a/README.md +++ b/README.md @@ -159,6 +159,16 @@ server = Server(connect="tunnel", tunnel_config=config) # ✓ No public IPs required ``` +In a cluster, don't bake a key into your image — mint one per start from a key +broker (`$SANDD_KEYBROKER_URL`) and take the headscale URL from +`$SANDD_TUNNEL_SERVER`: + +```python +from sandd import Server, tunnel_config_from_env + +server = Server(connect="tunnel", tunnel_config=tunnel_config_from_env()) +``` + See [Tunnel Mode Guide](./docs/proposals/TUNNEL.md) for setup instructions. ## Documentation diff --git a/docs/proposals/TUNNEL.md b/docs/proposals/TUNNEL.md index 9abf729..7ec33ce 100644 --- a/docs/proposals/TUNNEL.md +++ b/docs/proposals/TUNNEL.md @@ -259,6 +259,26 @@ config = TunnelConfig( server = Server(connect="tunnel", tunnel_config=config) ``` +#### With a key broker + +A static `authkey` in your source is fine for a laptop, but in a cluster the key is +minted per-controller by a key broker that holds the only headscale admin +credentials (Nebula deploys one as `nebula-keybroker`). `tunnel_config_from_env` +does that mint for you — it reads the headscale URL from `$SANDD_TUNNEL_SERVER`, +POSTs `/keys?kind=controller` to `$SANDD_KEYBROKER_URL`, and returns a ready +`TunnelConfig`, so no key is ever written into an image or a manifest: + +```python +from sandd import Server, tunnel_config_from_env + +server = Server(connect="tunnel", tunnel_config=tunnel_config_from_env()) +``` + +It retries the mint (3 attempts, 1s/2s backoff) because the broker often runs as a +headscale sidecar whose socket may not be up when your controller starts. Pass +`broker_url=`, `attempts=`, or `timeout=` to override, and `kind="daemon"` for a +daemon-policy key. + ### Docker Image Use the tunnel-enabled image. Build it yourself like this: diff --git a/protocol/src/lib.rs b/protocol/src/lib.rs index 33fbdfe..007adda 100644 --- a/protocol/src/lib.rs +++ b/protocol/src/lib.rs @@ -25,6 +25,27 @@ pub enum Message { message: String, }, Heartbeat, + /// Server -> daemon: the outcome of a `Heartbeat`, mirroring `RegisterAck`. + /// + /// `success: false` means the daemon is NOT in the registry and must send `Register` + /// again on this same connection. That happens because the server reaps daemons + /// whose heartbeats stall past a threshold, which mesh churn can cause WITHOUT + /// breaking TCP — so a daemon can be evicted while its socket is still healthy. + /// `Register` is sent once per connection, and the daemon cannot detect the eviction + /// on its own: its heartbeat writes keep succeeding, so its dead-connection signal + /// never fires. Without this it stays invisible (no exec, no logs) until the socket + /// truly breaks, which its own heartbeats keep preventing. + /// + /// `success: true` is also sent, and is load-bearing: it is the daemon's only + /// application-level proof that the controller is still PROCESSING messages, not + /// merely accepting bytes into a socket buffer. The daemon reconnects when acks stop + /// arriving (see the ack-timeout check in sandd's serve loop). + /// + /// `reason` explains a failure and is empty on success. + HeartbeatAck { + success: bool, + reason: String, + }, Pong, ExecuteCommand { request_id: String, diff --git a/python/sandd/__init__.py b/python/sandd/__init__.py index 3897ac3..1921922 100644 --- a/python/sandd/__init__.py +++ b/python/sandd/__init__.py @@ -42,6 +42,7 @@ from .models import CommandResult, ServerStats, DaemonInfo from .server import Server from .async_server import AsyncServer +from .keys import tunnel_config_from_env try: from ._core import Session, TunnelConfig @@ -58,4 +59,5 @@ "ServerStats", "DaemonInfo", "TunnelConfig", + "tunnel_config_from_env", ] diff --git a/python/sandd/keys.py b/python/sandd/keys.py new file mode 100644 index 0000000..75d3e77 --- /dev/null +++ b/python/sandd/keys.py @@ -0,0 +1,147 @@ +"""Mesh auth-key helpers for tunnel mode. + +A controller in tunnel mode needs a headscale pre-auth key before it can join the +mesh. In a Nebula cluster that key is NOT a static secret: an in-cluster key broker +mints a fresh reusable+ephemeral one per caller, and the broker is the only component +holding headscale admin authority. Every controller therefore had to hand-roll the +same POST-and-parse against the broker before constructing a ``Server`` — see the +`sandd-controller` sample in the Nebula repo, which carried it inline. That +boilerplate lives here instead. + +Only the stdlib is used (``urllib``): the package declares no runtime dependencies, +and a controller image that must talk to the broker before it can do anything else +is the wrong place to require ``requests``. + +Keys are secrets. Nothing here logs, prints, or embeds a key in an exception +message — including on the error paths, where the HTTP status is the actionable +signal and the body may still be key material. +""" + +import json +import os +import time +import urllib.error +import urllib.parse +import urllib.request +from typing import Optional + +try: + from ._core import TunnelConfig +except ImportError as e: # pragma: no cover - mirrors server.py's import guard + raise ImportError( + "Failed to import Rust extension. Please build the package with: make install" + ) from e + +__all__ = ["mint_authkey", "tunnel_config_from_env"] + +#: Env var holding the key-broker root URL, e.g. +#: ``http://nebula-keybroker.nebula-system:8090``. +KEYBROKER_URL_ENV = "SANDD_KEYBROKER_URL" + +#: Env var holding the headscale URL the controller joins, e.g. +#: ``http://nebula-headscale.nebula-system``. Consumed by +#: :func:`tunnel_config_from_env` only. +TUNNEL_SERVER_ENV = "SANDD_TUNNEL_SERVER" + + +def mint_authkey( + kind: str = "controller", + broker_url: Optional[str] = None, + timeout: float = 10.0, + attempts: int = 3, +) -> str: + """Mint a fresh headscale pre-auth key from the key broker. + + Args: + kind: Key policy to request, "controller" or "daemon". The broker owns the + policy (reusability, ephemerality, expiry); the caller only names a role. + broker_url: Broker root URL. Defaults to ``$SANDD_KEYBROKER_URL``. + timeout: Per-attempt timeout in seconds. A healthy mint is sub-second — the + broker shells out to a local CLI — so this bounds a wedged broker rather + than a slow one. + attempts: Total tries, with 1s/2s/... backoff between them. Defaults to 3 + because the broker commonly runs as a headscale sidecar and its socket + may not be up yet when a controller starts; a startup race should not + crash-loop the pod. + + Returns: + The key. It is a secret — do not log or print it. + + Raises: + ValueError: If no broker URL was given or found in the environment. + RuntimeError: If every attempt failed, or the broker returned no key. + """ + if attempts < 1: + raise ValueError(f"attempts must be >= 1, got {attempts}") + + base = broker_url if broker_url is not None else os.environ.get(KEYBROKER_URL_ENV) + if not base or not base.strip(): + raise ValueError( + "no key-broker URL: pass broker_url= or set " + f"${KEYBROKER_URL_ENV} (e.g. http://nebula-keybroker.nebula-system:8090)" + ) + + url = f"{base.strip().rstrip('/')}/keys?kind={urllib.parse.quote(kind)}" + + last_error = None + for attempt in range(attempts): + if attempt: + time.sleep(float(attempt)) + try: + return _mint_once(url, timeout) + except (urllib.error.URLError, OSError, ValueError) as e: + # URLError covers HTTPError (a 4xx/5xx from the broker) and connection + # failures alike; both are worth retrying, since a 502 here means the + # broker's own call to headscale failed transiently. ValueError is a + # malformed/empty response body. + last_error = e + + raise RuntimeError( + f"minting {kind} key from key broker failed after {attempts} attempt(s): {last_error}" + ) + + +def _mint_once(url: str, timeout: float) -> str: + """POST once and return the key, or raise for the caller to retry.""" + req = urllib.request.Request(url, method="POST") + with urllib.request.urlopen(req, timeout=timeout) as resp: + # Never surface the body in an error: on success it is key material. + payload = json.load(resp) + key = payload.get("key") if isinstance(payload, dict) else None + if not key: + raise ValueError("key broker returned an empty key") + return key + + +def tunnel_config_from_env( + kind: str = "controller", server: Optional[str] = None, **kwargs +) -> TunnelConfig: + """Build a :class:`TunnelConfig` from the environment, minting the auth key. + + Reads the headscale URL from ``$SANDD_TUNNEL_SERVER`` and mints a key via + :func:`mint_authkey`, so a controller needs no key handling of its own: + + >>> from sandd import Server, tunnel_config_from_env + >>> server = Server(connect="tunnel", tunnel_config=tunnel_config_from_env()) + + Args: + kind: Passed through to :func:`mint_authkey`. + server: headscale URL to join. Defaults to ``$SANDD_TUNNEL_SERVER``, which + is how it is set in-cluster; pass it explicitly to override, mirroring + ``broker_url``. + **kwargs: Passed through to :func:`mint_authkey` (``broker_url``, + ``timeout``, ``attempts``). + + Raises: + ValueError: If no headscale URL was given or found in the environment, or + the broker URL is missing. + RuntimeError: If minting failed. + """ + if server is None: + server = os.environ.get(TUNNEL_SERVER_ENV) + if not server or not server.strip(): + raise ValueError( + f"no headscale URL: pass server= or set ${TUNNEL_SERVER_ENV} " + "(e.g. http://nebula-headscale.nebula-system)" + ) + return TunnelConfig(authkey=mint_authkey(kind, **kwargs), server=server.strip()) diff --git a/python/sandd/server.py b/python/sandd/server.py index 9392bb0..1b3c9df 100644 --- a/python/sandd/server.py +++ b/python/sandd/server.py @@ -45,6 +45,11 @@ class Server: ... ) >>> server = Server(connect="tunnel", tunnel_config=config) >>> result = server.exec("daemon-1", "hostname") + + >>> # Production, key minted from the in-cluster broker: reads + >>> # $SANDD_TUNNEL_SERVER and $SANDD_KEYBROKER_URL, no key handling here + >>> from sandd import tunnel_config_from_env + >>> server = Server(connect="tunnel", tunnel_config=tunnel_config_from_env()) """ def __init__( diff --git a/python/tests/test_e2e_tunnel.py b/python/tests/test_e2e_tunnel.py index 66cb2b6..0882ad2 100644 --- a/python/tests/test_e2e_tunnel.py +++ b/python/tests/test_e2e_tunnel.py @@ -83,15 +83,29 @@ def tunnel_stack(): capture_output=True, text=True, ) - key = subprocess.run( - [ - "docker", "exec", hs, "headscale", "preauthkeys", "create", - "--user", "sandd", "--reusable", "--expiration", "1h", - ], - check=True, - capture_output=True, - text=True, - ).stdout.strip().splitlines()[-1].strip() + key = ( + subprocess.run( + [ + "docker", + "exec", + hs, + "headscale", + "preauthkeys", + "create", + "--user", + "sandd", + "--reusable", + "--expiration", + "1h", + ], + check=True, + capture_output=True, + text=True, + ) + .stdout.strip() + .splitlines()[-1] + .strip() + ) assert key and " " not in key, f"unexpected preauthkey output: {key!r}" # 3. controller + daemon, with the freshly-minted key in their env. Compose diff --git a/python/tests/test_keys.py b/python/tests/test_keys.py new file mode 100644 index 0000000..2ec0833 --- /dev/null +++ b/python/tests/test_keys.py @@ -0,0 +1,207 @@ +"""Unit tests for the key-broker helpers (sandd.keys). + +A real loopback HTTP server stands in for the broker rather than a mock of +urllib: the helper's job is exactly the HTTP request/parse, so mocking that away +would test nothing. Each test drives a handler that records what it received. +""" + +import json +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer + +import pytest +from sandd import tunnel_config_from_env +from sandd.keys import KEYBROKER_URL_ENV, TUNNEL_SERVER_ENV, mint_authkey + + +class _Broker: + """A stub key broker on loopback. `responses` is a list of (status, body) + consumed one per request, so retry behaviour can be scripted.""" + + def __init__(self, responses): + self.responses = list(responses) + self.requests = [] # (method, path) per received request + + broker = self + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + broker.requests.append(("POST", self.path)) + broker._respond(self) + + def do_GET(self): + broker.requests.append(("GET", self.path)) + broker._respond(self) + + def log_message(self, *args): + pass # keep pytest output clean + + self._server = HTTPServer(("127.0.0.1", 0), Handler) + self.url = f"http://127.0.0.1:{self._server.server_port}" + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + + def _respond(self, handler): + status, body = ( + self.responses.pop(0) if self.responses else (500, "no scripted response") + ) + payload = body.encode() if isinstance(body, str) else json.dumps(body).encode() + handler.send_response(status) + handler.send_header("Content-Type", "application/json") + handler.send_header("Content-Length", str(len(payload))) + handler.end_headers() + handler.wfile.write(payload) + + def __enter__(self): + self._thread.start() + return self + + def __exit__(self, *exc): + self._server.shutdown() + self._server.server_close() + self._thread.join(timeout=5) + + +class TestMintAuthkey: + """Test mint_authkey""" + + def test_returns_key(self): + """A 200 with a key yields that key""" + with _Broker([(200, {"key": "nodekey-abc"})]) as broker: + assert mint_authkey(broker_url=broker.url) == "nodekey-abc" + + def test_posts_to_keys_with_kind(self): + """The broker is called with POST /keys?kind=, no doubled slash""" + with _Broker([(200, {"key": "k"})]) as broker: + mint_authkey(kind="daemon", broker_url=broker.url + "/") + assert broker.requests == [("POST", "/keys?kind=daemon")] + + def test_defaults_to_controller_kind(self): + """kind defaults to controller — the SDK's caller is the controller""" + with _Broker([(200, {"key": "k"})]) as broker: + mint_authkey(broker_url=broker.url) + assert broker.requests == [("POST", "/keys?kind=controller")] + + def test_reads_broker_url_from_env(self, monkeypatch): + """With no broker_url, $SANDD_KEYBROKER_URL is used""" + with _Broker([(200, {"key": "from-env"})]) as broker: + monkeypatch.setenv(KEYBROKER_URL_ENV, broker.url) + assert mint_authkey() == "from-env" + + def test_explicit_url_wins_over_env(self, monkeypatch): + with _Broker([(200, {"key": "explicit"})]) as broker: + monkeypatch.setenv(KEYBROKER_URL_ENV, "http://127.0.0.1:1/unused") + assert mint_authkey(broker_url=broker.url) == "explicit" + + def test_missing_url_raises(self, monkeypatch): + monkeypatch.delenv(KEYBROKER_URL_ENV, raising=False) + with pytest.raises(ValueError, match=KEYBROKER_URL_ENV): + mint_authkey() + + def test_blank_url_raises(self, monkeypatch): + monkeypatch.setenv(KEYBROKER_URL_ENV, " ") + with pytest.raises(ValueError, match=KEYBROKER_URL_ENV): + mint_authkey() + + def test_retries_then_succeeds(self): + """A transient 502 (broker's own headscale call failed) is retried. + + This is the startup race the broker sidecar actually exhibits, so the + retry must be real, not decorative. + """ + with _Broker([(502, "failed to mint key"), (200, {"key": "second-try"})]) as b: + assert mint_authkey(broker_url=b.url, attempts=2) == "second-try" + assert len(b.requests) == 2 + + def test_exhausted_attempts_raise(self): + with _Broker([(502, "nope"), (502, "nope")]) as broker: + with pytest.raises(RuntimeError, match="after 2 attempt"): + mint_authkey(broker_url=broker.url, attempts=2) + assert len(broker.requests) == 2 + + def test_empty_key_raises(self): + """A 200 whose body has no key is a failure, not an empty key""" + with _Broker([(200, {"key": ""})]) as broker: + with pytest.raises(RuntimeError): + mint_authkey(broker_url=broker.url, attempts=1) + + def test_malformed_body_raises(self): + with _Broker([(200, "not json")]) as broker: + with pytest.raises(RuntimeError): + mint_authkey(broker_url=broker.url, attempts=1) + + def test_unreachable_broker_raises(self): + """Port 1 on loopback refuses connections — no server needed""" + with pytest.raises(RuntimeError): + mint_authkey(broker_url="http://127.0.0.1:1", attempts=1) + + def test_zero_attempts_rejected(self): + with pytest.raises(ValueError, match="attempts"): + mint_authkey(broker_url="http://127.0.0.1:1", attempts=0) + + def test_key_never_appears_in_error(self): + """Keys are secrets: the failure path must not echo the body. + + A 200 with a key under the wrong field name is the nastiest case — the + body IS key material and the helper still has to fail. + """ + with _Broker([(200, {"authkey": "SECRET-KEY-MATERIAL"})]) as broker: + with pytest.raises(RuntimeError) as excinfo: + mint_authkey(broker_url=broker.url, attempts=1) + assert "SECRET-KEY-MATERIAL" not in str(excinfo.value) + + +class TestTunnelConfigFromEnv: + """Test tunnel_config_from_env""" + + def test_builds_config(self, monkeypatch): + with _Broker([(200, {"key": "minted"})]) as broker: + monkeypatch.setenv(KEYBROKER_URL_ENV, broker.url) + monkeypatch.setenv(TUNNEL_SERVER_ENV, "http://headscale.test") + config = tunnel_config_from_env() + assert config.authkey == "minted" + assert config.server == "http://headscale.test" + + def test_missing_tunnel_server_raises(self, monkeypatch): + """Fails BEFORE minting: no point burning a key we can't use""" + with _Broker([(200, {"key": "unused"})]) as broker: + monkeypatch.setenv(KEYBROKER_URL_ENV, broker.url) + monkeypatch.delenv(TUNNEL_SERVER_ENV, raising=False) + with pytest.raises(ValueError, match=TUNNEL_SERVER_ENV): + tunnel_config_from_env() + assert broker.requests == [] + + def test_explicit_server_wins_over_env(self, monkeypatch): + with _Broker([(200, {"key": "k"})]) as broker: + monkeypatch.setenv(KEYBROKER_URL_ENV, broker.url) + monkeypatch.setenv(TUNNEL_SERVER_ENV, "http://from-env") + config = tunnel_config_from_env(server="http://explicit") + assert config.server == "http://explicit" + + def test_explicit_server_without_env(self, monkeypatch): + """server= alone is enough — $SANDD_TUNNEL_SERVER need not be set""" + with _Broker([(200, {"key": "k"})]) as broker: + monkeypatch.delenv(TUNNEL_SERVER_ENV, raising=False) + config = tunnel_config_from_env( + server="http://headscale.test", broker_url=broker.url + ) + assert config.server == "http://headscale.test" + + def test_blank_server_raises(self, monkeypatch): + with _Broker([(200, {"key": "unused"})]) as broker: + monkeypatch.setenv(KEYBROKER_URL_ENV, broker.url) + with pytest.raises(ValueError, match=TUNNEL_SERVER_ENV): + tunnel_config_from_env(server=" ") + assert broker.requests == [] + + def test_forwards_kwargs_to_mint(self, monkeypatch): + """kind/attempts reach mint_authkey rather than being silently dropped""" + with _Broker([(502, "x"), (200, {"key": "k"})]) as broker: + monkeypatch.setenv(TUNNEL_SERVER_ENV, "http://headscale.test") + config = tunnel_config_from_env( + kind="daemon", broker_url=broker.url, attempts=2 + ) + assert config.authkey == "k" + assert broker.requests == [ + ("POST", "/keys?kind=daemon"), + ("POST", "/keys?kind=daemon"), + ] diff --git a/sandd/src/main.rs b/sandd/src/main.rs index 8540503..3dd3938 100644 --- a/sandd/src/main.rs +++ b/sandd/src/main.rs @@ -10,7 +10,7 @@ use futures_util::{SinkExt, StreamExt}; use sandd_protocol::Message; use std::collections::HashMap; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use sysinfo::System; use tokio_tungstenite::tungstenite::protocol::Message as WsMessage; use tracing::{debug, error, info, warn}; @@ -344,10 +344,12 @@ where labels, }; - // Send registration + // Send registration. Clone: `metadata` is needed again if the controller later + // reports us as unregistered and we re-register on this same socket (see the + // HeartbeatAck arm in the message loop below). let register_msg = Message::Register { daemon_id: daemon_id.to_string(), - metadata, + metadata: metadata.clone(), }; let register_json = serde_json::to_string(®ister_msg)?; ws_tx.send(WsMessage::Text(register_json)).await?; @@ -426,6 +428,23 @@ where // without being moved (it may be `!Unpin`). tokio::pin!(shutdown); tokio::pin!(dead_rx); + + // Last time the controller acked a heartbeat. A successful heartbeat WRITE only + // proves bytes reached a socket buffer, so a controller that is connected but no + // longer processing (wedged event loop, half-open path over DERP) looks perfectly + // healthy to `dead_rx`. An ack, by contrast, is proof of processing — so if acks + // stop arriving while writes keep succeeding, the connection is useless and we + // reconnect. Seeded at connection time: registration just completed, so the + // controller was responsive a moment ago. + let mut last_ack = Instant::now(); + // Allow several missed acks before giving up, for the same reason the controller's + // own reaper allows ~6: mesh churn (DERP peer reconfig, netmap propagation) can + // stall traffic for tens of seconds without anything being broken. Tie it to the + // heartbeat interval so the two stay in step if that is retuned. + let ack_timeout = Duration::from_secs(heartbeat_interval.saturating_mul(6).max(30)); + // Drives the deadline check below. Independent of the heartbeat interval: it only + // decides how promptly a breach is noticed, not how long the deadline is. + let mut ack_check = tokio::time::interval(Duration::from_secs(1)); let outcome = loop { tokio::select! { // Poll shutdown FIRST. With `biased`, tokio checks branches top to @@ -455,6 +474,27 @@ where break ServeOutcome::Disconnected; } + // Writes keep succeeding but the controller stopped acking => it is + // connected yet not processing (wedged, or a half-open path that only + // fails on read). Reconnect rather than sit on a socket that cannot + // deliver work. + // + // Ticks UNCONDITIONALLY and tests the deadline in the body, rather than + // gating the branch on `if last_ack.elapsed() >= ack_timeout`: a disabled + // branch is re-evaluated only when some OTHER branch wakes the loop, and a + // silent controller means ws_rx.next() blocks forever — so the guard would + // never be re-checked in exactly the case it exists to catch. + _ = ack_check.tick() => { + let silent_for = last_ack.elapsed(); + if silent_for >= ack_timeout { + warn!( + "No heartbeat ack for {}s (controller connected but unresponsive); reconnecting", + silent_for.as_secs() + ); + break ServeOutcome::Disconnected; + } + } + msg = ws_rx.next() => { let msg = match msg { Some(Ok(WsMessage::Text(text))) => text, @@ -481,6 +521,48 @@ where } }; + // Heartbeat acks are handled here, not in handle_message: a failure means + // re-registering, which needs `metadata` from this scope, and every ack + // refreshes the liveness deadline tracked by this loop. + if let Message::HeartbeatAck { success, ref reason } = message { + // Any ack proves the controller is PROCESSING, not just accepting + // bytes into a socket buffer — that is what makes the ack-timeout + // check below able to spot a hung-but-connected controller. + last_ack = Instant::now(); + + if !success { + // The controller reaped us while this socket stayed healthy (mesh + // churn can stall heartbeats past its threshold without breaking + // TCP). We are invisible to it — no exec, no logs — until we + // register again, and we cannot detect that any other way: our + // heartbeat writes keep succeeding, so the dead-connection signal + // never fires. Only we hold our metadata, so re-sending Register + // is what restores the entry faithfully. + // + // Re-register IN PLACE rather than reconnecting: the socket is + // demonstrably fine (this ack just arrived on it), so a reconnect + // would pay the mesh dial plus the backoff sleep to rebuild a + // connection we already have. + warn!("Controller rejected heartbeat ({}); re-registering", reason); + let register = Message::Register { + daemon_id: daemon_id.to_string(), + metadata: metadata.clone(), + }; + match serde_json::to_string(®ister) { + Ok(json) => { + let mut tx = ws_tx_clone.lock().await; + if tx.send(WsMessage::Text(json)).await.is_err() { + // The socket died as we replied; reconnect instead. + error!("Failed to re-register; reconnecting"); + break ServeOutcome::Disconnected; + } + } + Err(e) => error!("Failed to serialize re-registration: {}", e), + } + } + continue; + } + // Handle message inline if let Err(e) = handle_message( message, @@ -1127,6 +1209,202 @@ mod shutdown_tests { assert_eq!(outcome.unwrap(), ServeOutcome::Disconnected); } + /// Read frames until one that is not a `Heartbeat` arrives, or the wait times + /// out. The daemon's heartbeat interval fires immediately on its first tick, so + /// beats are routinely interleaved with whatever a test is actually looking for; + /// `None` means the daemon sent nothing but heartbeats. + async fn next_non_heartbeat(server: &mut WebSocketStream) -> Option { + loop { + let frame = tokio::time::timeout(Duration::from_millis(500), server.next()) + .await + .ok()?? + .ok()?; + let text = match frame { + WsMessage::Text(text) => text, + _ => continue, + }; + match serde_json::from_str::(&text) { + Ok(Message::Heartbeat) => continue, + Ok(msg) => return Some(msg), + Err(_) => continue, + } + } + } + + /// A rejected heartbeat must make the daemon re-send `Register` ON THE SAME + /// socket, and keep serving. The controller reaps daemons whose heartbeats + /// stall past its threshold, which mesh churn can cause without breaking TCP + /// — so the daemon can be evicted while its socket is healthy. Re-registering + /// is the only recovery, because only the daemon holds its metadata. + #[tokio::test] + async fn rejected_heartbeat_reregisters_in_place() { + let (client, mut server) = ws_pair().await; + + let shutdown = std::future::pending::<()>(); + // 3600s heartbeat interval: the daemon's own heartbeat never fires in-test, + // so the only Register after the handshake is the re-registration. + let daemon = serve(client, "test-daemon", 3600, HashMap::new(), shutdown); + + let controller = async { + ack_registration(&mut server).await; + + let nack = Message::HeartbeatAck { + success: false, + reason: "daemon is not registered".to_string(), + }; + server + .send(WsMessage::Text(serde_json::to_string(&nack).unwrap())) + .await + .unwrap(); + + // The re-registration must arrive on this same connection, carrying the + // daemon's own id and metadata rather than anything server-side. Skip + // past heartbeats: the daemon's interval fires immediately on the first + // tick, so a beat can be in flight ahead of the Register. + let reregistered = next_non_heartbeat(&mut server).await.is_some_and(|msg| { + matches!( + &msg, + Message::Register { daemon_id, metadata } + if daemon_id == "test-daemon" && !metadata.hostname.is_empty() + ) + }); + + // End the session so `serve` returns and the join below completes. + server.close(None).await.unwrap(); + reregistered + }; + + let (outcome, reregistered) = tokio::join!(daemon, controller); + assert!(reregistered, "daemon did not re-register after a rejected heartbeat"); + // A rejected heartbeat must NOT tear down a demonstrably working socket: + // serve stays in its loop and only ends here because the controller closed. + assert_eq!(outcome.unwrap(), ServeOutcome::Disconnected); + } + + /// End-to-end recovery over a real socket: the daemon must come back from an + /// eviction and STAY usable — accept work afterwards and keep heartbeating — + /// rather than merely emitting one Register and wedging. + /// + /// This is the daemon half of the server's + /// `daemon_evicted_by_a_dying_connection_recovers_on_its_next_heartbeat`. Together + /// they cover the whole loop: the controller rejects a heartbeat from a daemon it no + /// longer holds, and the daemon turns that rejection back into a working session on + /// the connection it already has. Two rejections in a row are exercised because an + /// eviction can recur (a flapping mesh path, a second stale-cleanup) and recovery + /// must not be one-shot. + #[tokio::test] + async fn daemon_recovers_and_keeps_serving_after_eviction() { + let (client, mut server) = ws_pair().await; + + let shutdown = std::future::pending::<()>(); + let daemon = serve(client, "test-daemon", 3600, HashMap::new(), shutdown); + + let controller = async { + ack_registration(&mut server).await; + + let nack = || { + serde_json::to_string(&Message::HeartbeatAck { + success: false, + reason: "daemon is not registered".to_string(), + }) + .unwrap() + }; + + // Evicted twice, with a recovery in between: recovery must be repeatable, + // not a one-shot latch. + let mut registers = 0; + for _ in 0..2 { + server.send(WsMessage::Text(nack())).await.unwrap(); + if matches!( + next_non_heartbeat(&mut server).await, + Some(Message::Register { .. }) + ) { + registers += 1; + } + } + + // Recovered daemons must still do WORK, not just re-register. Dispatch a + // command and require its output back on this same socket — proof the + // session is functional end to end, not merely present. + let exec = Message::ExecuteCommand { + request_id: "recovery-1".to_string(), + command: "echo recovered".to_string(), + timeout_secs: 30, + env: HashMap::new(), + cwd: None, + }; + server + .send(WsMessage::Text(serde_json::to_string(&exec).unwrap())) + .await + .unwrap(); + + let mut output = None; + // Skip heartbeats and any trailing Register while waiting for the result. + for _ in 0..5 { + match next_non_heartbeat(&mut server).await { + Some(Message::CommandOutput { + request_id, stdout, .. + }) => { + output = Some((request_id, stdout)); + break; + } + Some(_) => continue, + None => break, + } + } + + server.close(None).await.unwrap(); + (registers, output) + }; + + let (outcome, (registers, output)) = tokio::join!(daemon, controller); + assert_eq!(registers, 2, "daemon must re-register after EVERY eviction"); + let (request_id, stdout) = output.expect("recovered daemon never returned command output"); + assert_eq!(request_id, "recovery-1"); + assert_eq!(stdout.trim(), "recovered"); + // The session survived both evictions: it ended only because the controller + // closed, never because a rejection tore down a healthy socket. + assert_eq!(outcome.unwrap(), ServeOutcome::Disconnected); + } + + /// A successful heartbeat ack is not an error path and must be consumed + /// quietly: no re-registration, no disconnect. + #[tokio::test] + async fn successful_heartbeat_ack_is_ignored() { + let (client, mut server) = ws_pair().await; + + let shutdown = std::future::pending::<()>(); + let daemon = serve(client, "test-daemon", 3600, HashMap::new(), shutdown); + + let controller = async { + ack_registration(&mut server).await; + + let ack = Message::HeartbeatAck { + success: true, + reason: String::new(), + }; + server + .send(WsMessage::Text(serde_json::to_string(&ack).unwrap())) + .await + .unwrap(); + + // Give the daemon a chance to (wrongly) respond, then close. Only + // heartbeats should arrive: a success ack refreshes an internal deadline + // and nothing more — in particular it must not trigger a re-registration. + let replied = next_non_heartbeat(&mut server).await; + server.close(None).await.unwrap(); + replied + }; + + let (outcome, replied) = tokio::join!(daemon, controller); + assert!( + replied.is_none(), + "daemon replied to a successful heartbeat ack: {:?}", + replied + ); + assert_eq!(outcome.unwrap(), ServeOutcome::Disconnected); + } + /// `shutdown_signal()` must resolve when the process receives SIGTERM (what /// `kubectl delete pod` / `docker stop` send). Uses a real self-signal; unix /// only. This asserts the wiring, not the WebSocket behavior above. diff --git a/server/src/server.rs b/server/src/server.rs index 1fe5196..b9e937f 100644 --- a/server/src/server.rs +++ b/server/src/server.rs @@ -1,4 +1,3 @@ -use sandd_protocol::Message; use crate::registry::{DaemonConnection, DaemonRegistry}; use anyhow::{Context, Result}; use axum::{ @@ -12,6 +11,7 @@ use axum::{ Router, }; use futures_util::{SinkExt, StreamExt}; +use sandd_protocol::Message; use std::sync::Arc; use std::time::Duration; use tokio::sync::mpsc; @@ -160,6 +160,48 @@ async fn handle_websocket(ws: WebSocket, registry: Arc) { } } +/// Record a heartbeat. Returns whether the daemon was registered; `false` means it must +/// send `Register` again before it is reachable (it was reaped, or never registered on +/// this connection) and is the only way a heartbeat can fail. +/// +/// A daemon can be reaped while its socket is still healthy: mesh churn (DERP peer +/// reconfig, netmap propagation) stalls heartbeats past heartbeat_monitor's threshold +/// WITHOUT breaking TCP. `Register` is sent once per connection, and the daemon cannot +/// detect the eviction — its heartbeat writes keep succeeding, so it never trips its +/// own dead-connection signal. Left alone it stays invisible (no exec, no logs) until +/// the socket truly breaks, which its own heartbeats keep preventing. +/// +/// The daemon owns its metadata, so recovery is to fail the heartbeat and let it send +/// `Register` again, rather than reconstructing the entry server-side from a copy the +/// server would have to hold for every connection. +/// +/// Daemons predating `HeartbeatAck` effectively ignore it (it fails to deserialize) and +/// recover only when the connection eventually drops. They cannot be upgraded in place — +/// the binary is fetched from /releases/latest/ at instance boot — so this takes full effect +/// on instances provisioned after the daemon ships. +/// +/// Split out of handle_daemon_message so it is unit-testable: the parent needs a +/// SplitSink that cannot be constructed without a real socket. +fn handle_heartbeat(id: &str, registry: &Arc) -> bool { + match registry.get(id) { + Some(conn) => { + conn.update_heartbeat(); + debug!("Heartbeat from daemon {}", id); + true + } + // warn, not debug: reaching here means the reaper fired on a live socket. + // Recovering silently would hide the churn that caused it. + None => { + warn!( + "Heartbeat from unregistered daemon {} (reaped while connected); \ + asking it to register again", + id + ); + false + } + } +} + async fn handle_daemon_message( message: Message, daemon_id: &mut Option, @@ -199,9 +241,28 @@ async fn handle_daemon_message( Message::Heartbeat => { if let Some(ref id) = daemon_id { - if let Some(conn) = registry.get(id) { - conn.update_heartbeat(); - debug!("Heartbeat from daemon {}", id); + // Ack either way. On failure it tells the daemon to register again (only + // the daemon holds its metadata, so it re-sends Register rather than the + // server rebuilding the entry from a copy). On success it proves this + // controller is still PROCESSING messages — the daemon's own liveness + // check only sees whether its write reached a socket buffer, so a hung + // controller is indistinguishable from a healthy one without this. + let registered = handle_heartbeat(id, registry); + let ack = Message::HeartbeatAck { + success: registered, + reason: if registered { + String::new() + } else { + "daemon is not registered".to_string() + }, + }; + match serde_json::to_string(&ack) { + Ok(json) => { + if let Err(e) = ws_tx.send(axum::extract::ws::Message::Text(json)).await { + error!("Failed to ack heartbeat from daemon {}: {}", id, e); + } + } + Err(e) => error!("Failed to serialize heartbeat ack: {}", e), } } } @@ -304,9 +365,10 @@ async fn heartbeat_monitor(registry: Arc) { // 30s threshold against a 5s daemon heartbeat interval = ~6 missed beats before // reaping. That margin is deliberate: mesh churn (DERP peer reconfig, netmap // propagation) can stall heartbeats for tens of seconds WITHOUT the daemon being - // dead, and reaping a daemon whose socket is still open orphans it (its later - // heartbeats hit no registry entry and are ignored until the socket truly - // breaks). Detection is ~30-35s vs the old ~90-120s; clean disconnects are still + // dead. Reaping a daemon whose socket is still open no longer orphans it: its + // next heartbeat re-registers it (see handle_heartbeat), so a false reap costs + // one heartbeat interval of invisibility rather than lasting until the socket + // breaks. Detection is ~30-35s vs the old ~90-120s; clean disconnects are still // removed instantly on Close (see the remove() on the disconnect path above). let removed = registry.cleanup_stale(30); if removed > 0 { @@ -316,3 +378,180 @@ async fn heartbeat_monitor(registry: Arc) { info!("Active daemons: {} ", registry.count()); } } + +#[cfg(test)] +mod tests { + use super::*; + use sandd_protocol::DaemonMetadata; + use std::collections::HashMap; + + fn test_metadata() -> DaemonMetadata { + let mut labels = HashMap::new(); + labels.insert("env".to_string(), "prod".to_string()); + DaemonMetadata { + hostname: "gpu-box".to_string(), + platform: "linux".to_string(), + arch: "x86_64".to_string(), + version: "0.1.0".to_string(), + labels, + } + } + + fn registered(id: &str) -> (Arc, mpsc::UnboundedSender) { + let registry = Arc::new(DaemonRegistry::new()); + let (tx, _rx) = mpsc::unbounded_channel(); + registry.register(DaemonConnection::new( + id.to_string(), + test_metadata(), + tx.clone(), + )); + (registry, tx) + } + + // A daemon reaped while its socket stayed open must be ASKED to register again. + // Before the fix the beat was dropped and the daemon stayed invisible forever: + // Register is sent once per connection and the daemon cannot detect the eviction, + // so nothing ever prompted it. + #[test] + fn heartbeat_from_reaped_daemon_requires_register() { + let (registry, _tx) = registered("daemon-1"); + + // Reaper evicts it (heartbeats stalled past the threshold), socket still open. + registry.remove("daemon-1"); + assert_eq!(registry.count(), 0); + + assert!(!handle_heartbeat("daemon-1", ®istry)); + // The server does NOT fabricate an entry — the daemon owns its metadata and + // re-sends it, so the restored entry is faithful rather than a stale copy. + assert_eq!(registry.count(), 0); + } + + // A daemon that never registered on this connection is treated the same: ask it to + // register. No special case needed. + #[test] + fn heartbeat_from_unknown_daemon_requires_register() { + let registry = Arc::new(DaemonRegistry::new()); + + assert!(!handle_heartbeat("daemon-1", ®istry)); + } + + // The normal path: a registered daemon's heartbeat refreshes its timestamp so the + // reaper leaves it alone. + #[test] + fn heartbeat_refreshes_existing_daemon() { + let (registry, _tx) = registered("daemon-1"); + + assert!(handle_heartbeat("daemon-1", ®istry)); + assert_eq!(registry.count(), 1); + assert_eq!( + registry.get("daemon-1").unwrap().seconds_since_heartbeat(), + 0 + ); + } + + // The daemon's re-registration must land on the LIVE connection. Re-registering is + // an ordinary Register, so it replaces the entry — which is correct here because it + // arrives on the socket the daemon is actually using. + #[test] + fn reregister_routes_to_the_registering_connection() { + let (registry, mut stale_rx) = { + let registry = Arc::new(DaemonRegistry::new()); + let (tx, rx) = mpsc::unbounded_channel(); + registry.register(DaemonConnection::new( + "daemon-1".to_string(), + test_metadata(), + tx, + )); + (registry, rx) + }; + + // The daemon re-registers, carrying its own connection's channel. + let (new_tx, mut new_rx) = mpsc::unbounded_channel(); + registry.register(DaemonConnection::new( + "daemon-1".to_string(), + test_metadata(), + new_tx, + )); + + assert_eq!(registry.count(), 1); + registry + .get("daemon-1") + .unwrap() + .send_message(Message::Heartbeat) + .unwrap(); + assert!( + new_rx.try_recv().is_ok(), + "must route to the re-registered connection" + ); + assert!( + stale_rx.try_recv().is_err(), + "must not route to the old connection" + ); + } + + // A LIVE daemon evicted by a DYING one must recover. This is the stale-remove race: + // handle_websocket's cleanup is `registry.remove(&id)`, keyed on the id with no + // check of WHICH connection is stored there, so when two connections for one daemon + // overlap — a half-open socket that has not errored yet, plus the reconnect the + // daemon made over a working path — the old task's cleanup deletes the NEW task's + // entry. Reproduced here in the order the tasks actually interleave. + // + // The race window is still open (a ptr-identity-aware remove would close it); what + // this pins down is that it is no longer PERMANENT. Before HeartbeatAck the live + // daemon stayed invisible for the life of its connection — Register is sent once, it + // cannot observe the eviction (its writes still succeed), and its own heartbeats keep + // the socket from breaking. Now its next heartbeat is rejected and it re-registers, + // so the damage is bounded to one heartbeat interval. + #[test] + fn daemon_evicted_by_a_dying_connection_recovers_on_its_next_heartbeat() { + // The daemon's original connection. + let (registry, _stale_tx) = registered("daemon-1"); + + // Its mesh path half-dies. TCP does not fail fast, so the old task is still + // parked in ws_rx.next() while the daemon reconnects over a working path and + // registers again — an ordinary Register, which overwrites the entry. + let (live_tx, mut live_rx) = mpsc::unbounded_channel(); + registry.register(DaemonConnection::new( + "daemon-1".to_string(), + test_metadata(), + live_tx.clone(), + )); + assert_eq!(registry.count(), 1); + + // The old socket finally errors and its task runs the cleanup on the way out. + // Keyed only by id, it evicts the LIVE connection that replaced it. + registry.remove("daemon-1"); + assert_eq!( + registry.count(), + 0, + "the dying connection's cleanup evicted the live entry (the race being modeled)" + ); + + // The live daemon is now invisible even though its socket is fine — the exact + // state that used to persist forever. Its next heartbeat is rejected... + assert!( + !handle_heartbeat("daemon-1", ®istry), + "an evicted daemon's heartbeat must be rejected so it knows to re-register" + ); + + // ...so it re-registers on that same healthy socket, carrying its own channel. + registry.register(DaemonConnection::new( + "daemon-1".to_string(), + test_metadata(), + live_tx, + )); + + // Recovered: visible again, and reachable on the connection it is actually using. + assert_eq!(registry.count(), 1); + assert!(handle_heartbeat("daemon-1", ®istry)); + registry + .get("daemon-1") + .unwrap() + .send_message(Message::Heartbeat) + .unwrap(); + assert!( + live_rx.try_recv().is_ok(), + "work must route to the recovered daemon's live connection" + ); + } +}