diff --git a/AGENTS.md b/AGENTS.md index 7af04fa..69feb92 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -75,10 +75,11 @@ bench/build_transfer.sh rust-rNN-description-$(git rev-parse --short HEAD) --nod The cold first build still downloads the Rust base and compiles dependencies; judge the workflow by its warm edit/rebuild time. A dependency change is a -legitimate one-time cold build. GitHub Actions uses `Swatinem/rust-cache` for -dependency artifacts (including failed runs), and Drone fans Rust, Go, and the -GPU-free protocol suite out in parallel. If a no-dependency CI change still -does a cold compile, inspect the cache action before accepting the delay. +legitimate one-time cold build. Drone fetches dependencies once, then fans +Clippy and Rust tests into isolated target directories while Go, Python, and +Compose run beside them. Build #205 completed the full PR gate in 58s versus +69s for the prior serial Rust lane. If a no-dependency change regresses that +loop, inspect the step timings and cache/export path before accepting it. `Cargo.toml` deliberately limits the crate package to Rust sources, examples, compatibility fixtures, and Cargo manifests. This keeps edits under `bench/` @@ -211,6 +212,19 @@ fenced observe-only mode. Do not reconnect on each subsequent event: only an authoritative all-blocks clear or an attested incarnation change may restore trust without a complete replay from zero. +`mini-dynamo-snapshot-companion` is the standalone one-engine composition. It +must remain a separate binary and `Dockerfile.companion`; build the ordinary LB +with `--bin mini-dynamo` so companion changes do not add a second release relink +to the serving-image loop. Serve mode must preflight the session secret, digest +secret, and authenticated incarnation file before any TCP bind, ZMQ connect, or +UDS publication. Invalid attestation refreshes immediately remove authority; +logs and metrics expose only closed reason labels. The container publishes no +metrics port: loopback inside a bridge container is not host-scrapable. Add a +separate, permission-isolated metrics UDS and Caddy/Prometheus wiring before a +node06 rollout; never put Caddy in the snapshot/session authority group. The +executable is not ready for node06 until the host-side attestation provisioner +and production Compose wiring exist and pass the dual-domain validator. + `snapshot_reconnect` is the LB-side owner around the consumer. Normal attempts are serial; only an explicit bounded replacement may overlap a second session. Validate the trusted socket parent on every connect, use a fresh OS-random @@ -314,13 +328,13 @@ export BENCH_TOKEN=$(grep -o 'Bearer [A-Za-z0-9_-]*' /etc/caddy/Caddyfile | head - Drone intentionally runs once per PR (and once after merge on `main`), not again for every feature-branch push. Its Rust lane omits a redundant release build: the local pre-push gate builds release, and Drone's post-merge main - pipeline builds and publishes the release container to GHCR (the Actions - GITHUB_TOKEN is denied write on that package — issue #39; the `publish-image` - step reads `ghcr_username`/`ghcr_token` Drone secrets instead). GitHub - Actions keeps format, Clippy, test, and deploy-config validation only. The - Drone image build runs in a fresh dind daemon, so it pays a cold cargo - release build on every publish; node06 deploys should keep using the warm - local `bench/build_transfer.sh` path. + pipeline builds and publishes the release container to GHCR. The + `publish-image` step reads `ghcr_username`/`ghcr_token` only on a main push; + pull-request containers never receive those secrets. GitHub Actions is not a + second CI system for this repository. Drone's PR quality budget is currently + about 58s; inspect and record regressions instead of adding duplicate gates. + Node06 deploys should keep using the warm local `bench/build_transfer.sh` + path because it reuses the development BuildKit cache. - Build the LB locally and stream it to node06 when the local amd64 Docker cache is warm. A typical warm LB transfer is seconds and avoids consuming node06's scarce 8-9GiB available host memory: diff --git a/Dockerfile b/Dockerfile index b58aef2..7b99941 100644 --- a/Dockerfile +++ b/Dockerfile @@ -4,7 +4,8 @@ WORKDIR /src COPY . . RUN --mount=type=cache,id=mini-dynamo-cargo-registry,target=/usr/local/cargo/registry \ --mount=type=cache,id=mini-dynamo-target,target=/src/target \ - cargo build --release --locked && cp target/release/mini-dynamo /mini-dynamo + cargo build --release --locked --bin mini-dynamo \ + && cp target/release/mini-dynamo /mini-dynamo FROM gcr.io/distroless/cc-debian12 LABEL org.opencontainers.image.source="https://github.com/helixml/mini-dynamo" diff --git a/Dockerfile.companion b/Dockerfile.companion new file mode 100644 index 0000000..6e304ed --- /dev/null +++ b/Dockerfile.companion @@ -0,0 +1,13 @@ +# syntax=docker/dockerfile:1.7 +FROM rust:1.95-bookworm AS build +WORKDIR /src +COPY . . +RUN --mount=type=cache,id=mini-dynamo-cargo-registry,target=/usr/local/cargo/registry \ + --mount=type=cache,id=mini-dynamo-target,target=/src/target \ + cargo build --release --locked --bin mini-dynamo-snapshot-companion \ + && cp target/release/mini-dynamo-snapshot-companion /mini-dynamo-snapshot-companion + +FROM gcr.io/distroless/cc-debian12 +LABEL org.opencontainers.image.source="https://github.com/helixml/mini-dynamo" +COPY --from=build /mini-dynamo-snapshot-companion /mini-dynamo-snapshot-companion +ENTRYPOINT ["/mini-dynamo-snapshot-companion"] diff --git a/EXPERIMENTS.md b/EXPERIMENTS.md index 4e9c412..eea54d1 100644 --- a/EXPERIMENTS.md +++ b/EXPERIMENTS.md @@ -3606,3 +3606,33 @@ were identical before and after; mini-dynamo remained 2/2 healthy with both non-writable under `/tmp` on node06, but no tool container, CUDA context, LB recreate, or benchmark ran. The next gate remains the separately acknowledged 1MiB/two-GPU scout during a low-traffic window. + +## 2026-08-13 — r46 standalone single-engine snapshot companion + +The previously separate owner, source, producer, supervisor, and metrics pieces +now compose into the off-by-default `mini-dynamo-snapshot-companion` executable. +Serve mode accepts exactly one live/replay endpoint pair and one group geometry. +Before binding metrics, connecting ZMQ, or publishing its Unix socket, it +hardens and loads three distinct protected files: the session secret, digest +secret, and an HMAC-authenticated engine-incarnation envelope. The refresh +watch removes authority immediately on any filesystem, schema, field, or MAC +failure and restores it only after a valid changed envelope. Owner telemetry is +closed-label and content-free, including the stable `replay_too_large` +observe-only state required by long-lived node06 engine A. + +One absolute shutdown deadline covers the owner, snapshot runtime, metrics +server, and attestation watcher; tasks that miss it are aborted. The ordinary +LB Dockerfile now names only `--bin mini-dynamo`, while the companion has a +dedicated `Dockerfile.companion`. With one shared BuildKit target cache, the +measured source builds were 42.3s for the LB followed by 9.7s for the companion; +no-op warm builds were 2.16s and 3.38s. Final images were 14.27MB and 11.50MB. +This preserves the fast LB-only loop instead of paying the observed extra +roughly 41-second companion relink for every router image. + +The rebased gate passed 298 library tests, 31 integration/E2E tests, strict +Clippy, formatting, release builds, off-mode executable/container smoke, and +the dual-domain offline Compose validator. No production Compose, image push, +or node06 state changed. The executable is intentionally not deployment-ready +yet: a bounded host provisioner must derive and atomically write the authenticated +incarnation from current engine metadata before either per-engine service can be +enabled. diff --git a/ROADMAP.md b/ROADMAP.md index 0423d86..146b36b 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -603,10 +603,13 @@ node06). The design rationale for each lives in DESIGN.md. two isolated processes and authority domains: distinct companions, UIDs, tmpfs directories, sockets, secrets, clients, and readiness checks per engine. Its default profile renders zero services, and static fault projection proves - one failed pair cannot publish or substitute through its healthy peer. Next - wire one authenticated owner/socket per engine into an executable and run at - least 100,000 revision-stable shadow comparisons before placement can consume - it. + one failed pair cannot publish or substitute through its healthy peer. A + dedicated off-by-default executable now composes exactly one authenticated + incarnation watch, ZMQ owner, long-lived source, snapshot socket, and bounded + metrics endpoint. It is built separately from the LB so ordinary router edits + do not relink both binaries. Next add the host-side authenticated-attestation + provisioner and production-shaped Compose wiring, then run at least 100,000 + revision-stable shadow comparisons before placement can consume it. The off-by-default library runtime now composes typed config, hardened secret loading, bind-last safe socket publication, the bounded supervisor, producer, @@ -618,9 +621,11 @@ node06). The design rationale for each lives in DESIGN.md. The coordinator delegates one absolute snapshot-phase timeout and a resettable tail-idle/write timeout to the producer; the supervisor retains the two-client cap and immediate shutdown - cancellation without imposing a total lifetime on healthy progress. Next - inject the completed process owner and its authenticated engine-incarnation - watch into each isolated executable instance. + cancellation without imposing a total lifetime on healthy progress. The + standalone process now injects the completed owner and a refreshable, + HMAC-authenticated engine-incarnation watch; any invalid refresh fences + authority immediately. Deployment remains blocked on generating that manifest + from current engine metadata without exposing either digest key or identity. A true offline public-stack harness now proves initial publication, live store/remove, rolling handoff, LB owner restart, companion shutdown/socket diff --git a/src/bin/mini-dynamo-snapshot-companion.rs b/src/bin/mini-dynamo-snapshot-companion.rs new file mode 100644 index 0000000..c1f9119 --- /dev/null +++ b/src/bin/mini-dynamo-snapshot-companion.rs @@ -0,0 +1,91 @@ +use std::{ffi::OsString, path::PathBuf}; + +use anyhow::{Context, bail}; +use mini_dynamo::companion_service::{SingleEngineCompanionConfig, run_single_engine_companion}; +use mini_dynamo::snapshot_socket_path::validate_published_socket; +use tokio::sync::watch; +use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + match command(std::env::args_os().skip(1))? { + Command::Run => run().await, + Command::Healthcheck(path) => { + validate_published_socket(&path).context("snapshot companion healthcheck failed")?; + Ok(()) + } + } +} + +enum Command { + Run, + Healthcheck(PathBuf), +} + +fn command(mut arguments: impl Iterator) -> anyhow::Result { + let Some(first) = arguments.next() else { + return Ok(Command::Run); + }; + if first != "healthcheck" { + bail!("invalid snapshot companion command"); + } + let Some(path) = arguments.next() else { + bail!("snapshot companion healthcheck requires one socket path"); + }; + if arguments.next().is_some() { + bail!("snapshot companion healthcheck accepts exactly one socket path"); + } + Ok(Command::Healthcheck(PathBuf::from(path))) +} + +async fn run() -> anyhow::Result<()> { + tracing_subscriber::registry() + .with(EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"))) + .with(tracing_subscriber::fmt::layer().json()) + .init(); + + let config = SingleEngineCompanionConfig::from_env() + .context("invalid standalone snapshot companion configuration")?; + tracing::info!(enabled = config.enabled(), "snapshot companion starting"); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let service = run_single_engine_companion(config, shutdown_rx); + tokio::pin!(service); + let result = tokio::select! { + biased; + result = &mut service => result, + () = shutdown_signal() => { + let _ = shutdown_tx.send(true); + service.await + } + }; + match result { + Ok(report) => { + tracing::info!(report = ?report, "snapshot companion stopped"); + Ok(()) + } + Err(error) => { + tracing::error!(reason = error.reason(), "snapshot companion failed"); + Err(anyhow::anyhow!( + "snapshot companion failed: {}", + error.reason() + )) + } + } +} + +async fn shutdown_signal() { + #[cfg(unix)] + { + let mut terminate = + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("install SIGTERM handler"); + tokio::select! { + result = tokio::signal::ctrl_c() => { let _ = result; }, + _ = terminate.recv() => {}, + } + } + #[cfg(not(unix))] + { + let _ = tokio::signal::ctrl_c().await; + } +} diff --git a/src/companion_attestation.rs b/src/companion_attestation.rs new file mode 100644 index 0000000..ed3fb5e --- /dev/null +++ b/src/companion_attestation.rs @@ -0,0 +1,453 @@ +//! Authenticated, refreshable engine-incarnation input for the companion. +//! +//! The manifest is a bounded JSON envelope. Its HMAC covers a canonical +//! `MessagePack` encoding under a dedicated domain and uses the digest secret as +//! key material. Filesystem protection is necessary but not sufficient: a +//! syntactically valid manifest without the correct MAC never grants authority. + +use std::{path::Path, time::Duration}; + +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use thiserror::Error; +use tokio::sync::watch; + +use crate::{ + kv_snapshot::EngineIncarnation, + snapshot_secret_file::{ + SnapshotSecretFileError, SnapshotSecretFilePolicy, load_snapshot_control_file, + }, +}; + +const SECRET_BYTES: usize = 32; +const HMAC_BLOCK_BYTES: usize = 64; +const MAX_ATTESTATION_BYTES: usize = 16 * 1024; +const MAX_COMPONENT_BYTES: usize = 512; +const SCHEMA_VERSION: u16 = 1; +const DOMAIN: &[u8] = b"mini-dynamo:engine-incarnation-attestation:v1\0"; + +/// Directly owned digest key loaded from hardened storage. +pub struct CompanionDigestSecret([u8; SECRET_BYTES]); + +impl CompanionDigestSecret { + #[must_use] + pub const fn as_bytes(&self) -> &[u8; SECRET_BYTES] { + &self.0 + } +} + +impl std::fmt::Debug for CompanionDigestSecret { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("CompanionDigestSecret([REDACTED])") + } +} + +impl Drop for CompanionDigestSecret { + fn drop(&mut self) { + self.0.fill(0); + } +} + +#[derive(Clone, Debug, Error, Eq, PartialEq)] +pub enum CompanionAttestationError { + #[error("companion protected file validation failed")] + File(#[from] SnapshotSecretFileError), + #[error("companion digest secret has an invalid length")] + SecretLength, + #[error("engine incarnation attestation is malformed")] + Malformed, + #[error("engine incarnation attestation is unsupported")] + Unsupported, + #[error("engine incarnation attestation has invalid fields")] + InvalidIncarnation, + #[error("engine incarnation attestation authentication failed")] + Authentication, +} + +impl CompanionAttestationError { + #[must_use] + pub const fn reason(&self) -> &'static str { + match self { + Self::File(error) => error.reason(), + Self::SecretLength => "secret_length", + Self::Malformed => "malformed", + Self::Unsupported => "unsupported", + Self::InvalidIncarnation => "invalid_incarnation", + Self::Authentication => "authentication_failed", + } + } +} + +#[derive(Serialize)] +struct AttestationPayload<'a> { + schema_version: u16, + engine_incarnation: &'a EngineIncarnation, +} + +#[derive(Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +struct AttestationEnvelope { + schema_version: u16, + engine_incarnation: EngineIncarnation, + mac_sha256: Vec, +} + +/// Load an exact raw 256-bit digest secret from hardened storage. +/// +/// # Errors +/// +/// Rejects every file-policy failure and every size other than 32 raw bytes. +pub fn load_companion_digest_secret( + path: &Path, + policy: SnapshotSecretFilePolicy, +) -> Result { + let mut raw = load_snapshot_control_file(path, policy, SECRET_BYTES)?; + let converted = <[u8; SECRET_BYTES]>::try_from(raw.as_slice()); + raw.fill(0); + converted + .map(CompanionDigestSecret) + .map_err(|_| CompanionAttestationError::SecretLength) +} + +/// Load and authenticate one incarnation manifest. +/// +/// # Errors +/// +/// Fails closed for filesystem, size, JSON, schema, field, or MAC violations. +pub fn load_authenticated_engine_incarnation( + path: &Path, + policy: SnapshotSecretFilePolicy, + secret: &CompanionDigestSecret, +) -> Result { + let mut bytes = load_snapshot_control_file(path, policy, MAX_ATTESTATION_BYTES)?; + let decoded = serde_json::from_slice::(&bytes) + .map_err(|_| CompanionAttestationError::Malformed); + bytes.fill(0); + let envelope = decoded?; + if envelope.schema_version != SCHEMA_VERSION { + return Err(CompanionAttestationError::Unsupported); + } + validate_incarnation(&envelope.engine_incarnation)?; + if envelope.mac_sha256.len() != SECRET_BYTES { + return Err(CompanionAttestationError::Authentication); + } + let expected = authenticate(&envelope.engine_incarnation, secret)?; + if !constant_time_equal(&expected, &envelope.mac_sha256) { + return Err(CompanionAttestationError::Authentication); + } + Ok(envelope.engine_incarnation) +} + +/// Encode the canonical authenticated JSON envelope used by the file watcher. +/// This helper makes provisioning independent of private serialization details. +/// +/// # Errors +/// +/// Rejects invalid incarnation fields or serialization failure. +pub fn encode_authenticated_engine_incarnation( + incarnation: &EngineIncarnation, + secret: &CompanionDigestSecret, +) -> Result, CompanionAttestationError> { + validate_incarnation(incarnation)?; + let envelope = AttestationEnvelope { + schema_version: SCHEMA_VERSION, + engine_incarnation: incarnation.clone(), + mac_sha256: authenticate(incarnation, secret)?.to_vec(), + }; + serde_json::to_vec(&envelope).map_err(|_| CompanionAttestationError::Malformed) +} + +/// Poll an authenticated manifest until shutdown. Every invalid refresh sends +/// explicit authority loss; a later valid manifest restores authority. +pub async fn watch_authenticated_engine_incarnation( + path: &Path, + policy: SnapshotSecretFilePolicy, + secret: CompanionDigestSecret, + refresh: Duration, + authority: &watch::Sender>, + mut shutdown: watch::Receiver, +) { + let mut ticker = tokio::time::interval(refresh); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut last_failure = None; + loop { + tokio::select! { + biased; + () = wait_for_shutdown(&mut shutdown) => return, + _ = ticker.tick() => { + let refreshed = load_authenticated_engine_incarnation(path, policy, &secret); + let next = match refreshed { + Ok(incarnation) => { + last_failure = None; + Some(incarnation) + } + Err(error) => { + let reason = error.reason(); + if last_failure != Some(reason) { + tracing::warn!(reason, "engine incarnation authority unavailable"); + last_failure = Some(reason); + } + None + } + }; + authority.send_if_modified(|current| { + if *current == next { + false + } else { + *current = next; + true + } + }); + } + } + } +} + +fn validate_incarnation(incarnation: &EngineIncarnation) -> Result<(), CompanionAttestationError> { + let components = [ + incarnation.engine_id.as_bytes(), + incarnation.model_revision.as_bytes(), + incarnation.image_digest.as_bytes(), + ]; + if components + .iter() + .any(|component| component.is_empty() || component.len() > MAX_COMPONENT_BYTES) + || incarnation.process_started_unix_ns == 0 + || incarnation.attestation_sha256.len() != SECRET_BYTES + { + return Err(CompanionAttestationError::InvalidIncarnation); + } + Ok(()) +} + +fn authenticate( + incarnation: &EngineIncarnation, + secret: &CompanionDigestSecret, +) -> Result<[u8; SECRET_BYTES], CompanionAttestationError> { + let payload = rmp_serde::to_vec_named(&AttestationPayload { + schema_version: SCHEMA_VERSION, + engine_incarnation: incarnation, + }) + .map_err(|_| CompanionAttestationError::Malformed)?; + let mut hmac = HmacSha256::new(secret.as_bytes()); + hmac.update(DOMAIN); + hmac.update(&payload); + Ok(hmac.finalize()) +} + +fn constant_time_equal(expected: &[u8; SECRET_BYTES], actual: &[u8]) -> bool { + if actual.len() != SECRET_BYTES { + return false; + } + expected + .iter() + .zip(actual) + .fold(0_u8, |difference, (left, right)| { + difference | (left ^ right) + }) + == 0 +} + +struct HmacSha256 { + inner: Sha256, + outer_pad: [u8; HMAC_BLOCK_BYTES], +} + +impl HmacSha256 { + fn new(key: &[u8]) -> Self { + let mut normalized = [0_u8; HMAC_BLOCK_BYTES]; + normalized[..key.len()].copy_from_slice(key); + let mut inner_pad = [0x36_u8; HMAC_BLOCK_BYTES]; + let mut outer_pad = [0x5c_u8; HMAC_BLOCK_BYTES]; + for ((inner, outer), key_byte) in inner_pad.iter_mut().zip(&mut outer_pad).zip(normalized) { + *inner ^= key_byte; + *outer ^= key_byte; + } + let mut inner = Sha256::new(); + inner.update(inner_pad); + normalized.fill(0); + inner_pad.fill(0); + Self { inner, outer_pad } + } + + fn update(&mut self, bytes: &[u8]) { + self.inner.update(bytes); + } + + fn finalize(mut self) -> [u8; SECRET_BYTES] { + let inner = self.inner.finalize(); + let mut outer = Sha256::new(); + outer.update(self.outer_pad); + outer.update(inner); + self.outer_pad.fill(0); + outer.finalize().into() + } +} + +async fn wait_for_shutdown(shutdown: &mut watch::Receiver) { + loop { + if *shutdown.borrow() || shutdown.changed().await.is_err() { + return; + } + } +} + +#[cfg(test)] +mod tests { + use std::{ + fs, + os::unix::fs::{MetadataExt, PermissionsExt}, + sync::atomic::{AtomicU64, Ordering}, + }; + + use tokio::time::timeout; + + use super::*; + + static TEST_ID: AtomicU64 = AtomicU64::new(1); + + struct TestFiles { + directory: std::path::PathBuf, + secret: std::path::PathBuf, + attestation: std::path::PathBuf, + owner: u32, + } + + impl TestFiles { + fn new() -> Self { + let id = TEST_ID.fetch_add(1, Ordering::Relaxed); + let directory = std::env::temp_dir().join(format!( + "mini-dynamo-attestation-{}-{id}", + std::process::id() + )); + fs::create_dir(&directory).unwrap(); + fs::set_permissions(&directory, fs::Permissions::from_mode(0o700)).unwrap(); + let owner = fs::metadata(&directory).unwrap().uid(); + let secret = directory.join("digest-secret"); + fs::write(&secret, [0x51; SECRET_BYTES]).unwrap(); + fs::set_permissions(&secret, fs::Permissions::from_mode(0o600)).unwrap(); + Self { + attestation: directory.join("attestation.json"), + directory, + secret, + owner, + } + } + + fn policy(&self) -> SnapshotSecretFilePolicy { + SnapshotSecretFilePolicy { + expected_owner_uid: self.owner, + } + } + + fn write_attestation(&self, bytes: &[u8]) { + fs::write(&self.attestation, bytes).unwrap(); + fs::set_permissions(&self.attestation, fs::Permissions::from_mode(0o600)).unwrap(); + } + } + + impl Drop for TestFiles { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.directory); + } + } + + fn incarnation(started: u64) -> EngineIncarnation { + EngineIncarnation { + engine_id: "engine-a".to_owned(), + model_revision: "revision".to_owned(), + image_digest: "sha256:image".to_owned(), + process_started_unix_ns: started, + attestation_sha256: vec![7; SECRET_BYTES], + } + } + + #[test] + fn hardened_authenticated_manifest_round_trips_and_rejects_tamper() { + let files = TestFiles::new(); + let secret = load_companion_digest_secret(&files.secret, files.policy()).unwrap(); + assert_eq!(format!("{secret:?}"), "CompanionDigestSecret([REDACTED])"); + let encoded = encode_authenticated_engine_incarnation(&incarnation(42), &secret).unwrap(); + files.write_attestation(&encoded); + assert_eq!( + load_authenticated_engine_incarnation(&files.attestation, files.policy(), &secret) + .unwrap(), + incarnation(42) + ); + + let mut envelope: serde_json::Value = serde_json::from_slice(&encoded).unwrap(); + envelope["engine_incarnation"]["process_started_unix_ns"] = 43.into(); + files.write_attestation(&serde_json::to_vec(&envelope).unwrap()); + assert_eq!( + load_authenticated_engine_incarnation(&files.attestation, files.policy(), &secret), + Err(CompanionAttestationError::Authentication) + ); + } + + #[test] + fn protected_inputs_reject_unsafe_permissions_and_secret_sizes() { + let files = TestFiles::new(); + fs::set_permissions(&files.secret, fs::Permissions::from_mode(0o622)).unwrap(); + assert!(matches!( + load_companion_digest_secret(&files.secret, files.policy()), + Err(CompanionAttestationError::File( + SnapshotSecretFileError::UnsafePermissions + )) + )); + fs::set_permissions(&files.secret, fs::Permissions::from_mode(0o600)).unwrap(); + fs::write(&files.secret, [0x51; SECRET_BYTES - 1]).unwrap(); + assert!(load_companion_digest_secret(&files.secret, files.policy()).is_err()); + } + + #[tokio::test] + async fn watcher_fences_invalid_refresh_and_restores_authenticated_authority() { + let files = TestFiles::new(); + let secret = load_companion_digest_secret(&files.secret, files.policy()).unwrap(); + let good = encode_authenticated_engine_incarnation(&incarnation(42), &secret).unwrap(); + files.write_attestation(&good); + let (authority_tx, mut authority_rx) = watch::channel(Some(incarnation(42))); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let path = files.attestation.clone(); + let policy = files.policy(); + let task = tokio::spawn(async move { + watch_authenticated_engine_incarnation( + &path, + policy, + secret, + Duration::from_millis(10), + &authority_tx, + shutdown_rx, + ) + .await; + }); + + files.write_attestation(b"not authenticated"); + timeout(Duration::from_secs(1), async { + loop { + authority_rx.changed().await.unwrap(); + if authority_rx.borrow_and_update().is_none() { + break; + } + } + }) + .await + .unwrap(); + + let signing_secret = load_companion_digest_secret(&files.secret, files.policy()).unwrap(); + let refreshed = + encode_authenticated_engine_incarnation(&incarnation(43), &signing_secret).unwrap(); + files.write_attestation(&refreshed); + timeout(Duration::from_secs(1), async { + loop { + authority_rx.changed().await.unwrap(); + if authority_rx.borrow_and_update().as_ref() == Some(&incarnation(43)) { + break; + } + } + }) + .await + .unwrap(); + shutdown_tx.send(true).unwrap(); + task.await.unwrap(); + } +} diff --git a/src/companion_service.rs b/src/companion_service.rs new file mode 100644 index 0000000..cc51801 --- /dev/null +++ b/src/companion_service.rs @@ -0,0 +1,1159 @@ +//! Standalone, off-by-default single-engine snapshot companion composition. + +use std::{ + env, fmt, + net::{IpAddr, SocketAddr}, + path::PathBuf, + sync::Arc, + time::Duration, +}; + +use axum::{Router, body::Body, http::Response, routing::get}; +use prometheus::{CounterVec, Encoder, Gauge, HistogramVec, Opts, Registry, TextEncoder}; +use thiserror::Error; +use tokio::{ + net::TcpListener, + sync::watch, + task::{JoinError, JoinHandle}, + time::{Instant, timeout_at}, +}; + +use crate::{ + companion_attestation::{ + CompanionAttestationError, load_authenticated_engine_incarnation, + load_companion_digest_secret, watch_authenticated_engine_incarnation, + }, + companion_config::{ + SnapshotCompanionConfig, SnapshotCompanionConfigError, SnapshotCompanionMode, + }, + companion_index_owner::{ + CompanionIndexOwner, CompanionIndexOwnerConfig, CompanionIndexOwnerError, + CompanionIndexOwnerEvent, CompanionIndexOwnerObserver, CompanionIndexOwnerRebuildReason, + CompanionIndexOwnerReplayKind, CompanionIndexOwnerReplayOutcome, CompanionIndexOwnerReport, + ZmqCompanionKvEventConnector, + }, + companion_index_source::{ + CompanionIndexSource, CompanionIndexSourceConfig, CompanionIndexSourceError, + }, + companion_runtime::{ + SnapshotCompanionRunReport, SnapshotCompanionRuntimeError, run_snapshot_companion, + }, + digest_index::DigestIndexLimits, + kv_snapshot::{AttentionKind, GroupDisposition, GroupMetadata, SnapshotLimits}, + kv_transport::KvTransportConfig, + kv_wire::KvWireLimits, + snapshot_producer::SnapshotProducerSource, + snapshot_secret_file::{SnapshotSecretFileError, SnapshotSecretFilePolicy}, + snapshot_supervisor::MAX_ACTIVE_SNAPSHOT_CLIENTS, +}; + +const MAX_PATH_BYTES: usize = 4_096; +const MAX_REPLAY_BATCHES: usize = 100_000; +const MAX_BLOCK_SIZE: usize = 1_048_576; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ConfigAttentionKind { + Full, + Mla, + SinkFull, +} + +impl ConfigAttentionKind { + const fn wire(self) -> AttentionKind { + match self { + Self::Full => AttentionKind::FullAttention, + Self::Mla => AttentionKind::MlaAttention, + Self::SinkFull => AttentionKind::SinkFullAttention, + } + } +} + +#[derive(Clone)] +pub struct SingleEngineCompanionConfig { + pub snapshot: SnapshotCompanionConfig, + digest_secret_path: Option, + attestation_path: Option, + metrics_bind: SocketAddr, + attestation_refresh: Duration, + transport: Option, + owner: CompanionIndexOwnerConfig, + group: Option, +} + +impl fmt::Debug for SingleEngineCompanionConfig { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("SingleEngineCompanionConfig") + .field("snapshot", &self.snapshot) + .field( + "digest_secret_path", + &self.digest_secret_path.as_ref().map(|_| "[REDACTED]"), + ) + .field( + "attestation_path", + &self.attestation_path.as_ref().map(|_| "[REDACTED]"), + ) + .field("metrics_loopback", &self.metrics_bind.ip().is_loopback()) + .field("attestation_refresh", &self.attestation_refresh) + .field("transport_configured", &self.transport.is_some()) + .field("owner", &self.owner) + .field("group_configured", &self.group.is_some()) + .finish() + } +} + +impl SingleEngineCompanionConfig { + /// Load the standalone companion environment contract. + /// + /// # Errors + /// + /// Returns content-free typed configuration errors. Off mode does not read + /// or require any companion executable setting. + pub fn from_env() -> Result { + Self::from_lookup(|key| env::var(key).ok()) + } + + #[allow(clippy::too_many_lines)] + /// Deterministic lookup-based constructor for tests and embedders. + /// + /// # Errors + /// + /// Rejects invalid base settings, multi-engine sources, missing protected + /// paths/geometry, non-loopback metrics, and out-of-bound runtime values. + pub fn from_lookup( + mut get: impl FnMut(&str) -> Option, + ) -> Result { + let snapshot = SnapshotCompanionConfig::from_lookup(&mut get)?; + if snapshot.mode == SnapshotCompanionMode::Off { + return Ok(Self { + snapshot, + digest_secret_path: None, + attestation_path: None, + metrics_bind: parse_loopback("127.0.0.1:9091")?, + attestation_refresh: Duration::from_secs(1), + transport: None, + owner: CompanionIndexOwnerConfig { + replay_limit: 10_000, + reconnect_min: Duration::from_millis(250), + reconnect_max: Duration::from_secs(5), + }, + group: None, + }); + } + let metrics_bind = parse_loopback( + get("DS4_SNAPSHOT_METRICS_BIND") + .as_deref() + .unwrap_or("127.0.0.1:9091"), + )?; + let attestation_refresh = parse_duration( + &mut get, + "DS4_SNAPSHOT_ATTESTATION_REFRESH_MS", + 1_000, + 50, + 60_000, + )?; + let owner = CompanionIndexOwnerConfig { + replay_limit: u64::try_from(parse_usize( + &mut get, + "DS4_SNAPSHOT_REPLAY_LIMIT", + 10_000, + 1, + MAX_REPLAY_BATCHES, + )?) + .map_err(|_| invalid("DS4_SNAPSHOT_REPLAY_LIMIT", "a bounded batch count"))?, + reconnect_min: parse_duration( + &mut get, + "DS4_SNAPSHOT_RECONNECT_MIN_MS", + 250, + 1, + 60_000, + )?, + reconnect_max: parse_duration( + &mut get, + "DS4_SNAPSHOT_RECONNECT_MAX_MS", + 5_000, + 1, + 60_000, + )?, + }; + if owner.reconnect_min > owner.reconnect_max { + return Err(invalid( + "DS4_SNAPSHOT_RECONNECT_MAX_MS", + "at least the reconnect minimum", + )); + } + if snapshot.sources.len() != 1 { + return Err(invalid( + "DS4_SNAPSHOT_LIVE_ENDPOINTS", + "exactly one live/replay engine pair", + )); + } + if snapshot.max_clients != MAX_ACTIVE_SNAPSHOT_CLIENTS { + return Err(invalid( + "DS4_SNAPSHOT_MAX_CLIENTS", + "exactly two active clients", + )); + } + let digest_secret_path = required_path(&mut get, "DS4_SNAPSHOT_DIGEST_SECRET_PATH")?; + let attestation_path = required_path(&mut get, "DS4_SNAPSHOT_ATTESTATION_PATH")?; + let session_path = snapshot + .secret_path + .as_ref() + .ok_or_else(|| invalid("DS4_SNAPSHOT_SECRET_PATH", "a session secret path"))?; + if &digest_secret_path == session_path + || attestation_path == digest_secret_path + || &attestation_path == session_path + { + return Err(invalid( + "DS4_SNAPSHOT_ATTESTATION_PATH", + "three distinct protected file paths", + )); + } + let block_size = parse_usize(&mut get, "DS4_SNAPSHOT_BLOCK_SIZE", 0, 1, MAX_BLOCK_SIZE)?; + let attention = match get("DS4_SNAPSHOT_ATTENTION_KIND") + .as_deref() + .unwrap_or("mla") + { + "full" => ConfigAttentionKind::Full, + "mla" => ConfigAttentionKind::Mla, + "sink_full" => ConfigAttentionKind::SinkFull, + _ => { + return Err(invalid( + "DS4_SNAPSHOT_ATTENTION_KIND", + "full, mla, or sink_full", + )); + } + }; + let source = &snapshot.sources[0]; + let transport = KvTransportConfig { + live_endpoint: source.live_endpoint.clone(), + replay_endpoint: Some(source.replay_endpoint.clone()), + topic: snapshot.event_topic.clone(), + connect_timeout: parse_duration( + &mut get, + "DS4_SNAPSHOT_CONNECT_TIMEOUT_MS", + 2_000, + 1, + 60_000, + )?, + replay_timeout: parse_duration( + &mut get, + "DS4_SNAPSHOT_REPLAY_TIMEOUT_MS", + 30_000, + 1, + 900_000, + )?, + max_replay_batches: usize::try_from(owner.replay_limit) + .map_err(|_| invalid("DS4_SNAPSHOT_REPLAY_LIMIT", "a bounded batch count"))?, + max_replay_tail_batches: parse_usize( + &mut get, + "DS4_SNAPSHOT_REPLAY_TAIL_LIMIT", + 1_024, + 0, + MAX_REPLAY_BATCHES, + )?, + wire_limits: KvWireLimits { + max_payload_bytes: snapshot.max_batch_payload_bytes, + max_events: snapshot.max_batch_events, + ..KvWireLimits::default() + }, + }; + let data_parallel_rank = parse_u32(&mut get, "DS4_SNAPSHOT_DATA_PARALLEL_RANK", 0)?; + let group_idx = parse_u32(&mut get, "DS4_SNAPSHOT_GROUP_INDEX", 0)?; + Ok(Self { + snapshot, + digest_secret_path: Some(digest_secret_path), + attestation_path: Some(attestation_path), + metrics_bind, + attestation_refresh, + transport: Some(transport), + owner, + group: Some(GroupMetadata { + data_parallel_rank, + group_idx, + attention_kind: attention.wire(), + disposition: GroupDisposition::Indexed, + block_size: u32::try_from(block_size) + .map_err(|_| invalid("DS4_SNAPSHOT_BLOCK_SIZE", "a bounded block size"))?, + }), + }) + } + + #[must_use] + pub fn enabled(&self) -> bool { + self.snapshot.mode == SnapshotCompanionMode::Serve + } +} + +#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)] +pub enum SingleEngineCompanionConfigError { + #[error("snapshot companion base configuration failed")] + Base(#[from] SnapshotCompanionConfigError), + #[error("invalid standalone companion setting {key}: expected {reason}")] + Invalid { + key: &'static str, + reason: &'static str, + }, + #[error("missing standalone companion setting {key}: expected {reason}")] + Missing { + key: &'static str, + reason: &'static str, + }, +} + +impl SingleEngineCompanionConfigError { + #[must_use] + pub const fn reason(self) -> &'static str { + match self { + Self::Base(_) => "base_config", + Self::Invalid { .. } => "invalid_setting", + Self::Missing { .. } => "missing_setting", + } + } +} + +#[derive(Debug, Error)] +pub enum SingleEngineCompanionError { + #[error("standalone companion session secret validation failed")] + SessionSecret(#[from] SnapshotSecretFileError), + #[error("standalone companion protected input failed")] + Attestation(#[from] CompanionAttestationError), + #[error("standalone companion source initialization failed")] + Source(#[from] CompanionIndexSourceError), + #[error("standalone companion owner failed")] + Owner(#[from] CompanionIndexOwnerError), + #[error("standalone companion snapshot server failed")] + Runtime(#[from] SnapshotCompanionRuntimeError), + #[error("standalone companion metric registration failed")] + Metrics(#[from] prometheus::Error), + #[error("standalone companion metrics listener failed")] + MetricsIo, + #[error("standalone companion configuration is incomplete")] + InvalidConfig, + #[error("standalone companion task failed")] + TaskFailed, + #[error("standalone companion shutdown timed out")] + ShutdownTimeout, +} + +impl SingleEngineCompanionError { + #[must_use] + pub const fn reason(&self) -> &'static str { + match self { + Self::SessionSecret(error) => error.reason(), + Self::Attestation(error) => error.reason(), + Self::Source(_) => "source", + Self::Owner(error) => error.reason(), + Self::Runtime(error) => error.reason(), + Self::Metrics(_) => "metrics", + Self::MetricsIo => "metrics_io", + Self::InvalidConfig => "invalid_config", + Self::TaskFailed => "task_failed", + Self::ShutdownTimeout => "shutdown_timeout", + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum SingleEngineCompanionReport { + Off, + Stopped { + owner: CompanionIndexOwnerReport, + snapshot: SnapshotCompanionRunReport, + }, +} + +enum ServiceExit { + Shutdown, + Owner(Result, JoinError>), + Snapshot(Result, JoinError>), + Metrics(Result, JoinError>), + Watcher(Result<(), JoinError>), +} + +/// Run the isolated single-engine companion until shutdown. +/// +/// # Errors +/// +/// Fails before socket publication for unsafe secrets, unauthenticated +/// incarnation input, invalid source geometry, or metrics bind failure. +#[allow(clippy::too_many_lines)] +pub async fn run_single_engine_companion( + config: SingleEngineCompanionConfig, + mut shutdown: watch::Receiver, +) -> Result { + if !config.enabled() { + return Ok(SingleEngineCompanionReport::Off); + } + let digest_path = config + .digest_secret_path + .as_deref() + .ok_or(SingleEngineCompanionError::InvalidConfig)?; + let attestation_path = config + .attestation_path + .as_deref() + .ok_or(SingleEngineCompanionError::InvalidConfig)?; + let transport = config + .transport + .clone() + .ok_or(SingleEngineCompanionError::InvalidConfig)?; + let group = config + .group + .clone() + .ok_or(SingleEngineCompanionError::InvalidConfig)?; + let policy = SnapshotSecretFilePolicy { + expected_owner_uid: config.snapshot.secret_owner_uid, + }; + // Preflight every protected input before binding metrics, connecting to the + // engine, or publishing the UDS. The snapshot runtime reopens this file at + // publication time, closing the validation/use race rather than trusting + // this first check. + let session_path = config + .snapshot + .secret_path + .as_deref() + .ok_or(SingleEngineCompanionError::InvalidConfig)?; + drop(crate::snapshot_secret_file::load_snapshot_session_secret( + session_path, + policy, + )?); + let digest_secret = load_companion_digest_secret(digest_path, policy)?; + let initial = load_authenticated_engine_incarnation(attestation_path, policy, &digest_secret)?; + let source = Arc::new(CompanionIndexSource::new( + CompanionIndexSourceConfig { + group, + index_limits: DigestIndexLimits::default(), + snapshot_limits: SnapshotLimits { + max_frame_bytes: config.snapshot.max_snapshot_frame_bytes, + max_payload_bytes: config.snapshot.max_snapshot_frame_bytes, + ..SnapshotLimits::default() + }, + max_active_sessions: config.snapshot.max_clients, + }, + initial.clone(), + 1, + digest_secret.as_bytes(), + )?); + + let registry = Arc::new(Registry::new()); + let observer = Arc::new(OwnerObserver::new(®istry)?); + let connector = Arc::new(ZmqCompanionKvEventConnector::new(transport)); + let owner = CompanionIndexOwner::new(config.owner, Arc::clone(&source), connector, observer); + let (authority_tx, authority_rx) = watch::channel(Some(initial)); + let (internal_tx, internal_rx) = watch::channel(false); + + let metrics_listener = TcpListener::bind(config.metrics_bind) + .await + .map_err(|_| SingleEngineCompanionError::MetricsIo)?; + let metrics_registry = Arc::clone(®istry); + let mut metrics_shutdown = internal_rx.clone(); + let mut metrics_task = tokio::spawn(async move { + axum::serve( + metrics_listener, + Router::new().route( + "/metrics", + get(move || { + let registry = Arc::clone(&metrics_registry); + async move { metrics_response(®istry) } + }), + ), + ) + .with_graceful_shutdown(async move { wait_for_shutdown(&mut metrics_shutdown).await }) + .await + .map_err(|_| SingleEngineCompanionError::MetricsIo) + }); + let mut owner_task = tokio::spawn(owner.run(authority_rx, internal_rx.clone())); + let producer_source: Arc = source; + let snapshot_config = config.snapshot.clone(); + let runtime_registry = Arc::clone(®istry); + let mut snapshot_task = tokio::spawn(async move { + run_snapshot_companion( + snapshot_config, + &runtime_registry, + Some(producer_source), + internal_rx, + ) + .await + .map_err(SingleEngineCompanionError::Runtime) + }); + let watcher_path = attestation_path.to_path_buf(); + let watcher_shutdown = internal_tx.subscribe(); + let refresh = config.attestation_refresh; + let mut watcher_task = tokio::spawn(async move { + watch_authenticated_engine_incarnation( + &watcher_path, + policy, + digest_secret, + refresh, + &authority_tx, + watcher_shutdown, + ) + .await; + }); + + let exit = tokio::select! { + biased; + () = wait_for_shutdown(&mut shutdown) => ServiceExit::Shutdown, + result = &mut owner_task => ServiceExit::Owner(result), + result = &mut snapshot_task => ServiceExit::Snapshot(result), + result = &mut metrics_task => ServiceExit::Metrics(result), + result = &mut watcher_task => ServiceExit::Watcher(result), + }; + let _ = internal_tx.send(true); + let deadline = Instant::now() + config.snapshot.shutdown_deadline; + + match exit { + ServiceExit::Shutdown => { + let (owner, snapshot, metrics, watcher) = tokio::join!( + await_task(&mut owner_task, deadline), + await_task(&mut snapshot_task, deadline), + await_task(&mut metrics_task, deadline), + await_task(&mut watcher_task, deadline), + ); + let owner = owner??; + let snapshot = snapshot??; + metrics??; + watcher?; + Ok(SingleEngineCompanionReport::Stopped { owner, snapshot }) + } + ServiceExit::Owner(result) => { + finish_remaining( + &mut snapshot_task, + &mut metrics_task, + &mut watcher_task, + deadline, + ) + .await?; + Err(result + .map_err(|_| SingleEngineCompanionError::TaskFailed)? + .err() + .map_or(SingleEngineCompanionError::InvalidConfig, Into::into)) + } + ServiceExit::Snapshot(result) => { + finish_owner_and_remaining( + &mut owner_task, + &mut metrics_task, + &mut watcher_task, + deadline, + ) + .await?; + Err(result + .map_err(|_| SingleEngineCompanionError::TaskFailed)? + .err() + .unwrap_or(SingleEngineCompanionError::InvalidConfig)) + } + ServiceExit::Metrics(result) => { + finish_owner_and_snapshot( + &mut owner_task, + &mut snapshot_task, + &mut watcher_task, + deadline, + ) + .await?; + Err(result + .map_err(|_| SingleEngineCompanionError::TaskFailed)? + .err() + .unwrap_or(SingleEngineCompanionError::MetricsIo)) + } + ServiceExit::Watcher(result) => { + finish_owner_and_snapshot( + &mut owner_task, + &mut snapshot_task, + &mut metrics_task, + deadline, + ) + .await?; + result.map_err(|_| SingleEngineCompanionError::TaskFailed)?; + Err(SingleEngineCompanionError::InvalidConfig) + } + } +} + +struct OwnerObserver { + events: CounterVec, + authority: Gauge, + replay_duration: HistogramVec, +} + +impl OwnerObserver { + fn new(registry: &Registry) -> Result { + let observer = Self { + events: CounterVec::new( + Opts::new( + "ds4proxy_snapshot_companion_owner_events_total", + "Companion owner transitions by bounded event and reason", + ), + &["event", "reason"], + )?, + authority: Gauge::with_opts(Opts::new( + "ds4proxy_snapshot_companion_authority", + "Whether authenticated engine-incarnation authority is available", + ))?, + replay_duration: HistogramVec::new( + prometheus::HistogramOpts::new( + "ds4proxy_snapshot_companion_owner_replay_duration_seconds", + "Transport replay duration by bounded kind and outcome", + ), + &["kind", "outcome"], + )?, + }; + registry.register(Box::new(observer.events.clone()))?; + registry.register(Box::new(observer.authority.clone()))?; + registry.register(Box::new(observer.replay_duration.clone()))?; + Ok(observer) + } +} + +impl CompanionIndexOwnerObserver for OwnerObserver { + fn observe(&self, event: CompanionIndexOwnerEvent) { + let (name, reason) = owner_event_labels(&event); + self.events.with_label_values(&[name, reason]).inc(); + match &event { + CompanionIndexOwnerEvent::AuthorityAvailable => self.authority.set(1.0), + CompanionIndexOwnerEvent::AuthorityLost | CompanionIndexOwnerEvent::Shutdown => { + self.authority.set(0.0); + } + CompanionIndexOwnerEvent::Replay { + kind, + outcome, + profile: Some(profile), + } => self + .replay_duration + .with_label_values(&[replay_kind_label(*kind), replay_outcome_label(*outcome)]) + .observe(profile.elapsed.as_secs_f64()), + _ => {} + } + tracing::debug!(event = name, reason, "snapshot companion owner transition"); + } +} + +fn owner_event_labels(event: &CompanionIndexOwnerEvent) -> (&'static str, &'static str) { + match event { + CompanionIndexOwnerEvent::AuthorityAvailable => ("authority", "available"), + CompanionIndexOwnerEvent::AuthorityLost => ("authority", "lost"), + CompanionIndexOwnerEvent::ConnectAttempt => ("connect", "attempt"), + CompanionIndexOwnerEvent::Connected => ("connect", "connected"), + CompanionIndexOwnerEvent::Rebuild(reason) => ("rebuild", rebuild_label(*reason)), + CompanionIndexOwnerEvent::Replay { outcome, .. } => { + ("replay", replay_outcome_label(*outcome)) + } + CompanionIndexOwnerEvent::Ready => ("source", "ready"), + CompanionIndexOwnerEvent::LiveApplied => ("live", "applied"), + CompanionIndexOwnerEvent::LiveDuplicate => ("live", "duplicate"), + CompanionIndexOwnerEvent::Shutdown => ("owner", "shutdown"), + } +} + +const fn rebuild_label(reason: CompanionIndexOwnerRebuildReason) -> &'static str { + match reason { + CompanionIndexOwnerRebuildReason::Startup => "startup", + CompanionIndexOwnerRebuildReason::AuthorityChanged => "authority_changed", + CompanionIndexOwnerRebuildReason::AuthorityLost => "authority_lost", + CompanionIndexOwnerRebuildReason::Disconnected => "disconnected", + CompanionIndexOwnerRebuildReason::Transport => "transport", + CompanionIndexOwnerRebuildReason::Replay => "replay", + CompanionIndexOwnerRebuildReason::ReplayInvalid => "replay_invalid", + CompanionIndexOwnerRebuildReason::ReplayTooLarge => "replay_too_large", + CompanionIndexOwnerRebuildReason::Apply => "apply", + } +} + +const fn replay_kind_label(kind: CompanionIndexOwnerReplayKind) -> &'static str { + match kind { + CompanionIndexOwnerReplayKind::Full => "full", + } +} + +const fn replay_outcome_label(outcome: CompanionIndexOwnerReplayOutcome) -> &'static str { + match outcome { + CompanionIndexOwnerReplayOutcome::Complete => "complete", + CompanionIndexOwnerReplayOutcome::TransportFailed => "transport_failed", + CompanionIndexOwnerReplayOutcome::Invalid => "invalid", + CompanionIndexOwnerReplayOutcome::Cancelled => "cancelled", + } +} + +fn metrics_response(registry: &Registry) -> Response { + let mut output = Vec::new(); + let encoder = TextEncoder::new(); + if encoder.encode(®istry.gather(), &mut output).is_err() { + return Response::builder() + .status(500) + .body(Body::from("metrics encoding failed")) + .expect("valid response"); + } + Response::builder() + .status(200) + .header("content-type", encoder.format_type()) + .body(Body::from(output)) + .expect("valid response") +} + +async fn await_task( + task: &mut JoinHandle, + deadline: Instant, +) -> Result { + if let Ok(result) = timeout_at(deadline, &mut *task).await { + result.map_err(|_| SingleEngineCompanionError::TaskFailed) + } else { + task.abort(); + let _ = task.await; + Err(SingleEngineCompanionError::ShutdownTimeout) + } +} + +async fn finish_remaining( + first: &mut JoinHandle, + second: &mut JoinHandle, + third: &mut JoinHandle, + deadline: Instant, +) -> Result<(), SingleEngineCompanionError> { + let (first, second, third) = tokio::join!( + await_task(first, deadline), + await_task(second, deadline), + await_task(third, deadline), + ); + let _ = first?; + let _ = second?; + let _ = third?; + Ok(()) +} + +async fn finish_owner_and_remaining( + first: &mut JoinHandle, + second: &mut JoinHandle, + third: &mut JoinHandle, + deadline: Instant, +) -> Result<(), SingleEngineCompanionError> { + finish_remaining(first, second, third, deadline).await +} + +async fn finish_owner_and_snapshot( + first: &mut JoinHandle, + second: &mut JoinHandle, + third: &mut JoinHandle, + deadline: Instant, +) -> Result<(), SingleEngineCompanionError> { + finish_remaining(first, second, third, deadline).await +} + +async fn wait_for_shutdown(shutdown: &mut watch::Receiver) { + loop { + if *shutdown.borrow() || shutdown.changed().await.is_err() { + return; + } + } +} + +fn required_path( + get: &mut impl FnMut(&str) -> Option, + key: &'static str, +) -> Result { + let raw = get(key).filter(|value| !value.is_empty()).ok_or( + SingleEngineCompanionConfigError::Missing { + key, + reason: "an absolute normalized protected file path", + }, + )?; + let path = PathBuf::from(&raw); + if raw.len() > MAX_PATH_BYTES + || !path.is_absolute() + || path.file_name().is_none() + || raw.ends_with('/') + || raw[1..] + .split('/') + .any(|component| component.is_empty() || matches!(component, "." | "..")) + { + return Err(invalid(key, "an absolute normalized protected file path")); + } + Ok(path) +} + +fn parse_loopback(raw: &str) -> Result { + let address = raw + .parse::() + .map_err(|_| invalid("DS4_SNAPSHOT_METRICS_BIND", "a loopback IP socket address"))?; + if !matches!(address.ip(), IpAddr::V4(_) | IpAddr::V6(_)) || !address.ip().is_loopback() { + return Err(invalid( + "DS4_SNAPSHOT_METRICS_BIND", + "a loopback IP socket address", + )); + } + Ok(address) +} + +fn parse_duration( + get: &mut impl FnMut(&str) -> Option, + key: &'static str, + fallback: u64, + minimum: u64, + maximum: u64, +) -> Result { + let value = get(key) + .as_deref() + .map(str::parse::) + .transpose() + .map_err(|_| invalid(key, "bounded milliseconds"))? + .unwrap_or(fallback); + if !(minimum..=maximum).contains(&value) { + return Err(invalid(key, "bounded milliseconds")); + } + Ok(Duration::from_millis(value)) +} + +fn parse_usize( + get: &mut impl FnMut(&str) -> Option, + key: &'static str, + fallback: usize, + minimum: usize, + maximum: usize, +) -> Result { + let value = get(key) + .as_deref() + .map(str::parse::) + .transpose() + .map_err(|_| invalid(key, "a bounded integer"))? + .unwrap_or(fallback); + if !(minimum..=maximum).contains(&value) { + return Err(invalid(key, "a bounded integer")); + } + Ok(value) +} + +fn parse_u32( + get: &mut impl FnMut(&str) -> Option, + key: &'static str, + fallback: u32, +) -> Result { + get(key) + .as_deref() + .map(str::parse::) + .transpose() + .map_err(|_| invalid(key, "a non-negative 32-bit integer")) + .map(|value| value.unwrap_or(fallback)) +} + +const fn invalid(key: &'static str, reason: &'static str) -> SingleEngineCompanionConfigError { + SingleEngineCompanionConfigError::Invalid { key, reason } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::HashMap, + fs, + os::unix::fs::{MetadataExt, PermissionsExt}, + sync::atomic::{AtomicU64, Ordering}, + }; + + use super::*; + use crate::{ + companion_attestation::{ + encode_authenticated_engine_incarnation, load_companion_digest_secret, + }, + kv_snapshot::EngineIncarnation, + }; + use tokio::time::timeout; + + static TEST_ID: AtomicU64 = AtomicU64::new(1); + // Drone runs this suite as root. Keep its protected files root-owned while + // exercising the production contract with distinct non-root process UIDs. + const ROOT_CONTAINER_COMPANION_UID: u32 = 12_001; + const ROOT_CONTAINER_CLIENT_UID: u32 = 12_002; + + struct TestFiles { + directory: PathBuf, + socket: PathBuf, + session: PathBuf, + digest: PathBuf, + attestation: PathBuf, + owner: u32, + companion_uid: u32, + client_uid: u32, + } + + impl TestFiles { + fn new() -> Self { + let id = TEST_ID.fetch_add(1, Ordering::Relaxed); + // The production UDS contract deliberately caps paths at 64 bytes. + // Do not inherit a potentially long Cargo `TMPDIR` in this test. + let directory = + PathBuf::from("/tmp").join(format!("md-svc-{}-{id}", std::process::id())); + fs::create_dir(&directory).unwrap(); + fs::set_permissions(&directory, fs::Permissions::from_mode(0o700)).unwrap(); + let owner = fs::metadata(&directory).unwrap().uid(); + let (companion_uid, client_uid) = if owner == 0 { + (ROOT_CONTAINER_COMPANION_UID, ROOT_CONTAINER_CLIENT_UID) + } else { + (owner, owner.saturating_add(1).max(1)) + }; + let session = directory.join("session"); + let digest = directory.join("digest"); + for (path, byte) in [(&session, 0x31), (&digest, 0x51)] { + fs::write(path, [byte; 32]).unwrap(); + fs::set_permissions(path, fs::Permissions::from_mode(0o600)).unwrap(); + } + let attestation = directory.join("attestation"); + let secret = load_companion_digest_secret( + &digest, + SnapshotSecretFilePolicy { + expected_owner_uid: owner, + }, + ) + .unwrap(); + let encoded = encode_authenticated_engine_incarnation(&incarnation(), &secret).unwrap(); + fs::write(&attestation, encoded).unwrap(); + fs::set_permissions(&attestation, fs::Permissions::from_mode(0o600)).unwrap(); + Self { + socket: directory.join("companion.sock"), + directory, + session, + digest, + attestation, + owner, + companion_uid, + client_uid, + } + } + + fn values(&self) -> HashMap<&'static str, String> { + HashMap::from([ + ("DS4_SNAPSHOT_COMPANION_MODE", "serve".to_owned()), + ( + "DS4_SNAPSHOT_SOCKET_PATH", + self.socket.to_string_lossy().into_owned(), + ), + ("DS4_SNAPSHOT_COMPANION_UID", self.companion_uid.to_string()), + ("DS4_SNAPSHOT_CLIENT_UID", self.client_uid.to_string()), + ( + "DS4_SNAPSHOT_SECRET_PATH", + self.session.to_string_lossy().into_owned(), + ), + ("DS4_SNAPSHOT_SECRET_OWNER_UID", self.owner.to_string()), + ( + "DS4_SNAPSHOT_LIVE_ENDPOINTS", + "tcp://127.0.0.1:45171".to_owned(), + ), + ( + "DS4_SNAPSHOT_REPLAY_ENDPOINTS", + "tcp://127.0.0.1:45172".to_owned(), + ), + ( + "DS4_SNAPSHOT_DIGEST_SECRET_PATH", + self.digest.to_string_lossy().into_owned(), + ), + ( + "DS4_SNAPSHOT_ATTESTATION_PATH", + self.attestation.to_string_lossy().into_owned(), + ), + ("DS4_SNAPSHOT_BLOCK_SIZE", "64".to_owned()), + ("DS4_SNAPSHOT_METRICS_BIND", "127.0.0.1:0".to_owned()), + ("DS4_SNAPSHOT_RECONNECT_MIN_MS", "10".to_owned()), + ("DS4_SNAPSHOT_RECONNECT_MAX_MS", "20".to_owned()), + ("DS4_SNAPSHOT_ATTESTATION_REFRESH_MS", "50".to_owned()), + ]) + } + } + + impl Drop for TestFiles { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.directory); + } + } + + fn incarnation() -> EngineIncarnation { + EngineIncarnation { + engine_id: "engine-a".to_owned(), + model_revision: "revision".to_owned(), + image_digest: "sha256:image".to_owned(), + process_started_unix_ns: 42, + attestation_sha256: vec![7; 32], + } + } + + fn load(values: &HashMap<&str, String>) -> SingleEngineCompanionConfig { + SingleEngineCompanionConfig::from_lookup(|key| values.get(key).cloned()).unwrap() + } + + #[test] + fn off_is_the_default_and_requires_no_files_or_engine() { + let config = SingleEngineCompanionConfig::from_lookup(|key| match key { + // Standalone-only settings are not parsed while disabled. + "DS4_SNAPSHOT_METRICS_BIND" => Some("not-an-address".to_owned()), + "DS4_SNAPSHOT_RECONNECT_MIN_MS" => Some("not-a-duration".to_owned()), + _ => None, + }) + .unwrap(); + assert!(!config.enabled()); + assert!(config.digest_secret_path.is_none()); + assert!(config.attestation_path.is_none()); + assert!(config.transport.is_none()); + assert!(config.group.is_none()); + } + + #[tokio::test] + async fn off_runtime_has_no_listener_or_filesystem_side_effects() { + let mut values = HashMap::new(); + values.insert("DS4_SNAPSHOT_METRICS_BIND", "127.0.0.1:1".to_owned()); + let config = load(&values); + let (_shutdown_tx, shutdown_rx) = watch::channel(false); + assert_eq!( + run_single_engine_companion(config, shutdown_rx) + .await + .unwrap(), + SingleEngineCompanionReport::Off + ); + } + + #[tokio::test] + async fn bounded_join_aborts_stalled_tasks_and_maps_panics() { + let mut stalled = tokio::spawn(std::future::pending::<()>()); + assert!(matches!( + await_task(&mut stalled, Instant::now() + Duration::from_millis(10)).await, + Err(SingleEngineCompanionError::ShutdownTimeout) + )); + assert!(stalled.is_finished()); + + let mut panicked = tokio::spawn(async { panic!("test task panic") }); + assert!(matches!( + await_task(&mut panicked, Instant::now() + Duration::from_secs(1)).await, + Err(SingleEngineCompanionError::TaskFailed) + )); + } + + #[test] + fn serve_requires_one_engine_explicit_geometry_and_loopback_metrics() { + let files = TestFiles::new(); + let mut values = files.values(); + let config = load(&values); + assert!(config.enabled()); + assert_eq!(config.snapshot.sources.len(), 1); + assert_eq!(config.group.as_ref().unwrap().block_size, 64); + + values.remove("DS4_SNAPSHOT_BLOCK_SIZE"); + assert!(SingleEngineCompanionConfig::from_lookup(|key| values.get(key).cloned()).is_err()); + let mut values = files.values(); + values.insert( + "DS4_SNAPSHOT_LIVE_ENDPOINTS", + "tcp://127.0.0.1:1,tcp://127.0.0.1:2".to_owned(), + ); + values.insert( + "DS4_SNAPSHOT_REPLAY_ENDPOINTS", + "tcp://127.0.0.1:3,tcp://127.0.0.1:4".to_owned(), + ); + assert!(SingleEngineCompanionConfig::from_lookup(|key| values.get(key).cloned()).is_err()); + let mut values = files.values(); + values.insert("DS4_SNAPSHOT_METRICS_BIND", "0.0.0.0:9091".to_owned()); + assert!(SingleEngineCompanionConfig::from_lookup(|key| values.get(key).cloned()).is_err()); + let mut values = files.values(); + values.insert("DS4_SNAPSHOT_MAX_CLIENTS", "1".to_owned()); + assert!(matches!( + SingleEngineCompanionConfig::from_lookup(|key| values.get(key).cloned()), + Err(SingleEngineCompanionConfigError::Invalid { + key: "DS4_SNAPSHOT_MAX_CLIENTS", + reason: "exactly two active clients", + }) + )); + } + + #[test] + fn stable_observe_only_fence_has_a_bounded_metric_reason() { + let event = + CompanionIndexOwnerEvent::Rebuild(CompanionIndexOwnerRebuildReason::ReplayTooLarge); + assert_eq!(owner_event_labels(&event), ("rebuild", "replay_too_large")); + + let registry = Registry::new(); + let observer = OwnerObserver::new(®istry).unwrap(); + observer.observe(event); + let family = registry + .gather() + .into_iter() + .find(|family| family.name() == "ds4proxy_snapshot_companion_owner_events_total") + .expect("owner metric family"); + assert_eq!(family.get_metric().len(), 1); + let labels = family.get_metric()[0].get_label(); + assert!( + labels + .iter() + .any(|label| { label.name() == "event" && label.value() == "rebuild" }) + ); + assert!( + labels + .iter() + .any(|label| { label.name() == "reason" && label.value() == "replay_too_large" }) + ); + } + + #[test] + fn debug_and_errors_never_expose_protected_paths_or_endpoints() { + let files = TestFiles::new(); + let config = load(&files.values()); + let debug = format!("{config:?}"); + for forbidden in [ + files.digest.to_string_lossy(), + files.attestation.to_string_lossy(), + std::borrow::Cow::Borrowed("45171"), + ] { + assert!(!debug.contains(forbidden.as_ref())); + } + } + + #[tokio::test] + async fn unauthenticated_startup_fails_before_socket_publication() { + let files = TestFiles::new(); + fs::write(&files.attestation, b"malformed").unwrap(); + let config = load(&files.values()); + let (_shutdown_tx, shutdown_rx) = watch::channel(false); + assert!(matches!( + run_single_engine_companion(config, shutdown_rx).await, + Err(SingleEngineCompanionError::Attestation(_)) + )); + assert!(!files.socket.exists()); + } + + #[tokio::test] + async fn unsafe_session_secret_fails_before_socket_publication() { + let files = TestFiles::new(); + fs::set_permissions(&files.session, fs::Permissions::from_mode(0o622)).unwrap(); + let config = load(&files.values()); + let (_shutdown_tx, shutdown_rx) = watch::channel(false); + assert!(matches!( + run_single_engine_companion(config, shutdown_rx).await, + Err(SingleEngineCompanionError::SessionSecret( + SnapshotSecretFileError::UnsafePermissions + )) + )); + assert!(!files.socket.exists()); + } + + #[tokio::test] + async fn full_composition_publishes_socket_and_shuts_down_cleanly() { + let files = TestFiles::new(); + if files.owner == 0 { + // The production contract intentionally prohibits root service + // identity; root-only test containers cannot exercise peer UIDs. + return; + } + let config = load(&files.values()); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let task = tokio::spawn(run_single_engine_companion(config, shutdown_rx)); + timeout(Duration::from_secs(2), async { + while !files.socket.exists() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("snapshot socket must be published"); + shutdown_tx.send(true).unwrap(); + let report = timeout(Duration::from_secs(2), task) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!(matches!( + report, + SingleEngineCompanionReport::Stopped { .. } + )); + assert!(!files.socket.exists()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 33f99f6..523d44d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,10 +1,12 @@ pub mod block_digest; +pub mod companion_attestation; pub mod companion_config; pub mod companion_index_owner; pub mod companion_index_source; pub mod companion_metrics; pub mod companion_runtime; +pub mod companion_service; pub mod compat; pub mod config; pub mod digest_index; diff --git a/src/snapshot_secret_file.rs b/src/snapshot_secret_file.rs index f94cc15..7d48650 100644 --- a/src/snapshot_secret_file.rs +++ b/src/snapshot_secret_file.rs @@ -17,6 +17,7 @@ use thiserror::Error; use crate::snapshot_session::{SNAPSHOT_SESSION_SECRET_BYTES, SnapshotSessionSecret}; const MAX_SECRET_READ_BYTES: usize = SNAPSHOT_SESSION_SECRET_BYTES + 1; +const MAX_CONTROL_FILE_BYTES: usize = 64 * 1024; const GROUP_OR_WORLD_WRITE: u32 = 0o022; const STICKY_BIT: u32 = 0o1000; @@ -144,6 +145,67 @@ pub fn load_snapshot_session_secret( Ok(SnapshotSessionSecret::new(secret_bytes)) } +/// Load a bounded, protected control-plane file with the same ownership, +/// symlink, link-count, permission, and inode-stability policy as a secret. +/// +/// This is used for authenticated incarnation envelopes. It returns bytes only +/// after the post-read path and parent checks succeed. +/// +/// # Errors +/// +/// Returns a content-free error if `max_bytes` is zero/too large, the file is +/// empty or oversized, or any hardened file invariant fails. +pub fn load_snapshot_control_file( + path: &Path, + policy: SnapshotSecretFilePolicy, + max_bytes: usize, +) -> Result, SnapshotSecretFileError> { + if max_bytes == 0 || max_bytes > MAX_CONTROL_FILE_BYTES { + return Err(SnapshotSecretFileError::InvalidLength); + } + validate_trusted_parents(path, policy)?; + let before = + fs::symlink_metadata(path).map_err(|_| SnapshotSecretFileError::InvalidMetadata)?; + validate_control_metadata(&before, policy, max_bytes)?; + let mut file = File::open(path).map_err(|_| SnapshotSecretFileError::OpenFailed)?; + let opened = file + .metadata() + .map_err(|_| SnapshotSecretFileError::InvalidMetadata)?; + validate_control_metadata(&opened, policy, max_bytes)?; + require_same_file(&before, &opened)?; + + let opened_len = usize::try_from(opened.len()).unwrap_or(max_bytes); + let mut bytes = Vec::with_capacity(max_bytes.min(opened_len).saturating_add(1)); + let read_result = file + .by_ref() + .take( + u64::try_from(max_bytes) + .unwrap_or(u64::MAX) + .saturating_add(1), + ) + .read_to_end(&mut bytes); + if read_result.is_err() { + bytes.fill(0); + return Err(SnapshotSecretFileError::ReadFailed); + } + if bytes.is_empty() || bytes.len() > max_bytes { + bytes.fill(0); + return Err(SnapshotSecretFileError::InvalidLength); + } + let post_read = fs::symlink_metadata(path) + .map_err(|_| SnapshotSecretFileError::InvalidMetadata) + .and_then(|after| { + validate_control_metadata(&after, policy, max_bytes) + .and_then(|()| require_same_file(&opened, &after)) + }) + .and_then(|()| validate_trusted_parents(path, policy)); + if let Err(error) = post_read { + bytes.fill(0); + return Err(error); + } + Ok(bytes) +} + fn read_bounded( file: &mut File, destination: &mut [u8; MAX_SECRET_READ_BYTES], @@ -182,6 +244,29 @@ fn validate_secret_metadata( Ok(()) } +fn validate_control_metadata( + metadata: &Metadata, + policy: SnapshotSecretFilePolicy, + max_bytes: usize, +) -> Result<(), SnapshotSecretFileError> { + if !metadata.file_type().is_file() { + return Err(SnapshotSecretFileError::NotRegularFile); + } + if metadata.nlink() != 1 { + return Err(SnapshotSecretFileError::InvalidLinkCount); + } + if metadata.uid() != policy.expected_owner_uid { + return Err(SnapshotSecretFileError::UnexpectedOwner); + } + if metadata.mode() & GROUP_OR_WORLD_WRITE != 0 { + return Err(SnapshotSecretFileError::UnsafePermissions); + } + if metadata.len() == 0 || metadata.len() > u64::try_from(max_bytes).unwrap_or(u64::MAX) { + return Err(SnapshotSecretFileError::InvalidLength); + } + Ok(()) +} + fn require_same_file(left: &Metadata, right: &Metadata) -> Result<(), SnapshotSecretFileError> { if left.dev() != right.dev() || left.ino() != right.ino() { return Err(SnapshotSecretFileError::FileIdentityChanged); diff --git a/src/snapshot_socket_path.rs b/src/snapshot_socket_path.rs index a611e4f..f32ecb5 100644 --- a/src/snapshot_socket_path.rs +++ b/src/snapshot_socket_path.rs @@ -66,6 +66,8 @@ pub enum SnapshotSocketPathError { InvalidPath, #[error("snapshot socket parent metadata is unavailable")] ParentMetadata, + #[error("snapshot companion process metadata is unavailable")] + ProcessMetadata, #[error("snapshot socket parent contains a symbolic link")] ParentSymlink, #[error("snapshot socket parent is not a directory")] @@ -82,6 +84,16 @@ pub enum SnapshotSocketPathError { PermissionFailed, #[error("snapshot socket metadata is invalid")] InvalidSocketMetadata, + #[error("snapshot socket target metadata is unavailable")] + SocketTargetMetadata, + #[error("snapshot socket target is not a Unix socket")] + SocketType, + #[error("snapshot socket owner does not match its parent owner")] + SocketOwnerMismatch, + #[error("snapshot socket permissions are invalid")] + SocketPermissions, + #[error("snapshot socket has an unexpected link count")] + SocketLinkCount, #[error("snapshot socket atomic publication failed")] PublishFailed, #[error("snapshot socket temporary cleanup failed")] @@ -94,6 +106,7 @@ impl SnapshotSocketPathError { match self { Self::InvalidPath => "invalid_path", Self::ParentMetadata => "parent_metadata", + Self::ProcessMetadata => "process_metadata", Self::ParentSymlink => "parent_symlink", Self::ParentNotDirectory => "parent_not_directory", Self::ParentOwnerMismatch => "parent_owner_mismatch", @@ -102,6 +115,11 @@ impl SnapshotSocketPathError { Self::BindFailed => "bind_failed", Self::PermissionFailed => "permission_failed", Self::InvalidSocketMetadata => "invalid_socket_metadata", + Self::SocketTargetMetadata => "socket_target_metadata", + Self::SocketType => "socket_type", + Self::SocketOwnerMismatch => "socket_owner_mismatch", + Self::SocketPermissions => "socket_permissions", + Self::SocketLinkCount => "socket_link_count", Self::PublishFailed => "publish_failed", Self::TemporaryCleanupFailed => "temporary_cleanup_failed", } @@ -240,6 +258,60 @@ pub fn validate_socket_client_path( validate_socket_parent(normalized_parent(public_path)?, policy) } +/// Validate the filesystem readiness contract of a published companion socket. +/// +/// This deliberately does not connect: the companion's own UID is not an +/// authorized router UID, so an active probe would create rejected sessions +/// and consume accept-loop telemetry. Container process liveness plus the +/// exclusively-owned parent and exact published-socket metadata prove that the +/// correctly-identified standalone service reached socket publication. +/// Authenticated authority readiness remains an operational metric rather than +/// Docker health. +/// +/// # Errors +/// +/// Returns [`SnapshotSocketPathError`] for a relative/non-normalized path, an +/// unsafe or symlinked parent not owned by this process, or a target that is +/// not a singly-linked mode-0660 Unix socket owned by the parent owner. +pub fn validate_published_socket( + public_path: &Path, +) -> Result { + let parent = normalized_parent(public_path)?; + let parent_metadata = + fs::symlink_metadata(parent).map_err(|_| SnapshotSocketPathError::ParentMetadata)?; + // The companion currently targets Linux, where the followed `/proc/self` + // process directory is owned by the effective UID. Keeping this check in + // safe std code avoids adding a libc dependency or an unsafe `geteuid(2)` + // call merely for the health probe. A missing procfs fails closed. + let process_uid = fs::metadata("/proc/self") + .map_err(|_| SnapshotSocketPathError::ProcessMetadata)? + .uid(); + validate_socket_parent( + parent, + SocketParentPolicy { + owner_uid: process_uid, + }, + )?; + let metadata = fs::symlink_metadata(public_path) + .map_err(|_| SnapshotSocketPathError::SocketTargetMetadata)?; + if !metadata.file_type().is_socket() { + return Err(SnapshotSocketPathError::SocketType); + } + if metadata.uid() != parent_metadata.uid() { + return Err(SnapshotSocketPathError::SocketOwnerMismatch); + } + if metadata.mode() & 0o777 != SOCKET_MODE { + return Err(SnapshotSocketPathError::SocketPermissions); + } + if metadata.nlink() != 1 { + return Err(SnapshotSocketPathError::SocketLinkCount); + } + Ok(SocketIdentity { + device: metadata.dev(), + inode: metadata.ino(), + }) +} + /// Bind a unique private socket, set mode `0660`, and atomically publish it /// without replacing any existing public target. /// @@ -458,6 +530,10 @@ mod tests { assert_eq!(metadata.mode() & 0o777, SOCKET_MODE); assert_eq!(metadata.dev(), published.identity().device()); assert_eq!(metadata.ino(), published.identity().inode()); + assert_eq!( + validate_published_socket(&path).unwrap(), + published.identity() + ); let client = UnixStream::connect(&path).unwrap(); let _server = published.listener().accept().unwrap().0; @@ -542,6 +618,42 @@ mod tests { bind_and_publish(Path::new("relative.sock"), directory.policy()), Err(SnapshotSocketPathError::InvalidPath) )); + assert!(matches!( + validate_published_socket(Path::new("relative.sock")), + Err(SnapshotSocketPathError::InvalidPath) + )); + } + + #[test] + fn healthcheck_rejects_missing_non_socket_linked_and_wrong_mode_targets() { + let directory = TestDirectory::new(); + let path = directory.path().join("companion.sock"); + assert!(matches!( + validate_published_socket(&path), + Err(SnapshotSocketPathError::SocketTargetMetadata) + )); + + fs::write(&path, b"not a socket").unwrap(); + assert!(matches!( + validate_published_socket(&path), + Err(SnapshotSocketPathError::SocketType) + )); + fs::remove_file(&path).unwrap(); + + let listener = UnixListener::bind(&path).unwrap(); + fs::set_permissions(&path, fs::Permissions::from_mode(0o600)).unwrap(); + assert!(matches!( + validate_published_socket(&path), + Err(SnapshotSocketPathError::SocketPermissions) + )); + fs::set_permissions(&path, fs::Permissions::from_mode(SOCKET_MODE)).unwrap(); + let linked = directory.path().join("linked.sock"); + fs::hard_link(&path, &linked).unwrap(); + assert!(matches!( + validate_published_socket(&path), + Err(SnapshotSocketPathError::SocketLinkCount) + )); + drop(listener); } #[test] diff --git a/tests/snapshot_companion_cli.rs b/tests/snapshot_companion_cli.rs new file mode 100644 index 0000000..9497110 --- /dev/null +++ b/tests/snapshot_companion_cli.rs @@ -0,0 +1,87 @@ +use std::{ + fs, + os::unix::{fs::PermissionsExt, net::UnixListener}, + path::{Path, PathBuf}, + process::{Command, Output}, + sync::atomic::{AtomicU64, Ordering}, +}; + +static TEST_SEQUENCE: AtomicU64 = AtomicU64::new(1); + +struct TestDirectory(PathBuf); + +impl TestDirectory { + fn new() -> Self { + let sequence = TEST_SEQUENCE.fetch_add(1, Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "md-companion-cli-{:x}-{sequence:x}", + std::process::id() + )); + fs::create_dir(&path).unwrap(); + fs::set_permissions(&path, fs::Permissions::from_mode(0o700)).unwrap(); + Self(path) + } + + fn path(&self) -> &Path { + &self.0 + } +} + +impl Drop for TestDirectory { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } +} + +fn companion(arguments: &[&str]) -> Output { + Command::new(env!("CARGO_BIN_EXE_mini-dynamo-snapshot-companion")) + .args(arguments) + .env_remove("DS4_SNAPSHOT_COMPANION_MODE") + .output() + .unwrap() +} + +#[test] +fn healthcheck_accepts_only_exact_published_socket_metadata() { + let directory = TestDirectory::new(); + let socket = directory.path().join("companion.sock"); + let _listener = UnixListener::bind(&socket).unwrap(); + fs::set_permissions(&socket, fs::Permissions::from_mode(0o660)).unwrap(); + + let output = companion(&["healthcheck", socket.to_str().unwrap()]); + assert!(output.status.success(), "{output:?}"); + assert!(output.stdout.is_empty()); + assert!(output.stderr.is_empty()); +} + +#[test] +fn healthcheck_rejects_missing_path_and_extra_arguments() { + let directory = TestDirectory::new(); + let missing = directory.path().join("missing.sock"); + assert!( + !companion(&["healthcheck", missing.to_str().unwrap()]) + .status + .success() + ); + assert!(!companion(&["healthcheck"]).status.success()); + assert!( + !companion(&["healthcheck", "/socket", "extra"]) + .status + .success() + ); +} + +#[test] +fn fixture_and_unknown_subcommands_fail_closed() { + assert!(!companion(&["snapshot-companion-fixture"]).status.success()); + assert!(!companion(&["unknown"]).status.success()); +} + +#[test] +fn no_arguments_preserves_default_off_mode() { + let output = companion(&[]); + assert!(output.status.success(), "{output:?}"); + let stdout = String::from_utf8(output.stdout).unwrap(); + assert!(stdout.contains("\"enabled\":false")); + assert!(stdout.contains("\"report\":\"Off\"")); +}