From a5982a8a71b746516cc8b0779b9556996aa951d0 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:32:02 -0700 Subject: [PATCH 01/13] fix(gateway): validate replicated instance records before they reach the data plane MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Instance records synced from a peer went into ProxyState — and from there into the rendered wg.conf — without re-running any of the checks the registration path applies. That made a single malformed record a node-wide failure: `wg syncconf` rejects the *entire* config file when one peer key is malformed, and a key containing a newline can inject `Endpoint=`/`AllowedIPs=` directives. Last-writer-wins replication also cannot enforce invariants that span keys, so a synced instance could carry the gateway's own wg IP or an IP/public key already claimed by another instance. All KV instance records now pass through `kv::import`, which re-runs the registration checks (public key is 32 base64-encoded bytes, IP and key unique, address not one of this gateway's own) and skips only the offending record, never the batch. Conflicts resolve by registration time so every node reaches the same decision from the same KV contents. The renderer re-checks each peer as a last line of defense before values reach `wg`. The address check deliberately says nothing about which pool an address came from. A CVM registers with one gateway but is handed *every* gateway as a WireGuard server, so each node carries peers for the CVMs registered on the other nodes, and each node allocates from its own `client_ip_range`. Nothing in a node's config describes the other nodes' pools, and the deployments do not agree on a shape that could be inferred: `deploy-to-vmm.sh` puts every pool inside one /16 that each interface covers, while `test-run/cluster.sh` and the e2e configs give each node a /24 that no other node's interface covers. So `is_valid_client_ip` keeps governing allocation, and the import boundary and the renderer use `is_routable_client_ip`, which asserts only what a node can know on its own: an ordinary unicast address that is not one of its own. What keeps the peer list coherent is the uniqueness pass, which runs over the whole KV contents and therefore holds cluster-wide. Refs #1029 --- dstack/gateway/src/config.rs | 42 +++ dstack/gateway/src/kv/import.rs | 318 ++++++++++++++++++ dstack/gateway/src/kv/mod.rs | 3 +- dstack/gateway/src/main_service.rs | 100 +++--- ...ateway__main_service__tests__config-2.snap | 2 +- ...ateway__main_service__tests__config-3.snap | 4 +- ..._gateway__main_service__tests__config.snap | 2 +- dstack/gateway/src/main_service/tests.rs | 158 +++++++-- dstack/gateway/src/models.rs | 45 +-- 9 files changed, 564 insertions(+), 110 deletions(-) create mode 100644 dstack/gateway/src/kv/import.rs diff --git a/dstack/gateway/src/config.rs b/dstack/gateway/src/config.rs index ae7837bff..f89153a8c 100644 --- a/dstack/gateway/src/config.rs +++ b/dstack/gateway/src/config.rs @@ -31,6 +31,48 @@ impl WgConfig { fn validate(&self) -> Result<()> { validate(self.ip, &self.reserved_net, self.client_ip_range) } + + /// Whether this gateway may allocate `ip` to a CVM registering with it. + /// + /// Narrower than [`Self::is_routable_client_ip`]: `client_ip_range` is this + /// node's *share* of the cluster's address space, and handing out an address + /// from outside it would collide with whichever node owns that share. + pub fn is_valid_client_ip(&self, ip: Ipv4Addr) -> bool { + self.client_ip_range.contains(&ip) && self.is_routable_client_ip(ip) + } + + /// Whether `ip` may appear as a WireGuard peer address on this gateway. + /// + /// Deliberately says nothing about *which pool* the address came from. A + /// CVM registers with one gateway but is handed every gateway as a + /// WireGuard server, so each node carries peers for the CVMs registered on + /// the other nodes — and each node allocates from its own + /// `client_ip_range`. Nothing in this node's config describes the other + /// nodes' pools, and the deployments do not even agree on a shape that + /// could be inferred: `dstack-app/deploy-to-vmm.sh` puts every pool inside + /// one /16 that each interface covers, while `test-run/cluster.sh` and the + /// e2e configs give each node a /24 that no other node's interface covers. + /// Judging a replicated address by local topology refuses legitimate peers + /// under the second shape, so this is limited to what a node can assert on + /// its own: an ordinary unicast address that is not one of *this* gateway's. + /// + /// What keeps the peer list coherent is not this check but the uniqueness + /// pass in `kv::import` — no two instances may claim the same address — + /// which holds cluster-wide because it runs over the whole KV contents. + pub fn is_routable_client_ip(&self, ip: Ipv4Addr) -> bool { + if ip.is_unspecified() || ip.is_loopback() || ip.is_multicast() || ip.is_broadcast() { + return false; + } + // This gateway's own addresses: handing them to a peer would point the + // interface's traffic into a tunnel. + if self.ip.addr() == ip || self.ip.broadcast() == ip { + return false; + } + if self.reserved_net.iter().any(|net| net.contains(&ip)) { + return false; + } + true + } } fn validate(ip: Ipv4Net, reserved_net: &[Ipv4Net], client_ip_range: Ipv4Net) -> Result<()> { diff --git a/dstack/gateway/src/kv/import.rs b/dstack/gateway/src/kv/import.rs new file mode 100644 index 000000000..49f6715b7 --- /dev/null +++ b/dstack/gateway/src/kv/import.rs @@ -0,0 +1,318 @@ +// SPDX-FileCopyrightText: © 2024-2025 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Validation boundary between replicated KV records and the local data plane. +//! +//! Every `inst/` record that reaches `ProxyState` — and through it the rendered +//! `wg.conf` — passes through [`accept_instances`] first. The registration-path +//! checks are re-run here because: +//! +//! - a record can arrive from a peer without ever passing through this node's +//! registration RPC, so its checks are not a boundary for synced data; +//! - last-writer-wins replication cannot enforce invariants that span keys, so +//! IP and public-key uniqueness have to be re-established on import; +//! - `wg syncconf` rejects the *whole* config file when a single peer key is +//! malformed, which turns one bad record into a node-wide WireGuard freeze. +//! +//! Validation never aborts the batch: an offending record is skipped and +//! reported, every other instance keeps its routing. + +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::net::Ipv4Addr; + +use anyhow::{bail, ensure, Result}; +use base64::{engine::general_purpose::STANDARD, Engine}; + +use super::InstanceData; +use crate::config::WgConfig; + +/// A WireGuard public key is 32 raw bytes, base64-encoded by `wg`. +const WG_PUBLIC_KEY_BYTES: usize = 32; + +/// Upper bound on identifier fields carried in a KV record. Real values are +/// hex-encoded hashes (40-64 chars); the bound keeps a corrupt record that +/// still decodes from pushing an arbitrarily long string into the logs and the +/// rendered config. +const MAX_ID_LEN: usize = 128; + +/// A record that failed validation, along with the reason. +pub struct RejectedInstance { + pub instance_id: String, + pub reason: anyhow::Error, +} + +/// Records accepted for import, plus the ones that were skipped. +pub struct AcceptedInstances { + pub instances: BTreeMap, + pub rejected: Vec, +} + +/// Validate a WireGuard public key as `wg` itself would accept it. +pub fn validate_wg_public_key(public_key: &str) -> Result<()> { + ensure!(!public_key.is_empty(), "public key is empty"); + ensure!( + !public_key.contains(|c: char| c.is_whitespace() || c.is_control()), + "public key contains whitespace or control characters" + ); + let decoded = STANDARD + .decode(public_key) + .map_err(|err| anyhow::anyhow!("public key is not valid base64: {err}"))?; + ensure!( + decoded.len() == WG_PUBLIC_KEY_BYTES, + "public key decodes to {} bytes, expected {WG_PUBLIC_KEY_BYTES}", + decoded.len() + ); + Ok(()) +} + +fn validate_id(field: &str, value: &str) -> Result<()> { + ensure!(!value.is_empty(), "{field} is empty"); + ensure!( + value.len() <= MAX_ID_LEN, + "{field} is {} bytes, limit is {MAX_ID_LEN}", + value.len() + ); + ensure!( + !value.contains(|c: char| c.is_whitespace() || c.is_control()), + "{field} contains whitespace or control characters" + ); + Ok(()) +} + +/// Per-record checks that do not depend on any other record. +fn validate_instance(wg: &WgConfig, instance_id: &str, data: &InstanceData) -> Result<()> { + validate_id("instance_id", instance_id)?; + validate_id("app_id", &data.app_id)?; + validate_wg_public_key(&data.public_key)?; + // The routable network, not this node's allocation share: in a cluster + // every node carries peers for the CVMs registered on the other nodes, and + // those hold addresses from the other nodes' shares by design. + ensure!( + wg.is_routable_client_ip(data.ip), + "ip {} is outside the WireGuard network", + data.ip + ); + Ok(()) +} + +/// Filter KV instance records down to the ones safe to apply locally. +/// +/// Conflicts on IP or public key are resolved by registration time (oldest +/// wins, ties broken by instance ID) so that every node reaches the same +/// decision from the same KV contents — the local registration path resolves +/// them the same way, by refusing the newcomer. +pub fn accept_instances( + wg: &WgConfig, + records: BTreeMap, +) -> AcceptedInstances { + let mut ordered: Vec<(String, InstanceData)> = records.into_iter().collect(); + ordered.sort_by(|(left_id, left), (right_id, right)| { + left.reg_time + .cmp(&right.reg_time) + .then_with(|| left_id.cmp(right_id)) + }); + + let mut instances = BTreeMap::new(); + let mut rejected = Vec::new(); + let mut claimed_ips: HashMap = HashMap::new(); + let mut claimed_keys: HashSet = HashSet::new(); + + for (instance_id, data) in ordered { + let checked = validate_instance(wg, &instance_id, &data).and_then(|()| { + if let Some(owner) = claimed_ips.get(&data.ip) { + bail!("ip {} is already assigned to instance {owner}", data.ip); + } + if claimed_keys.contains(&data.public_key) { + bail!("public key is already registered to another instance"); + } + Ok(()) + }); + if let Err(reason) = checked { + rejected.push(RejectedInstance { + instance_id, + reason, + }); + continue; + } + claimed_ips.insert(data.ip, instance_id.clone()); + claimed_keys.insert(data.public_key.clone()); + instances.insert(instance_id, data); + } + + AcceptedInstances { + instances, + rejected, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ipnet::Ipv4Net; + + fn wg_config() -> WgConfig { + WgConfig { + public_key: "gateway".to_string(), + private_key: "gateway".to_string(), + listen_port: 51820, + ip: "10.0.0.1/24".parse::().unwrap(), + reserved_net: vec!["10.0.0.0/28".parse::().unwrap()], + client_ip_range: "10.0.0.0/24".parse::().unwrap(), + interface: "wg0".to_string(), + config_path: "/tmp/wg.conf".to_string(), + endpoint: "127.0.0.1:51820".to_string(), + } + } + + fn key(seed: u8) -> String { + STANDARD.encode([seed; WG_PUBLIC_KEY_BYTES]) + } + + fn instance(ip: &str, public_key: &str, reg_time: u64) -> InstanceData { + InstanceData { + app_id: "0123456789abcdef".to_string(), + ip: ip.parse().unwrap(), + public_key: public_key.to_string(), + reg_time, + port_policy: None, + port_policy_hash: String::new(), + admin_port_policy: None, + } + } + + fn accept(records: Vec<(&str, InstanceData)>) -> AcceptedInstances { + let records = records + .into_iter() + .map(|(id, data)| (id.to_string(), data)) + .collect(); + accept_instances(&wg_config(), records) + } + + #[test] + fn accepts_well_formed_records() { + let accepted = accept(vec![ + ("a", instance("10.0.0.20", &key(1), 100)), + ("b", instance("10.0.0.21", &key(2), 200)), + ]); + assert!(accepted.rejected.is_empty()); + assert_eq!(accepted.instances.len(), 2); + } + + #[test] + fn rejects_keys_wg_would_refuse() { + for bad in [ + "", + "not base64!", + &STANDARD.encode([1u8; 16]), + &format!("{}\nEndpoint = 10.0.0.9:1234", key(3)), + ] { + assert!( + validate_wg_public_key(bad).is_err(), + "accepted public key {bad:?}" + ); + } + assert!(validate_wg_public_key(&key(3)).is_ok()); + } + + #[test] + fn one_bad_record_does_not_drop_the_others() { + let accepted = accept(vec![ + ("bad-key", instance("10.0.0.20", "not-a-key", 100)), + ("good", instance("10.0.0.21", &key(2), 200)), + ]); + assert_eq!(accepted.rejected.len(), 1); + assert_eq!(accepted.rejected[0].instance_id, "bad-key"); + assert!(accepted.instances.contains_key("good")); + } + + #[test] + fn rejects_addresses_belonging_to_this_gateway() { + // The gateway's own wg address, an address in its reserved net, and + // non-unicast garbage must never reach the peer list. + for ip in ["10.0.0.1", "10.0.0.5", "127.0.0.1", "224.0.0.1", "0.0.0.0"] { + let accepted = accept(vec![("x", instance(ip, &key(1), 100))]); + assert!(accepted.instances.is_empty(), "accepted ip {ip}"); + } + } + + #[test] + fn duplicate_ip_and_key_claims_resolve_to_the_older_registration() { + let accepted = accept(vec![ + ("new", instance("10.0.0.20", &key(9), 300)), + ("old", instance("10.0.0.20", &key(1), 100)), + ]); + assert_eq!(accepted.instances.len(), 1); + assert!(accepted.instances.contains_key("old")); + + let accepted = accept(vec![ + ("new", instance("10.0.0.21", &key(1), 300)), + ("old", instance("10.0.0.20", &key(1), 100)), + ]); + assert_eq!(accepted.instances.len(), 1); + assert!(accepted.instances.contains_key("old")); + } + + #[test] + fn conflict_resolution_does_not_depend_on_iteration_order() { + let forward = accept(vec![ + ("a", instance("10.0.0.20", &key(1), 100)), + ("b", instance("10.0.0.20", &key(2), 100)), + ]); + let reverse = accept(vec![ + ("b", instance("10.0.0.20", &key(2), 100)), + ("a", instance("10.0.0.20", &key(1), 100)), + ]); + assert_eq!( + forward.instances.keys().collect::>(), + reverse.instances.keys().collect::>() + ); + assert!(forward.instances.contains_key("a")); + } + + #[test] + fn a_cluster_peers_address_is_not_judged_by_this_nodes_pool() { + // Every deployment gives each node its own slice, and a CVM is handed + // every gateway as a WireGuard server — so each node must carry peers + // holding addresses out of the other nodes' slices. The two shapes in + // tree disagree on whether those slices sit inside a node's own + // interface network, so neither the pool nor the interface network can + // decide this. + let shapes = [ + // deploy-to-vmm.sh: /18 pools inside a shared /16 interface. + ("10.8.0.1/16", "10.8.0.0/18", "10.8.0.5", "10.8.64.5"), + // test-run/cluster.sh and e2e/configs: a /24 per node, and no + // node's interface covers another's. + ("10.0.41.1/24", "10.0.41.0/24", "10.0.41.5", "10.0.42.5"), + ]; + for (ip, pool, mine, peers) in shapes { + let gateway_addr = ip.split('/').next().unwrap_or_default(); + let wg = WgConfig { + ip: ip.parse::().unwrap(), + reserved_net: vec![format!("{gateway_addr}/32").parse::().unwrap()], + client_ip_range: pool.parse::().unwrap(), + ..wg_config() + }; + let records = [ + ("mine", instance(mine, &key(1), 100)), + ("peers", instance(peers, &key(2), 100)), + // Still refused: this gateway's own address. + ("steals-gateway-ip", instance(gateway_addr, &key(3), 100)), + ] + .into_iter() + .map(|(id, data)| (id.to_string(), data)) + .collect(); + let accepted = accept_instances(&wg, records); + assert!(accepted.instances.contains_key("mine"), "{ip}"); + assert!( + accepted.instances.contains_key("peers"), + "{ip}: a peer node's instance was refused, which would leave every \ + gateway serving only its own CVMs" + ); + assert!( + !accepted.instances.contains_key("steals-gateway-ip"), + "{ip}" + ); + } + } +} diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 0fa35907e..ba2202848 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -29,11 +29,12 @@ //! - `last_seen/node/{node_id}/{seen_by_node_id}` → u64 (timestamp) mod https_client; +pub mod import; mod sync_service; pub use https_client::{AppIdValidator, HttpsClientConfig}; pub use sync_service::{fetch_peers_from_bootnode, WaveKvSyncService}; -use tracing::warn; +use tracing::{error, warn}; use std::{collections::BTreeMap, net::Ipv4Addr, path::Path, time::Duration}; diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index de0d033a6..91c28f2e0 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -12,7 +12,6 @@ use std::{ use anyhow::{bail, Context, Result}; use auth_client::AuthClient; -use base64::{engine::general_purpose::STANDARD, Engine as _}; use crate::distributed_certbot::DistributedCertBot; use cmd_lib::run_cmd as cmd; @@ -40,10 +39,10 @@ use crate::{ cert_store::{CertResolver, CertStoreBuilder}, config::{Config, TlsConfig}, kv::{ - fetch_peers_from_bootnode, AppIdValidator, CertData, HttpsClientConfig, InstanceData, - KvStore, NodeData, NodeStatus, PortPolicy, WaveKvSyncService, + fetch_peers_from_bootnode, import, AppIdValidator, CertData, HttpsClientConfig, + InstanceData, KvStore, NodeData, NodeStatus, PortPolicy, WaveKvSyncService, }, - models::{InstanceInfo, PortPolicyView, WgConf}, + models::{InstanceInfo, PortPolicyView, WgConf, WgPeer}, proxy::{create_acceptor_with_cert_resolver, AddressGroup, AddressInfo, AppAddressResolver}, }; @@ -52,16 +51,6 @@ mod handshakes; use handshakes::LatestHandshakesCache; -fn validate_wireguard_public_key(public_key: &str) -> Result<()> { - let decoded = STANDARD - .decode(public_key) - .context("invalid WireGuard public key encoding")?; - if decoded.len() != 32 { - bail!("WireGuard public key must decode to 32 bytes"); - } - Ok(()) -} - #[derive(Clone)] pub struct Proxy { _inner: Arc, @@ -180,7 +169,7 @@ impl ProxyInner { instances.len(), nodes.len() ); - let state = build_state_from_kv_store(instances); + let state = build_state_from_kv_store(&config, instances); // This node's own records are written *after* the bootstrap below, not // here. A local write allocates a sequence number, and after a @@ -529,7 +518,7 @@ impl Proxy { if instance_id.is_empty() { bail!("[{instance_id}] instance id is empty"); } - validate_wireguard_public_key(client_public_key) + import::validate_wg_public_key(client_public_key) .with_context(|| format!("[{instance_id}] invalid client public key"))?; let client_info = state .new_client_by_id( @@ -581,11 +570,31 @@ impl Proxy { } } -fn build_state_from_kv_store(instances: BTreeMap) -> ProxyStateMut { +/// Log the records a KV import refused, one line each. +/// +/// A refused record makes its CVM invisible to this node, so it must never be +/// a silent skip. +fn report_rejected_instances(rejected: Vec) { + for import::RejectedInstance { + instance_id, + reason, + } in rejected + { + error!("ignoring KV instance record {instance_id}: {reason:#}"); + } +} + +fn build_state_from_kv_store( + config: &Config, + instances: BTreeMap, +) -> ProxyStateMut { let mut state = ProxyStateMut::default(); + let accepted = import::accept_instances(&config.wg, instances); + report_rejected_instances(accepted.rejected); + // Build instances - for (instance_id, data) in instances { + for (instance_id, data) in accepted.instances { let info = InstanceInfo { id: instance_id.clone(), app_id: data.app_id.clone(), @@ -877,7 +886,9 @@ fn start_wavekv_watch_task(proxy: Proxy) -> Result<()> { } fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()> { - let instances = store.load_all_instances(); + let accepted = import::accept_instances(&proxy.config.wg, store.load_all_instances()); + report_rejected_instances(accepted.rejected); + let instances = accepted.instances; let mut state = proxy.lock(); let mut wg_changed = false; @@ -934,26 +945,7 @@ fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()> impl ProxyState { fn valid_ip(&self, ip: Ipv4Addr) -> bool { - // Must be within client IP range - if !self.config.wg.client_ip_range.contains(&ip) { - return false; - } - if self.config.wg.ip.broadcast() == ip { - return false; - } - if self.config.wg.ip.addr() == ip { - return false; - } - if self - .config - .wg - .reserved_net - .iter() - .any(|net| net.contains(&ip)) - { - return false; - } - true + self.config.wg.is_valid_client_ip(ip) } fn alloc_ip(&mut self) -> Option { for ip in self.config.wg.client_ip_range.hosts() { @@ -983,9 +975,10 @@ impl ProxyState { if app_id.is_empty() { bail!("app_id is empty"); } - if public_key.is_empty() { - bail!("public_key is empty"); - } + // Checked here as well as on the KV import path: a key `wg` refuses + // makes it reject the whole config file, so it must never enter the + // instance table in the first place. + import::validate_wg_public_key(public_key).context("invalid WireGuard public key")?; if self .state .instances @@ -1174,10 +1167,31 @@ impl ProxyState { } fn generate_wg_config(&self) -> Result { + // Last check before the values leave Rust: `wg syncconf` refuses the + // entire file when one peer is malformed, so a record that somehow got + // past the import boundary must cost only its own routing. + let mut peers = Vec::with_capacity(self.state.instances.len()); + for info in self.state.instances.values() { + if let Err(err) = import::validate_wg_public_key(&info.public_key) { + error!("excluding instance {} from wg config: {err:#}", info.id); + continue; + } + if !self.config.wg.is_routable_client_ip(info.ip) { + error!( + "excluding instance {} from wg config: ip {} is outside the wg network", + info.id, info.ip + ); + continue; + } + peers.push(WgPeer { + public_key: &info.public_key, + ip: info.ip, + }); + } let model = WgConf { private_key: &self.config.wg.private_key, listen_port: self.config.wg.listen_port, - peers: (&self.state.instances).into(), + peers, }; Ok(model.render()?) } diff --git a/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config-2.snap b/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config-2.snap index 78baa9249..de3d7e1d2 100644 --- a/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config-2.snap +++ b/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config-2.snap @@ -7,7 +7,7 @@ InstanceInfo { id: "test-id-1", app_id: "app-id-1", ip: 10.0.0.3, - public_key: "test-pubkey-1", + public_key: "dGVzdC1wdWJrZXktMXRlc3QtcHVia2V5LTF0ZXN0LXA=", reg_time: SystemTime { tv_sec: 0, tv_nsec: 0, diff --git a/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config-3.snap b/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config-3.snap index df39fb280..0f9e2a19c 100644 --- a/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config-3.snap +++ b/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config-3.snap @@ -10,11 +10,11 @@ ListenPort = 51820 [Peer] -PublicKey = test-pubkey-0 +PublicKey = dGVzdC1wdWJrZXktMHRlc3QtcHVia2V5LTB0ZXN0LXA= AllowedIPs = 10.0.0.2/32 PersistentKeepalive = 25 [Peer] -PublicKey = test-pubkey-1 +PublicKey = dGVzdC1wdWJrZXktMXRlc3QtcHVia2V5LTF0ZXN0LXA= AllowedIPs = 10.0.0.3/32 PersistentKeepalive = 25 diff --git a/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config.snap b/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config.snap index 1f03bd3fb..c9fa988e2 100644 --- a/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config.snap +++ b/dstack/gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config.snap @@ -7,7 +7,7 @@ InstanceInfo { id: "test-id-0", app_id: "app-id-0", ip: 10.0.0.2, - public_key: "test-pubkey-0", + public_key: "dGVzdC1wdWJrZXktMHRlc3QtcHVia2V5LTB0ZXN0LXA=", reg_time: SystemTime { tv_sec: 0, tv_nsec: 0, diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs index d80db04ba..2ce98cec5 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -6,6 +6,7 @@ use super::*; use crate::config::{load_config_figment, Config, MutualConfig}; use crate::kv::PortFlags; use crate::proxy::port_policy::is_port_allowed; +use base64::Engine as _; use tempfile::TempDir; struct TestState { @@ -88,6 +89,19 @@ async fn wg_config_is_written_owner_only() { ); } +/// Deterministic stand-in for a real WireGuard public key. +/// +/// Registration and the wg-config renderer both reject keys `wg` itself would +/// refuse, so test fixtures have to be 32 base64-encoded bytes like the real +/// thing. +fn test_pubkey(label: &str) -> String { + let mut key = [0u8; 32]; + for (slot, byte) in key.iter_mut().zip(label.bytes().cycle()) { + *slot = byte; + } + base64::engine::general_purpose::STANDARD.encode(key) +} + fn policy(restrict: bool, ports: &[u16]) -> PortPolicy { PortPolicy { ports: ports @@ -100,9 +114,12 @@ fn policy(restrict: bool, ports: &[u16]) -> PortPolicy { #[test] fn test_validate_wireguard_public_key() { - assert!(validate_wireguard_public_key("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=").is_ok()); - assert!(validate_wireguard_public_key("not-a-wireguard-key").is_err()); - assert!(validate_wireguard_public_key("AQID").is_err()); + assert!(crate::kv::import::validate_wg_public_key( + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" + ) + .is_ok()); + assert!(crate::kv::import::validate_wg_public_key("not-a-wireguard-key").is_err()); + assert!(crate::kv::import::validate_wg_public_key("AQID").is_err()); } #[tokio::test] @@ -137,7 +154,7 @@ async fn test_port_policy_restrict_mode_allows_listed_only() { .new_client_by_id( "inst-allow", "app-allow", - "pubkey-allow", + &test_pubkey("pubkey-allow"), "hash-allow", Some(policy(true, &[8080, 9090])), ) @@ -155,7 +172,7 @@ async fn test_port_policy_disabled_allows_all() { .new_client_by_id( "inst-open", "app-open", - "pubkey-open", + &test_pubkey("pubkey-open"), "hash-open", // restrict_mode = false, but with `ports` listed: still open. Some(policy(false, &[8080])), @@ -174,7 +191,7 @@ async fn test_port_policy_unknown_fails_closed() { .new_client_by_id( "inst-legacy", "app-legacy", - "pubkey-legacy", + &test_pubkey("pubkey-legacy"), "hash-legacy", None, ) @@ -199,7 +216,7 @@ async fn test_admin_override_takes_precedence() { .new_client_by_id( "inst-ovr", "app-ovr", - "pubkey-ovr", + &test_pubkey("pubkey-ovr"), "hash-ovr", Some(policy(true, &[8080])), ) @@ -223,7 +240,7 @@ async fn test_admin_override_can_open_what_instance_restricts() { .new_client_by_id( "inst-lock", "app-lock", - "pubkey-lock", + &test_pubkey("pubkey-lock"), "hash-lock", Some(policy(true, &[])), ) @@ -245,7 +262,7 @@ async fn test_clear_admin_override_reverts_to_instance_policy() { .new_client_by_id( "inst-revert", "app-revert", - "pubkey-revert", + &test_pubkey("pubkey-revert"), "hash-revert", Some(policy(true, &[8080])), ) @@ -285,7 +302,7 @@ async fn test_admin_override_survives_compose_hash_change() { .new_client_by_id( "inst-upgrade", "app-upgrade", - "pubkey-upgrade", + &test_pubkey("pubkey-upgrade"), "hash-v1", Some(policy(true, &[8080])), ) @@ -301,7 +318,7 @@ async fn test_admin_override_survives_compose_hash_change() { .new_client_by_id( "inst-upgrade", "app-upgrade", - "pubkey-upgrade", + &test_pubkey("pubkey-upgrade"), "hash-v2", Some(policy(true, &[7070, 8080])), ) @@ -317,14 +334,26 @@ async fn test_config() { let state = create_test_state().await; let mut info = state .lock() - .new_client_by_id("test-id-0", "app-id-0", "test-pubkey-0", "", None) + .new_client_by_id( + "test-id-0", + "app-id-0", + &test_pubkey("test-pubkey-0"), + "", + None, + ) .unwrap(); info.reg_time = SystemTime::UNIX_EPOCH; insta::assert_debug_snapshot!(info); let mut info1 = state .lock() - .new_client_by_id("test-id-1", "app-id-1", "test-pubkey-1", "", None) + .new_client_by_id( + "test-id-1", + "app-id-1", + &test_pubkey("test-pubkey-1"), + "", + None, + ) .unwrap(); info1.reg_time = SystemTime::UNIX_EPOCH; insta::assert_debug_snapshot!(info1); @@ -346,17 +375,17 @@ async fn gateway_top_n_batch_007_cache_health_and_invalidation() { .new_client_by_id( &format!("top-instance-{index}"), "top-app", - &format!("top-key-{index}"), + &test_pubkey(&format!("top-key-{index}")), "", Some(policy(false, &[])), ) .unwrap(); } proxy.handshake_cache.set_for_test(BTreeMap::from([ - ("top-key-0".to_string(), now), - ("top-key-1".to_string(), now - 1), - ("top-key-2".to_string(), now - 2), - ("top-key-3".to_string(), now - 3600), + (test_pubkey("top-key-0"), now), + (test_pubkey("top-key-1"), now - 1), + (test_pubkey("top-key-2"), now - 2), + (test_pubkey("top-key-3"), now - 3600), ])); let selected = proxy.select_top_n_hosts("top-app").unwrap(); let selected_ids = selected @@ -370,10 +399,10 @@ async fn gateway_top_n_batch_007_cache_health_and_invalidation() { assert_eq!(proxy.state.top_n.len(), 1); proxy.handshake_cache.set_for_test(BTreeMap::from([ - ("top-key-0".to_string(), now - 3600), - ("top-key-1".to_string(), now - 3600), - ("top-key-2".to_string(), now - 3600), - ("top-key-3".to_string(), now), + (test_pubkey("top-key-0"), now - 3600), + (test_pubkey("top-key-1"), now - 3600), + (test_pubkey("top-key-2"), now - 3600), + (test_pubkey("top-key-3"), now), ])); let cached = proxy.select_top_n_hosts("top-app").unwrap(); assert_eq!( @@ -388,18 +417,18 @@ async fn gateway_top_n_batch_007_cache_health_and_invalidation() { .new_client_by_id( "top-instance-4", "top-app", - "top-key-4", + &test_pubkey("top-key-4"), "", Some(policy(false, &[])), ) .unwrap(); assert!(proxy.state.top_n.is_empty()); proxy.handshake_cache.set_for_test(BTreeMap::from([ - ("top-key-0".to_string(), now - 3600), - ("top-key-1".to_string(), now - 3600), - ("top-key-2".to_string(), now - 3600), - ("top-key-3".to_string(), now), - ("top-key-4".to_string(), now - 1), + (test_pubkey("top-key-0"), now - 3600), + (test_pubkey("top-key-1"), now - 3600), + (test_pubkey("top-key-2"), now - 3600), + (test_pubkey("top-key-3"), now), + (test_pubkey("top-key-4"), now - 1), ])); let refreshed = proxy.select_top_n_hosts("top-app").unwrap(); assert_eq!(refreshed.len(), 2); @@ -415,3 +444,76 @@ async fn gateway_top_n_batch_007_cache_health_and_invalidation() { assert!(proxy.select_top_n_hosts("other-app").is_err()); } } + +/// Write a record straight into the KV store, bypassing registration, the way +/// a peer's sync round would. +fn sync_from_peer(state: &TestState, instance_id: &str, ip: &str, public_key: &str) { + state + .kv_store + .sync_instance( + instance_id, + &InstanceData { + app_id: "peer-app".to_string(), + ip: ip.parse().unwrap(), + public_key: public_key.to_string(), + reg_time: 1, + port_policy: None, + port_policy_hash: String::new(), + admin_port_policy: None, + }, + ) + .unwrap(); +} + +#[tokio::test] +async fn a_poisoned_peer_record_costs_only_its_own_instance() { + let state = create_test_state().await; + sync_from_peer(&state, "good", "10.0.0.41", &test_pubkey("good-key")); + // A key `wg` refuses makes `wg syncconf` reject the entire config file, so + // this record must never reach ProxyState or the rendered peer list. + sync_from_peer( + &state, + "poisoned", + "10.0.0.42", + "not-a-key\nEndpoint = 1.2.3.4:1", + ); + // The gateway's own wg address, claimed by an instance. + sync_from_peer( + &state, + "steals-gateway-ip", + "10.0.0.1", + &test_pubkey("other"), + ); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + + let proxy = state.lock(); + assert!(proxy.state.instances.contains_key("good")); + assert!(!proxy.state.instances.contains_key("poisoned")); + assert!(!proxy.state.instances.contains_key("steals-gateway-ip")); + + let rendered = proxy.generate_wg_config().unwrap(); + assert!(rendered.contains(&test_pubkey("good-key"))); + assert!(!rendered.contains("Endpoint = 1.2.3.4:1")); +} + +#[tokio::test] +async fn a_cvm_registered_on_another_node_becomes_a_wg_peer_here() { + let state = create_test_state().await; + // What a peer node allocated out of its own slice. `test-run/cluster.sh` + // and the e2e configs give each node a /24 of its own, so a peer's address + // is outside this node's pool *and* outside its interface network — yet + // every CVM is handed every gateway as a WireGuard server, so this node + // still has to carry it. + let peer_ip = "10.0.42.5"; + assert!(!state.config.wg.is_valid_client_ip(peer_ip.parse().unwrap())); + sync_from_peer(&state, "peer-node-cvm", peer_ip, &test_pubkey("far")); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + + let proxy = state.lock(); + assert!( + proxy.state.instances.contains_key("peer-node-cvm"), + "refusing a peer node's instance leaves each gateway serving only its own CVMs" + ); + let rendered = proxy.generate_wg_config().unwrap(); + assert!(rendered.contains(peer_ip), "peer missing from wg.conf"); +} diff --git a/dstack/gateway/src/models.rs b/dstack/gateway/src/models.rs index 318869334..c7c82f0d1 100644 --- a/dstack/gateway/src/models.rs +++ b/dstack/gateway/src/models.rs @@ -6,7 +6,6 @@ use dstack_gateway_rpc::{AcmeInfoResponse, ProxyAccelStatus, StatusResponse}; use rinja::Template; use serde::{Deserialize, Serialize}; use std::{ - collections::{btree_map::Iter, BTreeMap}, net::Ipv4Addr, sync::{ atomic::{AtomicU64, Ordering}, @@ -23,38 +22,6 @@ mod filters { } } -pub struct MapValues<'a, K, V>(pub &'a BTreeMap); -impl Copy for MapValues<'_, K, V> {} -impl Clone for MapValues<'_, K, V> { - fn clone(&self) -> Self { - *self - } -} -impl<'a, K, V> From<&'a BTreeMap> for MapValues<'a, K, V> { - fn from(map: &'a BTreeMap) -> Self { - MapValues(map) - } -} - -pub struct MapValuesIter<'a, K, V>(Iter<'a, K, V>); - -impl<'a, K, V> IntoIterator for MapValues<'a, K, V> { - type Item = &'a V; - type IntoIter = MapValuesIter<'a, K, V>; - - fn into_iter(self) -> Self::IntoIter { - MapValuesIter(self.0.iter()) - } -} - -impl<'a, K, V> Iterator for MapValuesIter<'a, K, V> { - type Item = &'a V; - - fn next(&mut self) -> Option { - self.0.next().map(|(_, v)| v) - } -} - #[derive(Clone, Debug, Serialize, Deserialize)] pub struct InstanceInfo { pub id: String, @@ -155,12 +122,22 @@ impl Drop for EnteredCounter { } } +/// One `[Peer]` stanza of the rendered WireGuard config. +/// +/// Built by the caller rather than borrowed straight from the instance table so +/// that a record `wg` would refuse can be left out — the template renders with +/// `escape = "none"`, and `wg syncconf` rejects the whole file on one bad line. +pub struct WgPeer<'a> { + pub public_key: &'a str, + pub ip: Ipv4Addr, +} + #[derive(Template)] #[template(path = "wg.conf", escape = "none")] pub struct WgConf<'a> { pub private_key: &'a str, pub listen_port: u16, - pub peers: MapValues<'a, String, InstanceInfo>, + pub peers: Vec>, } #[derive(Template)] From 386209ebd0d99bb768dd80277e0aa8ca1b29670e Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:32:27 -0700 Subject: [PATCH 02/13] fix(gateway): drop instances deleted on another node MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `reload_instances_from_kv_store` only ever upserted. An instance recycled or deregistered on node A stayed routable on node B — in ProxyState, in the top-N selection and in B's WireGuard config — until B's own recycle timeout expired, which is 10h by default. The reload now also removes instances that are present locally but gone from the KV store. Records that merely failed validation are left alone: the last known-good state of an instance is better than no state. A registration newer than the grace window is also kept, so an instance whose KV write failed is not evicted before the CVM's next registration refresh. Refs #1029 --- dstack/gateway/src/main_service.rs | 56 +++++++++++++++++++----- dstack/gateway/src/main_service/tests.rs | 49 +++++++++++++++++++++ 2 files changed, 93 insertions(+), 12 deletions(-) diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index 91c28f2e0..e0add7fac 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -885,6 +885,14 @@ fn start_wavekv_watch_task(proxy: Proxy) -> Result<()> { Ok(()) } +/// Grace period protecting a freshly registered local instance from being +/// dropped by the "gone from KV" pass. +/// +/// A registration writes to ProxyState and to the KV store under the same lock, +/// so the two cannot normally disagree — but if that write failed, the instance +/// would otherwise be evicted before the CVM's next registration refresh. +const LOCAL_REGISTRATION_GRACE: Duration = Duration::from_secs(60); + fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()> { let accepted = import::accept_instances(&proxy.config.wg, store.load_all_instances()); report_rejected_instances(accepted.rejected); @@ -892,6 +900,26 @@ fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()> let mut state = proxy.lock(); let mut wg_changed = false; + // Instances deleted (or recycled) on another node must stop being routable + // here too, rather than lingering until this node's own recycle timeout. + // Records that merely failed validation are left alone: the last known-good + // state of an instance is better than no state at all. + let removed: Vec = state + .state + .instances + .iter() + .filter(|(id, info)| { + !instances.contains_key(*id) + && info.reg_time.elapsed().unwrap_or_default() > LOCAL_REGISTRATION_GRACE + }) + .map(|(id, _)| id.clone()) + .collect(); + for instance_id in removed { + info!("WaveKV: instance {instance_id} was deleted remotely, dropping it"); + state.forget_instance(&instance_id); + wg_changed = true; + } + for (instance_id, data) in instances { let new_info = InstanceInfo { id: instance_id.clone(), @@ -1325,18 +1353,12 @@ impl ProxyState { self.handshake_cache.latest(stale_timeout) } - fn remove_instance(&mut self, id: &str) -> Result<()> { - let info = self - .state - .instances - .remove(id) - .context("instance not found")?; - - // Sync deletion to KvStore - if let Err(err) = self.kv_store.sync_delete_instance(id) { - error!("Failed to sync instance deletion to KvStore: {err:?}"); - } - + /// Drop an instance from the local state only. + /// + /// Used when the KV store already says the instance is gone; the syncing + /// counterpart is [`Self::remove_instance`]. + fn forget_instance(&mut self, id: &str) -> Option { + let info = self.state.instances.remove(id)?; self.state.allocated_addresses.remove(&info.ip); self.state.top_n.remove(&info.app_id); if let Some(app_instances) = self.state.apps.get_mut(&info.app_id) { @@ -1345,6 +1367,16 @@ impl ProxyState { self.state.apps.remove(&info.app_id); } } + Some(info) + } + + fn remove_instance(&mut self, id: &str) -> Result<()> { + self.forget_instance(id).context("instance not found")?; + + // Sync deletion to KvStore + if let Err(err) = self.kv_store.sync_delete_instance(id) { + error!("Failed to sync instance deletion to KvStore: {err:?}"); + } Ok(()) } diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs index 2ce98cec5..d265cb90c 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -517,3 +517,52 @@ async fn a_cvm_registered_on_another_node_becomes_a_wg_peer_here() { let rendered = proxy.generate_wg_config().unwrap(); assert!(rendered.contains(peer_ip), "peer missing from wg.conf"); } + +#[tokio::test] +async fn an_instance_deleted_on_another_node_stops_being_routable_here() { + let state = create_test_state().await; + sync_from_peer( + &state, + "peer-instance", + "10.0.0.40", + &test_pubkey("peer-key"), + ); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + assert!(state.lock().state.instances.contains_key("peer-instance")); + + // The remote node recycled the CVM; until the deletion is applied locally + // the deregistered instance keeps receiving proxied traffic. + state + .kv_store + .sync_delete_instance("peer-instance") + .unwrap(); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + + let proxy = state.lock(); + assert!(!proxy.state.instances.contains_key("peer-instance")); + assert!(!proxy.state.apps.contains_key("peer-app")); + assert!(!proxy + .state + .allocated_addresses + .contains(&"10.0.0.40".parse().unwrap())); +} + +#[tokio::test] +async fn a_local_registration_survives_a_reload_that_cannot_see_it_yet() { + let state = create_test_state().await; + state + .lock() + .new_client_by_id("fresh", "fresh-app", &test_pubkey("fresh-key"), "", None) + .unwrap(); + // Stand in for a KV write that failed: the instance exists locally only. + // Evicting it would black-hole a live CVM until it re-registers. + state + .kv_store + .persistent() + .write() + .delete(crate::kv::keys::inst("fresh")) + .unwrap(); + + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + assert!(state.lock().state.instances.contains_key("fresh")); +} From 75bd4dd2136e2709839ca0c6974128fe229cf24c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:33:00 -0700 Subject: [PATCH 03/13] fix(gateway): fail closed on corrupt global KV records `get_acme_credentials()` already distinguished missing from corrupt, but its siblings folded a corrupt record into `None`, which silently changes global behavior: `get_certbot_config()` fell back to the defaults, switching `acme_url` to Let's Encrypt production and resetting every renewal interval; a corrupt `dns_cred_default` or per-credential record made the certbot issue through the wrong DNS account or none at all; a corrupt ACME attestation reported an attested account as unattested. The three-state read (missing/tombstoned vs. decodable vs. corrupt) is now a `decode_strict` helper on the KV codec, applied to every global record whose corruption must not silently change behavior. The renewal loop skips its round and retries instead of proceeding with defaults. Refs #1029 --- dstack/gateway/src/admin_service.rs | 24 ++-- dstack/gateway/src/distributed_certbot.rs | 22 ++-- dstack/gateway/src/kv/mod.rs | 137 ++++++++++++++++++---- dstack/gateway/src/main_service.rs | 16 ++- 4 files changed, 152 insertions(+), 47 deletions(-) diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index c1e988d19..11313356a 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -310,7 +310,7 @@ impl AdminRpc for AdminRpcHandler { .into_iter() .map(dns_cred_to_proto) .collect(); - let default_id = kv_store.get_default_dns_credential_id(); + let default_id = kv_store.get_default_dns_credential_id()?; Ok(ListDnsCredentialsResponse { credentials, default_id, @@ -323,7 +323,7 @@ impl AdminRpc for AdminRpcHandler { ) -> Result { let kv_store = self.state.kv_store(); let cred = kv_store - .get_dns_credential(&request.id) + .get_dns_credential(&request.id)? .context("dns credential not found")?; Ok(dns_cred_to_proto(cred)) } @@ -383,7 +383,7 @@ impl AdminRpc for AdminRpcHandler { let kv_store = self.state.kv_store(); let mut cred = kv_store - .get_dns_credential(&request.id) + .get_dns_credential(&request.id)? .context("dns credential not found")?; // Update name if provided @@ -414,7 +414,7 @@ impl AdminRpc for AdminRpcHandler { let kv_store = self.state.kv_store(); // Check if this is the default credential - if let Some(default_id) = kv_store.get_default_dns_credential_id() { + if let Some(default_id) = kv_store.get_default_dns_credential_id()? { if default_id == request.id { bail!("cannot delete the default DNS credential; set a different default first"); } @@ -438,8 +438,12 @@ impl AdminRpc for AdminRpcHandler { async fn get_default_dns_credential(self) -> Result { let kv_store = self.state.kv_store(); - let default_id = kv_store.get_default_dns_credential_id().unwrap_or_default(); - let credential = kv_store.get_default_dns_credential().map(dns_cred_to_proto); + let default_id = kv_store + .get_default_dns_credential_id()? + .unwrap_or_default(); + let credential = kv_store + .get_default_dns_credential()? + .map(dns_cred_to_proto); Ok(GetDefaultDnsCredentialResponse { default_id, credential, @@ -454,7 +458,7 @@ impl AdminRpc for AdminRpcHandler { // Verify the credential exists kv_store - .get_dns_credential(&request.id) + .get_dns_credential(&request.id)? .context("dns credential not found")?; kv_store.set_default_dns_credential_id(&request.id)?; @@ -609,7 +613,7 @@ impl AdminRpc for AdminRpcHandler { // ==================== Global Certbot Configuration ==================== async fn get_certbot_config(self) -> Result { - let config = self.state.kv_store().get_certbot_config(); + let config = self.state.kv_store().get_certbot_config()?; Ok(CertbotConfigResponse { renew_interval_secs: config.renew_interval.as_secs(), renew_before_expiration_secs: config.renew_before_expiration.as_secs(), @@ -620,7 +624,7 @@ impl AdminRpc for AdminRpcHandler { async fn set_certbot_config(self, request: SetCertbotConfigRequest) -> Result<()> { let kv_store = self.state.kv_store(); - let mut config = kv_store.get_certbot_config(); + let mut config = kv_store.get_certbot_config()?; // Update only the fields that are specified if let Some(secs) = request.renew_interval_secs { @@ -851,7 +855,7 @@ fn proto_to_zt_domain_config( // Validate DNS credential if specified if let Some(ref cred_id) = dns_cred_id { kv_store - .get_dns_credential(cred_id) + .get_dns_credential(cred_id)? .context("specified dns credential not found")?; } diff --git a/dstack/gateway/src/distributed_certbot.rs b/dstack/gateway/src/distributed_certbot.rs index a7dc07a95..715dcc3bf 100644 --- a/dstack/gateway/src/distributed_certbot.rs +++ b/dstack/gateway/src/distributed_certbot.rs @@ -100,7 +100,7 @@ impl DistributedCertBot { async fn do_rotate_acme_credentials(&self) -> Result<(String, usize)> { let configs = self.kv_store.list_zt_domain_configs(); - let certbot_config = self.config(); + let certbot_config = self.config()?; let acme_url = if certbot_config.acme_url.is_empty() { DEFAULT_ACME_URL } else { @@ -206,8 +206,12 @@ impl DistributedCertBot { Ok((account_uri, total)) } - /// Get the current certbot configuration from KV store - fn config(&self) -> crate::kv::GlobalCertbotConfig { + /// Get the current certbot configuration from KV store. + /// + /// Propagates a corrupt record instead of falling back to the defaults: + /// the default `acme_url` is Let's Encrypt production, so a silent + /// fallback would move issuance to a different ACME server. + fn config(&self) -> Result { self.kv_store.get_certbot_config() } @@ -337,7 +341,7 @@ impl DistributedCertBot { } else if let Some(ref data) = cert_data { let now = now_secs(); let expires_in = data.not_after.saturating_sub(now); - expires_in < self.config().renew_before_expiration.as_secs() + expires_in < self.config()?.renew_before_expiration.as_secs() } else { true }; @@ -422,7 +426,7 @@ impl DistributedCertBot { wildcard_domain ); let cert_pem = tokio::time::timeout( - self.config().renew_timeout, + self.config()?.renew_timeout, acme_client.request_new_certificate(&key_pem, &[wildcard_domain]), ) .await @@ -476,7 +480,7 @@ impl DistributedCertBot { wildcard_domain ); let new_cert_pem = tokio::time::timeout( - self.config().renew_timeout, + self.config()?.renew_timeout, // Note: we request a new cert rather than renew, since we have a new key acme_client.request_new_certificate(&key_pem, &[wildcard_domain]), ) @@ -524,7 +528,7 @@ impl DistributedCertBot { let dns01_client = self.dns_client(domain, &dns_cred).await?; // Use ACME URL from certbot config, fall back to default if not set - let config = self.config(); + let config = self.config()?; let acme_url = if config.acme_url.is_empty() { DEFAULT_ACME_URL } else { @@ -714,11 +718,11 @@ impl DistributedCertBot { fn dns_credential_for(kv_store: &KvStore, config: &ZtDomainConfig) -> Result { if let Some(ref cred_id) = config.dns_cred_id { kv_store - .get_dns_credential(cred_id) + .get_dns_credential(cred_id)? .context("specified DNS credential not found") } else { kv_store - .get_default_dns_credential() + .get_default_dns_credential()? .context("no default DNS credential configured") } } diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index ba2202848..62274d87c 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -424,6 +424,7 @@ pub fn decode Deserialize<'de>>(bytes: &[u8]) -> Result { trait GetPutCodec { fn decode serde::Deserialize<'de>>(&self, key: &str) -> Option; + fn decode_strict serde::Deserialize<'de>>(&self, key: &str) -> Result>; fn put_encoded(&mut self, key: String, value: &T) -> Result<()>; fn iter_decoded serde::Deserialize<'de>>( &self, @@ -447,6 +448,27 @@ impl GetPutCodec for NodeState { }) } + /// Three-state read: `Ok(None)` for a key that is missing or tombstoned, + /// `Ok(Some)` for a decodable value, `Err` for a stored value that no + /// longer decodes. + /// + /// [`Self::decode`] folds corruption into `None`, which is right for + /// per-instance records (skip the bad one, keep serving the rest) and + /// wrong for global records, where "absent" means "apply the default" and + /// a corrupt record would silently change cluster-wide behavior. + fn decode_strict serde::Deserialize<'de>>(&self, key: &str) -> Result> { + let Some(entry) = self.get(key) else { + return Ok(None); + }; + // A `None` value is a tombstone: the key was deliberately deleted. + let Some(value) = entry.value.as_ref() else { + return Ok(None); + }; + decode(value) + .map(Some) + .with_context(|| format!("corrupt record at KV key {key}")) + } + fn put_encoded(&mut self, key: String, value: &T) -> Result<()> { self.put(key.clone(), encode(value)?) .with_context(|| format!("failed to put key {key}"))?; @@ -777,9 +799,15 @@ impl KvStore { // ==================== DNS Credential Management ==================== - /// Get a DNS credential by ID - pub fn get_dns_credential(&self, cred_id: &str) -> Option { - self.persistent.read().decode(&keys::dns_cred(cred_id)) + /// Get a DNS credential by ID. + /// + /// Fails closed on a corrupt record: silently reading it as "no such + /// credential" would make the certbot fall back to the default credential + /// and issue the domain's certificate through the wrong DNS account. + pub fn get_dns_credential(&self, cred_id: &str) -> Result> { + self.persistent + .read() + .decode_strict(&keys::dns_cred(cred_id)) } /// Save a DNS credential @@ -804,9 +832,12 @@ impl KvStore { .collect() } - /// Get the default DNS credential ID - pub fn get_default_dns_credential_id(&self) -> Option { - self.persistent.read().decode(keys::DNS_CRED_DEFAULT) + /// Get the default DNS credential ID. + /// + /// Fails closed on a corrupt record for the same reason as + /// [`Self::get_dns_credential`]. + pub fn get_default_dns_credential_id(&self) -> Result> { + self.persistent.read().decode_strict(keys::DNS_CRED_DEFAULT) } /// Set the default DNS credential ID @@ -818,19 +849,26 @@ impl KvStore { } /// Get the default DNS credential (resolves the ID to the actual credential) - pub fn get_default_dns_credential(&self) -> Option { - let cred_id = self.get_default_dns_credential_id()?; + pub fn get_default_dns_credential(&self) -> Result> { + let Some(cred_id) = self.get_default_dns_credential_id()? else { + return Ok(None); + }; self.get_dns_credential(&cred_id) } // ==================== Global Certbot Config ==================== - /// Get global certbot configuration (returns default if not set) - pub fn get_certbot_config(&self) -> GlobalCertbotConfig { - self.persistent + /// Get global certbot configuration (returns default if not set). + /// + /// Fails closed on a corrupt record: falling back to the defaults would + /// silently switch `acme_url` back to Let's Encrypt production and reset + /// every renewal interval on this node. + pub fn get_certbot_config(&self) -> Result { + Ok(self + .persistent .read() - .decode(keys::GLOBAL_CERTBOT_CONFIG) - .unwrap_or_default() + .decode_strict(keys::GLOBAL_CERTBOT_CONFIG)? + .unwrap_or_default()) } /// Set global certbot configuration @@ -969,16 +1007,9 @@ impl KvStore { /// Treating corruption as absence would silently register a fresh ACME /// account that the existing account-bound CAA records refuse. pub fn get_acme_credentials(&self) -> Result> { - let state = self.persistent.read(); - let Some(entry) = state.get(keys::GLOBAL_ACME_CREDENTIALS) else { - return Ok(None); - }; - // A `None` value is a tombstone: the key was deliberately deleted. - let Some(value) = entry.value.as_ref() else { - return Ok(None); - }; - decode(value) - .map(Some) + self.persistent + .read() + .decode_strict(keys::GLOBAL_ACME_CREDENTIALS) .context("corrupt ACME credentials record in KvStore") } @@ -990,9 +1021,15 @@ impl KvStore { Ok(()) } - /// Get global ACME attestation (TDX quote of account URI) - pub fn get_acme_attestation(&self) -> Option { - self.persistent.read().decode(keys::GLOBAL_ACME_ATTESTATION) + /// Get global ACME attestation (TDX quote of account URI). + /// + /// Fails closed on a corrupt record: reporting "no attestation" for an + /// account that does have one lets a verifier conclude the ACME account is + /// unattested. + pub fn get_acme_attestation(&self) -> Result> { + self.persistent + .read() + .decode_strict(keys::GLOBAL_ACME_ATTESTATION) } /// Save global ACME attestation @@ -1428,6 +1465,54 @@ mod decompression_tests { } } +#[cfg(test)] +mod corruption_tests { + use super::*; + + fn test_kv(data_dir: &std::path::Path) -> KvStore { + KvStore::new(1, vec![], data_dir).expect("failed to create kv store") + } + + fn put_raw(kv: &KvStore, key: &str, value: &[u8]) { + kv.persistent + .write() + .put(key.to_string(), value.to_vec()) + .expect("raw put should succeed"); + } + + #[test] + fn a_corrupt_certbot_config_does_not_read_as_the_default() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let kv = test_kv(dir.path()); + // Absent means "use the defaults" — that part must keep working. + let default = kv + .get_certbot_config() + .expect("missing key should not error"); + assert!(default.acme_url.is_empty()); + + kv.set_certbot_config(&GlobalCertbotConfig { + acme_url: "https://acme-staging.example/directory".to_string(), + ..Default::default() + }) + .expect("save should succeed"); + put_raw(&kv, keys::GLOBAL_CERTBOT_CONFIG, b"not-messagepack"); + // Reading the corrupt record as the default would silently move + // issuance back to Let's Encrypt production. + assert!(kv.get_certbot_config().is_err()); + } + + #[test] + fn corrupt_global_records_fail_closed() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let kv = test_kv(dir.path()); + put_raw(&kv, keys::GLOBAL_ACME_ATTESTATION, b"not-messagepack"); + put_raw(&kv, keys::DNS_CRED_DEFAULT, b"not-messagepack"); + assert!(kv.get_acme_attestation().is_err()); + assert!(kv.get_default_dns_credential_id().is_err()); + assert!(kv.get_default_dns_credential().is_err()); + } +} + #[cfg(test)] mod peer_url_tests { use super::validate_peer_url; diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index e0add7fac..0a13b50ff 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -456,7 +456,9 @@ impl Proxy { // The account URI comes from the published credentials; the attestation // record is written best-effort and may lag behind a rotation, so it // only supplies the quote when it matches the current account. - let attestation = kv_store.get_acme_attestation(); + let attestation = kv_store + .get_acme_attestation() + .context("failed to read the ACME account attestation")?; let account_uri = kv_store .get_acme_credentials() .context("call RotateAcmeCredentials to replace the stored ACME credentials")? @@ -647,7 +649,17 @@ async fn start_certbot_task(proxy: Proxy) { loop { // Get current config from KV store (allows dynamic updates) - let renew_interval = proxy.kv_store.get_certbot_config().renew_interval; + let renew_interval = match proxy.kv_store.get_certbot_config() { + Ok(config) => config.renew_interval, + Err(err) => { + // Falling back to the defaults here would switch acme_url + // back to Let's Encrypt production; wait for an operator to + // repair the record instead. + error!("failed to read certbot config, skipping renewal round: {err:?}"); + tokio::time::sleep(Duration::from_secs(60)).await; + continue; + } + }; if renew_interval.is_zero() { // Check again later if disabled tokio::time::sleep(Duration::from_secs(60)).await; From 254474af9738e964fffb7f6006fe84a0844852eb Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:33:24 -0700 Subject: [PATCH 04/13] fix(gateway): ignore future-dated handshake and last_seen observations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `handshake/` and `last_seen/` records are wall-clock seconds written by whichever node made the observation, and the gateway aggregates them with `max`. One node with a fast clock — or a single corrupt record near `u64::MAX` — therefore kept a dead CVM "alive" on every node in the cluster: `recycle()` never fired and top-N routing kept steering traffic at it, with no way to correct the record until real time caught up. Observations dated more than 5 minutes ahead of local time are now dropped on read, logged with a count. The allowance is well above NTP-synced drift and well below the recycle timeout. Refs #1029 --- dstack/gateway/src/kv/mod.rs | 116 ++++++++++++++++++++++++++++++++--- 1 file changed, 108 insertions(+), 8 deletions(-) diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 62274d87c..accd51086 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -407,6 +407,47 @@ pub fn gunzip_bounded(data: &[u8], limit: usize) -> Result> { Ok(out) } +/// How far into the future a replicated observation may be timestamped before +/// this node ignores it. +/// +/// `handshake/` and `last_seen/` records are wall-clock seconds written by +/// whichever node made the observation, and the gateway aggregates them with +/// `max`. Without a horizon, a single node with a fast clock — or one corrupt +/// record near `u64::MAX` — keeps a dead CVM "alive" on every node forever: +/// `recycle()` never fires and top-N routing keeps steering traffic at it. +/// 5 minutes is well above the drift between NTP-synced hosts and well below +/// the recycle timeout. +pub const MAX_CLOCK_DRIFT_SECS: u64 = 300; + +fn now_secs() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +/// Drop observations timestamped beyond [`MAX_CLOCK_DRIFT_SECS`] into the +/// future, logging once per call with the number dropped. +fn drop_future_observations( + observations: impl Iterator, + timestamp: impl Fn(&T) -> u64, + kind: &str, +) -> Vec { + let horizon = now_secs().saturating_add(MAX_CLOCK_DRIFT_SECS); + let mut dropped = 0usize; + let kept = observations + .filter(|item| { + let plausible = timestamp(item) <= horizon; + dropped += usize::from(!plausible); + plausible + }) + .collect(); + if dropped > 0 { + warn!("ignored {dropped} {kind} observation(s) dated more than {MAX_CLOCK_DRIFT_SECS}s ahead of local time"); + } + kept +} + /// Encode a KV value as MessagePack. /// /// Structs are encoded as maps keyed by field name rather than as positional @@ -669,9 +710,13 @@ impl KvStore { Ok(()) } - /// Get all handshake observations for an instance (from all nodes) + /// Get all handshake observations for an instance (from all nodes). + /// + /// Observations dated into the future are dropped; see + /// [`MAX_CLOCK_DRIFT_SECS`]. pub fn get_instance_handshakes(&self, instance_id: &str) -> BTreeMap { - self.ephemeral + let observations = self + .ephemeral .read() .iter_decoded(&keys::handshake_prefix(instance_id)) .filter_map(|(key, ts)| { @@ -679,14 +724,22 @@ impl KvStore { let observer: NodeId = suffix.parse().ok()?; Some((observer, ts)) }) + .collect::>(); + drop_future_observations(observations.into_iter(), |(_, ts)| *ts, "handshake") + .into_iter() .collect() } - /// Get the latest handshake timestamp for an instance (max across all nodes) + /// Get the latest handshake timestamp for an instance (max across all + /// nodes), ignoring future-dated observations. pub fn get_instance_latest_handshake(&self, instance_id: &str) -> Option { - self.ephemeral + let observations = self + .ephemeral .read() .iter_decoded_values(&keys::handshake_prefix(instance_id)) + .collect::>(); + drop_future_observations(observations.into_iter(), |ts| *ts, "handshake") + .into_iter() .max() } @@ -698,9 +751,10 @@ impl KvStore { Ok(()) } - /// Get all observations of a node's last_seen + /// Get all observations of a node's last_seen, ignoring future-dated ones. pub fn get_node_last_seen_by_all(&self, node_id: NodeId) -> BTreeMap { - self.ephemeral + let observations = self + .ephemeral .read() .iter_decoded(&keys::last_seen_node_prefix(node_id)) .filter_map(|(key, ts)| { @@ -708,14 +762,22 @@ impl KvStore { let seen_by: NodeId = suffix.parse().ok()?; Some((seen_by, ts)) }) + .collect::>(); + drop_future_observations(observations.into_iter(), |(_, ts)| *ts, "node last_seen") + .into_iter() .collect() } - /// Get the latest last_seen timestamp for a node (max across all observers) + /// Get the latest last_seen timestamp for a node (max across all + /// observers), ignoring future-dated observations. pub fn get_node_latest_last_seen(&self, node_id: NodeId) -> Option { - self.ephemeral + let observations = self + .ephemeral .read() .iter_decoded_values(&keys::last_seen_node_prefix(node_id)) + .collect::>(); + drop_future_observations(observations.into_iter(), |ts| *ts, "node last_seen") + .into_iter() .max() } @@ -1511,6 +1573,44 @@ mod corruption_tests { assert!(kv.get_default_dns_credential_id().is_err()); assert!(kv.get_default_dns_credential().is_err()); } + + #[test] + fn future_dated_observations_are_ignored() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let kv = test_kv(dir.path()); + let now = now_secs(); + + kv.ephemeral + .write() + .put_encoded(keys::handshake("cvm", 2), &(now.saturating_sub(30))) + .unwrap(); + kv.ephemeral + .write() + .put_encoded(keys::handshake("cvm", 3), &u64::MAX) + .unwrap(); + // A peer with a broken clock must not keep a dead CVM alive forever. + let latest = kv + .get_instance_latest_handshake("cvm") + .expect("the plausible observation should survive"); + assert!(latest <= now, "kept a future-dated handshake: {latest}"); + assert_eq!(kv.get_instance_handshakes("cvm").len(), 1); + + kv.ephemeral + .write() + .put_encoded(keys::last_seen_node(7, 3), &u64::MAX) + .unwrap(); + assert_eq!(kv.get_node_latest_last_seen(7), None); + assert!(kv.get_node_last_seen_by_all(7).is_empty()); + + // Drift within the allowance stays usable: nodes are not perfectly + // synchronized and dropping every slightly-ahead record would make + // instances look stale. + kv.ephemeral + .write() + .put_encoded(keys::handshake("cvm", 4), &(now + MAX_CLOCK_DRIFT_SECS / 2)) + .unwrap(); + assert_eq!(kv.get_instance_handshakes("cvm").len(), 2); + } } #[cfg(test)] From f17e862016a4135132ff972d403f04e412c7215e Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:33:55 -0700 Subject: [PATCH 05/13] fix(gateway): quarantine an unreadable KV data dir instead of refusing to start MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Node::new_with_persistence` hard-fails on a checksum or deserialize error in the WAL, so a torn tail — the normal artifact of a crash — kept the gateway from starting at all, and the node served no traffic until an operator intervened. The persistent state is a cache: every record is replicated on the peers and re-fetched by the sync service. `KvStore::new` now moves the unreadable directory to `.corrupt.` and starts empty, logging loudly. Nothing is deleted, so the original bytes stay available for post-mortem. Refs #1029 --- dstack/gateway/src/kv/mod.rs | 99 ++++++++++++++++++++++++++++++++++-- 1 file changed, 95 insertions(+), 4 deletions(-) diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index accd51086..6662b07c9 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -564,15 +564,36 @@ pub struct KvStore { } impl KvStore { - /// Create a new sync store + /// Create a new sync store. + /// + /// If the on-disk WAL/snapshot cannot be read, the data directory is moved + /// aside and the store starts empty rather than refusing to boot: the + /// persistent state is replicated on every peer, a torn WAL tail is the + /// normal artifact of a crash, and a gateway that cannot start serves no + /// traffic at all. Nothing is deleted — the unreadable directory is kept + /// under `.corrupt.` for inspection. pub fn new( my_node_id: NodeId, peer_ids: Vec, data_dir: impl AsRef, ) -> Result { - let persistent = - Node::new_with_persistence(my_node_id, peer_ids.clone(), data_dir.as_ref()) - .context("failed to create persistent wavekv node")?; + let data_dir = data_dir.as_ref(); + let persistent = match Node::new_with_persistence(my_node_id, peer_ids.clone(), data_dir) { + Ok(node) => node, + Err(err) => { + let quarantined = quarantine_data_dir(data_dir).context( + "failed to open the WaveKV data dir and failed to move it aside for recovery", + )?; + error!( + "WaveKV data dir {} is unreadable ({err:#}); moved it to {} and started empty — \ + state will be re-fetched from peers", + data_dir.display(), + quarantined.display(), + ); + Node::new_with_persistence(my_node_id, peer_ids.clone(), data_dir) + .context("failed to create persistent wavekv node on a fresh data dir")? + } + }; // Get peers from persistent store (may have been restored from WAL) // and include them when creating ephemeral store @@ -1265,6 +1286,37 @@ impl KvStore { } } +/// Move an unreadable WaveKV data dir aside, returning the new path. +/// +/// Renaming keeps the bytes for post-mortem analysis and guarantees the +/// gateway never starts on half-readable state. +fn quarantine_data_dir(data_dir: &Path) -> Result { + anyhow::ensure!( + data_dir.exists(), + "WaveKV data dir {} does not exist", + data_dir.display() + ); + let stamp = now_secs(); + for attempt in 0..u32::MAX { + let suffix = if attempt == 0 { + format!("corrupt.{stamp}") + } else { + format!("corrupt.{stamp}.{attempt}") + }; + let mut target = data_dir.as_os_str().to_owned(); + target.push("."); + target.push(&suffix); + let target = std::path::PathBuf::from(target); + if target.exists() { + continue; + } + std::fs::rename(data_dir, &target) + .with_context(|| format!("failed to rename {}", data_dir.display()))?; + return Ok(target); + } + anyhow::bail!("no free quarantine path for {}", data_dir.display()) +} + fn validate_peer_url(url: &str) -> Result<()> { let parsed = reqwest::Url::parse(url).context("invalid peer URL")?; anyhow::ensure!( @@ -1611,6 +1663,45 @@ mod corruption_tests { .unwrap(); assert_eq!(kv.get_instance_handshakes("cvm").len(), 2); } + + #[test] + fn an_unreadable_data_dir_is_quarantined_instead_of_blocking_startup() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let data_dir = dir.path().join("kv"); + { + let kv = test_kv(&data_dir); + kv.sync_instance( + "cvm", + &InstanceData { + app_id: "app".to_string(), + ip: "10.0.0.20".parse().unwrap(), + public_key: "key".to_string(), + reg_time: 1, + port_policy: None, + port_policy_hash: String::new(), + admin_port_policy: None, + }, + ) + .expect("sync should succeed"); + kv.persist_if_dirty().expect("persist should succeed"); + } + // A torn WAL tail is the normal artifact of a crash and every record is + // replicated, so it must not keep the gateway from booting. + std::fs::write(data_dir.join("node_1.wal"), b"garbage").expect("failed to corrupt wal"); + + let kv = KvStore::new(1, vec![], &data_dir).expect("startup must survive a corrupt wal"); + assert!(kv.load_all_instances().is_empty()); + let quarantined: Vec<_> = std::fs::read_dir(dir.path()) + .expect("failed to read temp dir") + .filter_map(|entry| entry.ok()) + .filter(|entry| entry.file_name().to_string_lossy().contains(".corrupt.")) + .collect(); + assert_eq!( + quarantined.len(), + 1, + "the unreadable data dir must be kept for inspection" + ); + } } #[cfg(test)] From 6a25b18ce60b6c902aa8d4dade7451ec865ae967 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:34:19 -0700 Subject: [PATCH 06/13] fix(gateway): skip zt-domain configs whose key and value disagree MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `list_zt_domain_configs` returned the decoded value without checking it against the `cert/{domain}/config` key it was filed under. Everything downstream — certificate issuance, the DNS-01 challenge, `cert/{domain}/data` — is driven by the value, so one poisoned record could point the certbot at a domain nobody configured. Refs #1029 --- dstack/gateway/src/kv/mod.rs | 58 +++++++++++++++++++++++++++++++++--- 1 file changed, 54 insertions(+), 4 deletions(-) diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 6662b07c9..dca06d86e 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -987,7 +987,12 @@ impl KvStore { Ok(()) } - /// List all ZT-Domain configurations + /// List all ZT-Domain configurations. + /// + /// A record whose `domain` disagrees with the domain in its key is + /// skipped: everything downstream (certificate issuance, DNS-01 challenge, + /// `cert/{domain}/data`) is driven by the value, so honouring it would let + /// one poisoned record request a certificate for an unrelated domain. pub fn list_zt_domain_configs(&self) -> Vec { let state = self.persistent.read(); state @@ -998,13 +1003,22 @@ impl KvStore { return None; } let value = entry.value.as_ref()?; - match decode(value) { - Ok(config) => Some(config), + let config: ZtDomainConfig = match decode(value) { + Ok(config) => config, Err(e) => { warn!("failed to decode cert config for key {key}: {e:?}"); - None + return None; } + }; + let key_domain = keys::parse_cert_domain(key)?; + if key_domain != config.domain { + warn!( + "skipping cert config at key {key}: record claims domain {}", + config.domain + ); + return None; } + Some(config) }) .collect() } @@ -1702,6 +1716,42 @@ mod corruption_tests { "the unreadable data dir must be kept for inspection" ); } + + #[test] + fn a_cert_config_that_disagrees_with_its_key_is_skipped() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let kv = test_kv(dir.path()); + kv.save_zt_domain_config(&ZtDomainConfig { + domain: "good.example".to_string(), + dns_cred_id: None, + port: 443, + node: None, + priority: 0, + }) + .expect("save should succeed"); + // Same record filed under another domain's key: honouring the value + // would request a certificate for a domain nobody configured. + kv.persistent + .write() + .put_encoded( + keys::zt_domain_config("victim.example"), + &ZtDomainConfig { + domain: "attacker.example".to_string(), + dns_cred_id: None, + port: 443, + node: None, + priority: 100, + }, + ) + .expect("raw put should succeed"); + + let domains: Vec = kv + .list_zt_domain_configs() + .into_iter() + .map(|c| c.domain) + .collect(); + assert_eq!(domains, vec!["good.example".to_string()]); + } } #[cfg(test)] From bc0fdb331a85fd6813aac49b4696c5ca94565017 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 11 Aug 2026 00:27:04 -0700 Subject: [PATCH 07/13] fix(gateway): keep instances whose replicated record is unreadable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The reload pass claimed that "records that merely failed validation are left alone", but it filtered on `!accepted.instances.contains_key(id)` alone, and a rejected record is absent from `accepted.instances`. So a record that stopped validating — and, because `load_all_instances` folded decode failures into "key not present", one that stopped decoding too — evicted a healthy instance from ProxyState after the 60s grace, taking its wg peer with it. A single corrupt or hostile record was enough to black-hole a live CVM until it re-registered. `load_all_instances` now returns `LoadedInstances`, which keeps undecodable records separate from absent ones, and `import` labels every refusal with a `Rejection`: - `Unusable` (fails validation, or does not decode) says nothing about whether the instance still exists, so the instance keeps whatever the data plane already holds for it; - `LostConflict` (a well-formed record that lost an IP or key conflict to an older registration) does say the address belongs to someone else, so the loser stops being routable — keeping it would put the same address in `wg.conf` twice. The removal pass now exempts only the first kind. Three tests cover the split: a record that stops validating and one that stops decoding both keep their instance, while a conflict loser is still dropped. --- dstack/gateway/src/debug_service.rs | 1 + dstack/gateway/src/kv/import.rs | 163 ++++++++++++++++++----- dstack/gateway/src/kv/mod.rs | 80 +++++++++-- dstack/gateway/src/main_service.rs | 32 +++-- dstack/gateway/src/main_service/tests.rs | 76 ++++++++++- 5 files changed, 295 insertions(+), 57 deletions(-) diff --git a/dstack/gateway/src/debug_service.rs b/dstack/gateway/src/debug_service.rs index b00176a14..703ecb0d8 100644 --- a/dstack/gateway/src/debug_service.rs +++ b/dstack/gateway/src/debug_service.rs @@ -85,6 +85,7 @@ impl DebugRpc for DebugRpcHandler { // Get all instances let instances: Vec = kv_store .load_all_instances() + .decoded .into_iter() .map(|(instance_id, data)| InstanceEntry { instance_id, diff --git a/dstack/gateway/src/kv/import.rs b/dstack/gateway/src/kv/import.rs index 49f6715b7..e21ef896f 100644 --- a/dstack/gateway/src/kv/import.rs +++ b/dstack/gateway/src/kv/import.rs @@ -17,14 +17,20 @@ //! //! Validation never aborts the batch: an offending record is skipped and //! reported, every other instance keeps its routing. +//! +//! Refusals are not all the same, so [`Rejection`] tells the caller which kind +//! it is holding. A record this node cannot make sense of says nothing about +//! whether its instance still exists, so the instance keeps the state the data +//! plane already has; a record that lost an IP or key conflict says the address +//! belongs to someone else, so its instance has to stop being routable. use std::collections::{BTreeMap, HashMap, HashSet}; use std::net::Ipv4Addr; -use anyhow::{bail, ensure, Result}; +use anyhow::{ensure, Result}; use base64::{engine::general_purpose::STANDARD, Engine}; -use super::InstanceData; +use super::{InstanceData, LoadedInstances}; use crate::config::WgConfig; /// A WireGuard public key is 32 raw bytes, base64-encoded by `wg`. @@ -36,10 +42,27 @@ const WG_PUBLIC_KEY_BYTES: usize = 32; /// rendered config. const MAX_ID_LEN: usize = 128; +/// Whether a refused record should also cost the instance the state the data +/// plane already holds for it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Rejection { + /// The record is unusable: it fails validation, or its bytes no longer + /// decode. Either way this node cannot tell what the instance looks like + /// now, and the last known-good state is a better answer than none — so + /// whatever the data plane already knows about the instance stays. + Unusable, + /// The record is well-formed but lost a uniqueness conflict to an older + /// registration. Here the winner genuinely owns the IP or the key, so the + /// loser has to stop being routable; keeping it would put the same address + /// in `wg.conf` twice and hand it traffic that belongs to the winner. + LostConflict, +} + /// A record that failed validation, along with the reason. pub struct RejectedInstance { pub instance_id: String, pub reason: anyhow::Error, + pub rejection: Rejection, } /// Records accepted for import, plus the ones that were skipped. @@ -48,6 +71,19 @@ pub struct AcceptedInstances { pub rejected: Vec, } +impl AcceptedInstances { + /// Instance IDs that are absent from `instances` only because their record + /// was unreadable, and whose existing data-plane state must therefore be + /// left in place rather than treated as a remote deletion. + pub fn unreadable(&self) -> HashSet<&str> { + self.rejected + .iter() + .filter(|rejected| rejected.rejection == Rejection::Unusable) + .map(|rejected| rejected.instance_id.as_str()) + .collect() + } +} + /// Validate a WireGuard public key as `wg` itself would accept it. pub fn validate_wg_public_key(public_key: &str) -> Result<()> { ensure!(!public_key.is_empty(), "public key is empty"); @@ -102,11 +138,13 @@ fn validate_instance(wg: &WgConfig, instance_id: &str, data: &InstanceData) -> R /// wins, ties broken by instance ID) so that every node reaches the same /// decision from the same KV contents — the local registration path resolves /// them the same way, by refusing the newcomer. -pub fn accept_instances( - wg: &WgConfig, - records: BTreeMap, -) -> AcceptedInstances { - let mut ordered: Vec<(String, InstanceData)> = records.into_iter().collect(); +pub fn accept_instances(wg: &WgConfig, loaded: LoadedInstances) -> AcceptedInstances { + let LoadedInstances { + decoded, + undecodable, + } = loaded; + + let mut ordered: Vec<(String, InstanceData)> = decoded.into_iter().collect(); ordered.sort_by(|(left_id, left), (right_id, right)| { left.reg_time .cmp(&right.reg_time) @@ -114,24 +152,40 @@ pub fn accept_instances( }); let mut instances = BTreeMap::new(); - let mut rejected = Vec::new(); + let mut rejected: Vec = undecodable + .into_iter() + .map(|instance_id| RejectedInstance { + instance_id, + reason: anyhow::anyhow!("record does not decode"), + rejection: Rejection::Unusable, + }) + .collect(); let mut claimed_ips: HashMap = HashMap::new(); let mut claimed_keys: HashSet = HashSet::new(); for (instance_id, data) in ordered { - let checked = validate_instance(wg, &instance_id, &data).and_then(|()| { - if let Some(owner) = claimed_ips.get(&data.ip) { - bail!("ip {} is already assigned to instance {owner}", data.ip); - } - if claimed_keys.contains(&data.public_key) { - bail!("public key is already registered to another instance"); - } - Ok(()) - }); - if let Err(reason) = checked { + let checked = validate_instance(wg, &instance_id, &data) + .map_err(|reason| (Rejection::Unusable, reason)) + .and_then(|()| { + if let Some(owner) = claimed_ips.get(&data.ip) { + return Err(( + Rejection::LostConflict, + anyhow::anyhow!("ip {} is already assigned to instance {owner}", data.ip), + )); + } + if claimed_keys.contains(&data.public_key) { + return Err(( + Rejection::LostConflict, + anyhow::anyhow!("public key is already registered to another instance"), + )); + } + Ok(()) + }); + if let Err((rejection, reason)) = checked { rejected.push(RejectedInstance { instance_id, reason, + rejection, }); continue; } @@ -150,6 +204,7 @@ pub fn accept_instances( mod tests { use super::*; use ipnet::Ipv4Net; + use std::collections::BTreeSet; fn wg_config() -> WgConfig { WgConfig { @@ -181,12 +236,18 @@ mod tests { } } + fn loaded(records: Vec<(&str, InstanceData)>) -> LoadedInstances { + LoadedInstances { + decoded: records + .into_iter() + .map(|(id, data)| (id.to_string(), data)) + .collect(), + undecodable: BTreeSet::new(), + } + } + fn accept(records: Vec<(&str, InstanceData)>) -> AcceptedInstances { - let records = records - .into_iter() - .map(|(id, data)| (id.to_string(), data)) - .collect(); - accept_instances(&wg_config(), records) + accept_instances(&wg_config(), loaded(records)) } #[test] @@ -293,16 +354,15 @@ mod tests { client_ip_range: pool.parse::().unwrap(), ..wg_config() }; - let records = [ - ("mine", instance(mine, &key(1), 100)), - ("peers", instance(peers, &key(2), 100)), - // Still refused: this gateway's own address. - ("steals-gateway-ip", instance(gateway_addr, &key(3), 100)), - ] - .into_iter() - .map(|(id, data)| (id.to_string(), data)) - .collect(); - let accepted = accept_instances(&wg, records); + let accepted = accept_instances( + &wg, + loaded(vec![ + ("mine", instance(mine, &key(1), 100)), + ("peers", instance(peers, &key(2), 100)), + // Still refused: this gateway's own address. + ("steals-gateway-ip", instance(gateway_addr, &key(3), 100)), + ]), + ); assert!(accepted.instances.contains_key("mine"), "{ip}"); assert!( accepted.instances.contains_key("peers"), @@ -315,4 +375,41 @@ mod tests { ); } } + + #[test] + fn a_conflict_loser_is_dropped_but_an_unusable_record_keeps_its_instance() { + // The two rejection kinds drive opposite decisions in the reload pass: + // a conflict loser must lose its routing to the winner, while an + // instance whose record we cannot read keeps what the data plane holds. + let accepted = accept(vec![ + ("loser", instance("10.0.0.20", &key(9), 300)), + ("winner", instance("10.0.0.20", &key(1), 100)), + ("malformed", instance("10.0.0.30", "not-a-key", 100)), + ]); + let kind = |id: &str| { + accepted + .rejected + .iter() + .find(|rejected| rejected.instance_id == id) + .map(|rejected| rejected.rejection) + }; + assert_eq!(kind("loser"), Some(Rejection::LostConflict)); + assert_eq!(kind("malformed"), Some(Rejection::Unusable)); + assert_eq!(accepted.unreadable(), HashSet::from(["malformed"])); + } + + #[test] + fn an_undecodable_record_is_reported_as_unreadable_not_as_absent() { + let accepted = accept_instances( + &wg_config(), + LoadedInstances { + decoded: [("good".to_string(), instance("10.0.0.20", &key(1), 100))] + .into_iter() + .collect(), + undecodable: ["corrupt".to_string()].into_iter().collect(), + }, + ); + assert!(accepted.instances.contains_key("good")); + assert_eq!(accepted.unreadable(), HashSet::from(["corrupt"])); + } } diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index dca06d86e..10e9fd4ac 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -36,7 +36,12 @@ pub use https_client::{AppIdValidator, HttpsClientConfig}; pub use sync_service::{fetch_peers_from_bootnode, WaveKvSyncService}; use tracing::{error, warn}; -use std::{collections::BTreeMap, net::Ipv4Addr, path::Path, time::Duration}; +use std::{ + collections::{BTreeMap, BTreeSet}, + net::Ipv4Addr, + path::Path, + time::Duration, +}; use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; @@ -90,6 +95,22 @@ pub struct InstanceData { pub admin_port_policy: Option, } +/// The `inst/` records currently in the KV store, split by readability. +/// +/// A key that is absent or tombstoned does not appear here at all — that is the +/// signal that the instance was deleted. A key whose bytes no longer decode +/// lands in `undecodable`, which is deliberately *not* the same signal: the +/// record still exists, we just cannot read it, and dropping the instance from +/// the data plane on that basis would turn one unreadable record into an +/// outage. +#[derive(Debug, Default)] +pub struct LoadedInstances { + /// Records that decoded successfully, keyed by instance ID. + pub decoded: BTreeMap, + /// Instance IDs whose stored bytes are present but no longer decode. + pub undecodable: BTreeSet, +} + /// Gateway node status (stored separately for independent updates) #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] #[serde(rename_all = "snake_case")] @@ -475,6 +496,10 @@ trait GetPutCodec { &self, prefix: &str, ) -> impl Iterator; + fn iter_decoded_strict serde::Deserialize<'de>>( + &self, + prefix: &str, + ) -> impl Iterator)>; } impl GetPutCodec for NodeState { @@ -547,6 +572,25 @@ impl GetPutCodec for NodeState { Some(value) }) } + + /// Like [`Self::iter_decoded`], but surfaces undecodable records instead of + /// skipping them. + /// + /// Tombstoned keys are still skipped — a deleted record and an unreadable + /// one call for opposite responses, and only this form lets the caller tell + /// them apart. + fn iter_decoded_strict serde::Deserialize<'de>>( + &self, + prefix: &str, + ) -> impl Iterator)> { + self.iter_by_prefix(prefix).filter_map(|(key, entry)| { + let value = entry.value.as_ref()?; + Some(( + key.to_string(), + decode(value).with_context(|| format!("corrupt record at KV key {key}")), + )) + }) + } } /// Sync store wrapping two WaveKV Nodes (persistent and ephemeral). @@ -648,16 +692,28 @@ impl KvStore { Ok(()) } - /// Load all instances from sync store (for initial sync on startup) - pub fn load_all_instances(&self) -> BTreeMap { - self.persistent + /// Load all instances from the sync store. + pub fn load_all_instances(&self) -> LoadedInstances { + let mut loaded = LoadedInstances::default(); + for (key, result) in self + .persistent .read() - .iter_decoded(keys::INST_PREFIX) - .filter_map(|(key, data)| { - let instance_id = keys::parse_inst_key(&key)?; - Some((instance_id.into(), data)) - }) - .collect() + .iter_decoded_strict::(keys::INST_PREFIX) + { + let Some(instance_id) = keys::parse_inst_key(&key) else { + continue; + }; + match result { + Ok(data) => { + loaded.decoded.insert(instance_id.into(), data); + } + Err(err) => { + error!("{err:#}"); + loaded.undecodable.insert(instance_id.into()); + } + } + } + loaded } // ==================== Node Sync ==================== @@ -1704,7 +1760,9 @@ mod corruption_tests { std::fs::write(data_dir.join("node_1.wal"), b"garbage").expect("failed to corrupt wal"); let kv = KvStore::new(1, vec![], &data_dir).expect("startup must survive a corrupt wal"); - assert!(kv.load_all_instances().is_empty()); + let loaded = kv.load_all_instances(); + assert!(loaded.decoded.is_empty()); + assert!(loaded.undecodable.is_empty()); let quarantined: Vec<_> = std::fs::read_dir(dir.path()) .expect("failed to read temp dir") .filter_map(|entry| entry.ok()) diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index 0a13b50ff..bfd8e77e8 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -40,7 +40,8 @@ use crate::{ config::{Config, TlsConfig}, kv::{ fetch_peers_from_bootnode, import, AppIdValidator, CertData, HttpsClientConfig, - InstanceData, KvStore, NodeData, NodeStatus, PortPolicy, WaveKvSyncService, + InstanceData, KvStore, LoadedInstances, NodeData, NodeStatus, PortPolicy, + WaveKvSyncService, }, models::{InstanceInfo, PortPolicyView, WgConf, WgPeer}, proxy::{create_acceptor_with_cert_resolver, AddressGroup, AddressInfo, AppAddressResolver}, @@ -165,8 +166,9 @@ impl ProxyInner { let instances = kv_store.load_all_instances(); let nodes = kv_store.load_all_nodes(); info!( - "Loaded state from WaveKV: {} instances, {} nodes", - instances.len(), + "Loaded state from WaveKV: {} instances ({} unreadable), {} nodes", + instances.decoded.len(), + instances.undecodable.len(), nodes.len() ); let state = build_state_from_kv_store(&config, instances); @@ -576,24 +578,22 @@ impl Proxy { /// /// A refused record makes its CVM invisible to this node, so it must never be /// a silent skip. -fn report_rejected_instances(rejected: Vec) { +fn report_rejected_instances(rejected: &[import::RejectedInstance]) { for import::RejectedInstance { instance_id, reason, + .. } in rejected { error!("ignoring KV instance record {instance_id}: {reason:#}"); } } -fn build_state_from_kv_store( - config: &Config, - instances: BTreeMap, -) -> ProxyStateMut { +fn build_state_from_kv_store(config: &Config, instances: LoadedInstances) -> ProxyStateMut { let mut state = ProxyStateMut::default(); let accepted = import::accept_instances(&config.wg, instances); - report_rejected_instances(accepted.rejected); + report_rejected_instances(&accepted.rejected); // Build instances for (instance_id, data) in accepted.instances { @@ -907,21 +907,29 @@ const LOCAL_REGISTRATION_GRACE: Duration = Duration::from_secs(60); fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()> { let accepted = import::accept_instances(&proxy.config.wg, store.load_all_instances()); - report_rejected_instances(accepted.rejected); + report_rejected_instances(&accepted.rejected); + // An unreadable record is not a deletion. Its instance keeps whatever the + // data plane already holds, so it must be exempt from the removal pass + // below; a record that lost an IP or key conflict is not exempt, because + // the winner owns that IP or key and the loser has to stop being routable. + let unreadable: HashSet = accepted + .unreadable() + .into_iter() + .map(str::to_owned) + .collect(); let instances = accepted.instances; let mut state = proxy.lock(); let mut wg_changed = false; // Instances deleted (or recycled) on another node must stop being routable // here too, rather than lingering until this node's own recycle timeout. - // Records that merely failed validation are left alone: the last known-good - // state of an instance is better than no state at all. let removed: Vec = state .state .instances .iter() .filter(|(id, info)| { !instances.contains_key(*id) + && !unreadable.contains(id.as_str()) && info.reg_time.elapsed().unwrap_or_default() > LOCAL_REGISTRATION_GRACE }) .map(|(id, _)| id.clone()) diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs index d265cb90c..b4f4ab4c7 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -448,6 +448,16 @@ async fn gateway_top_n_batch_007_cache_health_and_invalidation() { /// Write a record straight into the KV store, bypassing registration, the way /// a peer's sync round would. fn sync_from_peer(state: &TestState, instance_id: &str, ip: &str, public_key: &str) { + sync_from_peer_at(state, instance_id, ip, public_key, 1); +} + +fn sync_from_peer_at( + state: &TestState, + instance_id: &str, + ip: &str, + public_key: &str, + reg_time: u64, +) { state .kv_store .sync_instance( @@ -456,7 +466,7 @@ fn sync_from_peer(state: &TestState, instance_id: &str, ip: &str, public_key: &s app_id: "peer-app".to_string(), ip: ip.parse().unwrap(), public_key: public_key.to_string(), - reg_time: 1, + reg_time, port_policy: None, port_policy_hash: String::new(), admin_port_policy: None, @@ -566,3 +576,67 @@ async fn a_local_registration_survives_a_reload_that_cannot_see_it_yet() { reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); assert!(state.lock().state.instances.contains_key("fresh")); } + +#[tokio::test] +async fn a_record_that_stops_validating_keeps_the_instance_it_describes() { + let state = create_test_state().await; + sync_from_peer(&state, "peer-instance", "10.0.0.40", &test_pubkey("good")); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + assert!(state.lock().state.instances.contains_key("peer-instance")); + + // A peer overwrites the record with one this node refuses. "I cannot read + // your current state" is not "you were deleted": the last known-good state + // keeps the CVM reachable until a usable record or a real deletion arrives. + sync_from_peer(&state, "peer-instance", "10.0.0.40", "not-a-key"); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + + let proxy = state.lock(); + let instance = proxy + .state + .instances + .get("peer-instance") + .expect("a rejected record must not evict the instance"); + assert_eq!(instance.public_key, test_pubkey("good")); + assert!(proxy.generate_wg_config().unwrap().contains("10.0.0.40")); +} + +#[tokio::test] +async fn an_undecodable_record_keeps_the_instance_it_describes() { + let state = create_test_state().await; + sync_from_peer(&state, "peer-instance", "10.0.0.40", &test_pubkey("good")); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + + // A torn write, or a record written by a build whose schema this one cannot + // read. Folding that into "absent" would let a rolling upgrade evict every + // instance the newer nodes registered. + state + .kv_store + .persistent() + .write() + .put( + crate::kv::keys::inst("peer-instance"), + b"not-messagepack".to_vec(), + ) + .unwrap(); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + + assert!(state.lock().state.instances.contains_key("peer-instance")); +} + +#[tokio::test] +async fn an_instance_that_lost_an_ip_conflict_stops_being_routable() { + let state = create_test_state().await; + sync_from_peer_at(&state, "loser", "10.0.0.40", &test_pubkey("loser"), 300); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + assert!(state.lock().state.instances.contains_key("loser")); + + // An older registration claims the same IP. Unlike an unreadable record, + // this one is readable and says the address belongs to someone else, so + // the loser has to go: two peers on one IP is a config `wg` will not load. + sync_from_peer_at(&state, "winner", "10.0.0.40", &test_pubkey("winner"), 100); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + + let proxy = state.lock(); + assert!(!proxy.state.instances.contains_key("loser")); + assert!(proxy.state.instances.contains_key("winner")); +} From d1688097f25c01983c593a37dac53d3ba04c5b34 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 11 Aug 2026 00:27:34 -0700 Subject: [PATCH 08/13] fix(gateway): reject instance records dated into the future MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P0.4 put a drift horizon on `handshake/` and `last_seen/`, but left `reg_time` unchecked, and `reg_time` feeds the same kind of arithmetic. It is the registering node's clock at one instant, so the clock only has to be wrong once — during the window before chrony converges, or across a time jump — for a future timestamp to be written into the KV, and it does not heal when the clock does. From then on that record is permanent: the reload's "gone from KV" pass and `recycle()` both age instances with `elapsed().unwrap_or_default()`, which reads a future timestamp as zero age. The instance is immune to remote deletion and to local recycling at the same time, and stays routable on every node that loaded it until that process restarts, even after an operator deletes the record. Import now holds `reg_time` to the same `MAX_CLOCK_DRIFT_SECS` horizon as the observations, which closes the path at the boundary and leaves the two `elapsed().unwrap_or_default()` call sites nothing to mishandle. Drift inside the horizon is still accepted, since nodes are not perfectly synchronized; the batch samples the clock once so every record in it is judged against the same instant. Also check a public key's base64 length before decoding it. `wg` writes the padded form and accepts nothing else, so 44 characters is part of the format rather than something to discover from the decode. --- dstack/gateway/src/kv/import.rs | 81 ++++++++++++++++++++++-- dstack/gateway/src/kv/mod.rs | 12 +++- dstack/gateway/src/main_service/tests.rs | 21 ++++++ 3 files changed, 105 insertions(+), 9 deletions(-) diff --git a/dstack/gateway/src/kv/import.rs b/dstack/gateway/src/kv/import.rs index e21ef896f..6f898c8e1 100644 --- a/dstack/gateway/src/kv/import.rs +++ b/dstack/gateway/src/kv/import.rs @@ -30,12 +30,15 @@ use std::net::Ipv4Addr; use anyhow::{ensure, Result}; use base64::{engine::general_purpose::STANDARD, Engine}; -use super::{InstanceData, LoadedInstances}; +use super::{now_secs, InstanceData, LoadedInstances, MAX_CLOCK_DRIFT_SECS}; use crate::config::WgConfig; /// A WireGuard public key is 32 raw bytes, base64-encoded by `wg`. const WG_PUBLIC_KEY_BYTES: usize = 32; +/// Padded base64 of 32 bytes is always 44 characters. +const WG_PUBLIC_KEY_B64_LEN: usize = 44; + /// Upper bound on identifier fields carried in a KV record. Real values are /// hex-encoded hashes (40-64 chars); the bound keeps a corrupt record that /// still decodes from pushing an arbitrarily long string into the logs and the @@ -91,6 +94,13 @@ pub fn validate_wg_public_key(public_key: &str) -> Result<()> { !public_key.contains(|c: char| c.is_whitespace() || c.is_control()), "public key contains whitespace or control characters" ); + // `wg` writes the padded form and accepts nothing else, so the length is + // part of the format rather than a consequence of the decode below. + ensure!( + public_key.len() == WG_PUBLIC_KEY_B64_LEN, + "public key is {} characters, expected {WG_PUBLIC_KEY_B64_LEN}", + public_key.len() + ); let decoded = STANDARD .decode(public_key) .map_err(|err| anyhow::anyhow!("public key is not valid base64: {err}"))?; @@ -117,7 +127,15 @@ fn validate_id(field: &str, value: &str) -> Result<()> { } /// Per-record checks that do not depend on any other record. -fn validate_instance(wg: &WgConfig, instance_id: &str, data: &InstanceData) -> Result<()> { +/// +/// `now` is the local wall clock in seconds, taken once per batch so every +/// record in a batch is judged against the same instant. +fn validate_instance( + wg: &WgConfig, + now: u64, + instance_id: &str, + data: &InstanceData, +) -> Result<()> { validate_id("instance_id", instance_id)?; validate_id("app_id", &data.app_id)?; validate_wg_public_key(&data.public_key)?; @@ -129,6 +147,21 @@ fn validate_instance(wg: &WgConfig, instance_id: &str, data: &InstanceData) -> R "ip {} is outside the WireGuard network", data.ip ); + // `reg_time` is the registering node's clock at that one instant, so a + // clock only has to be wrong once — during the window before chrony + // converges, or across a time jump — for the future timestamp to be written + // into the KV permanently. It does not heal when the clock does: both the + // "gone from KV" pass and `recycle()` age instances with + // `elapsed().unwrap_or_default()`, which reads a future timestamp as zero + // age, so the instance is immune to remote deletion and to local recycling + // until the process restarts. Same horizon as the handshake observations, + // which are ignored for the same reason. + let horizon = now.saturating_add(MAX_CLOCK_DRIFT_SECS); + ensure!( + data.reg_time <= horizon, + "reg_time {} is more than {MAX_CLOCK_DRIFT_SECS}s ahead of local time ({now})", + data.reg_time + ); Ok(()) } @@ -139,6 +172,10 @@ fn validate_instance(wg: &WgConfig, instance_id: &str, data: &InstanceData) -> R /// decision from the same KV contents — the local registration path resolves /// them the same way, by refusing the newcomer. pub fn accept_instances(wg: &WgConfig, loaded: LoadedInstances) -> AcceptedInstances { + accept_instances_at(wg, loaded, now_secs()) +} + +fn accept_instances_at(wg: &WgConfig, loaded: LoadedInstances, now: u64) -> AcceptedInstances { let LoadedInstances { decoded, undecodable, @@ -164,7 +201,7 @@ pub fn accept_instances(wg: &WgConfig, loaded: LoadedInstances) -> AcceptedInsta let mut claimed_keys: HashSet = HashSet::new(); for (instance_id, data) in ordered { - let checked = validate_instance(wg, &instance_id, &data) + let checked = validate_instance(wg, now, &instance_id, &data) .map_err(|reason| (Rejection::Unusable, reason)) .and_then(|()| { if let Some(owner) = claimed_ips.get(&data.ip) { @@ -206,9 +243,13 @@ mod tests { use ipnet::Ipv4Net; use std::collections::BTreeSet; + /// Wall clock the tests validate against; every fixture `reg_time` below is + /// well under it unless the test is about the future-timestamp horizon. + const NOW: u64 = 1_700_000_000; + fn wg_config() -> WgConfig { WgConfig { - public_key: "gateway".to_string(), + public_key: key(7), private_key: "gateway".to_string(), listen_port: 51820, ip: "10.0.0.1/24".parse::().unwrap(), @@ -247,7 +288,7 @@ mod tests { } fn accept(records: Vec<(&str, InstanceData)>) -> AcceptedInstances { - accept_instances(&wg_config(), loaded(records)) + accept_instances_at(&wg_config(), loaded(records), NOW) } #[test] @@ -376,6 +417,25 @@ mod tests { } } + #[test] + fn rejects_registrations_dated_into_the_future() { + // A future reg_time reads as zero age in every `elapsed()` check, which + // makes the instance immune to both remote deletion and local recycling. + let accepted = accept(vec![ + ( + "future", + instance("10.0.0.20", &key(1), NOW + MAX_CLOCK_DRIFT_SECS + 1), + ), + ( + "skewed", + instance("10.0.0.21", &key(2), NOW + MAX_CLOCK_DRIFT_SECS), + ), + ]); + assert!(!accepted.instances.contains_key("future")); + // Drift inside the horizon is ordinary skew between nodes. + assert!(accepted.instances.contains_key("skewed")); + } + #[test] fn a_conflict_loser_is_dropped_but_an_unusable_record_keeps_its_instance() { // The two rejection kinds drive opposite decisions in the reload pass: @@ -400,7 +460,7 @@ mod tests { #[test] fn an_undecodable_record_is_reported_as_unreadable_not_as_absent() { - let accepted = accept_instances( + let accepted = accept_instances_at( &wg_config(), LoadedInstances { decoded: [("good".to_string(), instance("10.0.0.20", &key(1), 100))] @@ -408,8 +468,17 @@ mod tests { .collect(), undecodable: ["corrupt".to_string()].into_iter().collect(), }, + NOW, ); assert!(accepted.instances.contains_key("good")); assert_eq!(accepted.unreadable(), HashSet::from(["corrupt"])); } + + #[test] + fn rejects_keys_of_the_wrong_length_before_decoding_them() { + let oversized = "A".repeat(4096); + assert!(validate_wg_public_key(&oversized).is_err()); + // 32 bytes unpadded is 43 characters; `wg` writes the padded form. + assert!(validate_wg_public_key(key(1).trim_end_matches('=')).is_err()); + } } diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 10e9fd4ac..4312e22c1 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -625,9 +625,15 @@ impl KvStore { let persistent = match Node::new_with_persistence(my_node_id, peer_ids.clone(), data_dir) { Ok(node) => node, Err(err) => { - let quarantined = quarantine_data_dir(data_dir).context( - "failed to open the WaveKV data dir and failed to move it aside for recovery", - )?; + // Keep the original open error in the context: if moving the + // directory aside also fails, the reason the open failed is the + // more useful half of the diagnosis and is otherwise lost. + let quarantined = quarantine_data_dir(data_dir).with_context(|| { + format!( + "failed to open the WaveKV data dir ({err:#}) and failed to \ + move it aside for recovery" + ) + })?; error!( "WaveKV data dir {} is unreadable ({err:#}); moved it to {} and started empty — \ state will be re-fetched from peers", diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs index b4f4ab4c7..a1528d5cb 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -475,6 +475,13 @@ fn sync_from_peer_at( .unwrap(); } +fn now_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() +} + #[tokio::test] async fn a_poisoned_peer_record_costs_only_its_own_instance() { let state = create_test_state().await; @@ -640,3 +647,17 @@ async fn an_instance_that_lost_an_ip_conflict_stops_being_routable() { assert!(!proxy.state.instances.contains_key("loser")); assert!(proxy.state.instances.contains_key("winner")); } + +#[tokio::test] +async fn a_future_dated_registration_cannot_park_itself_in_the_data_plane() { + let state = create_test_state().await; + // Both the removal pass and `recycle()` age instances with + // `elapsed().unwrap_or_default()`, which reads a future `reg_time` as zero + // age — an instance dated forward would survive deletion and recycling + // alike until this process restarts. + let far_future = now_secs() + 365 * 24 * 3600; + sync_from_peer_at(&state, "zombie", "10.0.0.40", &test_pubkey("z"), far_future); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + + assert!(!state.lock().state.instances.contains_key("zombie")); +} From 88369cc6eac706fe61b8859b1f2b924b2a2c47b5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 11 Aug 2026 00:27:59 -0700 Subject: [PATCH 09/13] fix(gateway): let an operator replace a corrupt certbot config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `start_certbot_task` refuses to run against an unreadable `global/certbot_config` and says to "wait for an operator to repair the record instead" — but there was no way to repair it. The key is a singleton with no delete RPC, so the only write path is the read-modify-write inside SetCertbotConfig, whose first statement is the fail-closed `get_certbot_config()`. One bad record therefore stopped renewal permanently, and since `do_rotate_acme_credentials` reads the same key, it took RotateAcmeCredentials down with it: unlike `global/acme_credentials`, this one had no way back. SetCertbotConfig is a partial update — a field the operator leaves unset keeps its stored value — so simply reading through to the defaults would not do either. `acme_url` defaults to empty, meaning Let's Encrypt production, so an operator who hit a corrupt record and then tuned `renew_interval` would silently move issuance off their staging or private ACME server and start burning real rate limits. That is the very switch the fail-closed reader exists to prevent. The merge now happens in `merge_certbot_config`, which keeps the stored values when the record is readable and, when it is not, requires the request to state every field before it will replace it. Nothing is ever inherited from a record we cannot read, so the repair path exists without any field being guessed. The error names the fields to resend. --- dstack/gateway/src/admin_service.rs | 138 ++++++++++++++++++++++++---- dstack/gateway/src/kv/mod.rs | 20 ++++ 2 files changed, 139 insertions(+), 19 deletions(-) diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index 11313356a..625af036e 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -5,7 +5,7 @@ use std::sync::atomic::Ordering; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use anyhow::{bail, Context, Result}; +use anyhow::{bail, ensure, Context, Result}; use dstack_gateway_rpc::{ admin_server::{AdminRpc, AdminServer}, CertAttestationInfo, CertbotConfigResponse, ClearInstancePortPolicyRequest, @@ -24,11 +24,14 @@ use dstack_gateway_rpc::{ WaveKvStatusResponse, ZtDomainCertStatus, ZtDomainConfig as ProtoZtDomainConfig, ZtDomainInfo, }; use ra_rpc::{CallContext, RpcCall}; -use tracing::info; +use tracing::{info, warn}; use wavekv::node::NodeStatus as WaveKvNodeStatus; use crate::{ - kv::{DnsCredential, DnsProvider, NodeStatus, PortFlags, PortPolicy, ZtDomainConfig}, + kv::{ + DnsCredential, DnsProvider, GlobalCertbotConfig, NodeStatus, PortFlags, PortPolicy, + ZtDomainConfig, + }, main_service::Proxy, models::PortPolicyView, proxy::{stats::accel_status, NUM_CONNECTIONS}, @@ -624,22 +627,7 @@ impl AdminRpc for AdminRpcHandler { async fn set_certbot_config(self, request: SetCertbotConfigRequest) -> Result<()> { let kv_store = self.state.kv_store(); - let mut config = kv_store.get_certbot_config()?; - - // Update only the fields that are specified - if let Some(secs) = request.renew_interval_secs { - config.renew_interval = Duration::from_secs(secs); - } - if let Some(secs) = request.renew_before_expiration_secs { - config.renew_before_expiration = Duration::from_secs(secs); - } - if let Some(secs) = request.renew_timeout_secs { - config.renew_timeout = Duration::from_secs(secs); - } - if let Some(url) = request.acme_url { - config.acme_url = url; - } - + let config = merge_certbot_config(kv_store.get_certbot_config(), request)?; kv_store.set_certbot_config(&config)?; info!( "Updated certbot config: renew_interval={:?}, renew_before_expiration={:?}, renew_timeout={:?}, acme_url={:?}", @@ -903,6 +891,118 @@ fn zt_domain_to_proto( } } +/// Apply a partial certbot-config update to the stored record. +/// +/// SetCertbotConfig is a merge: a field the operator leaves unset keeps its +/// stored value. That needs a readable base, and `global/certbot_config` is a +/// singleton with no delete RPC — so if an unreadable record simply failed the +/// call, the corruption would be permanent, and since `do_rotate_acme_credentials` +/// reads the same key it would keep RotateAcmeCredentials blocked along with it. +/// +/// Merging into the defaults instead is not the answer either: `acme_url` +/// defaults to empty, which means Let's Encrypt production. An operator who hit +/// a corrupt record and then tuned `renew_interval` would silently move issuance +/// off their staging or private ACME server and start burning real rate limits — +/// exactly the switch the fail-closed reader exists to prevent. +/// +/// So an unreadable record is repairable, but only by a request that states +/// every field. Nothing is ever inherited from a record we cannot read. +fn merge_certbot_config( + stored: Result, + request: SetCertbotConfigRequest, +) -> Result { + let mut config = match stored { + Ok(config) => config, + Err(err) => { + ensure!( + request.renew_interval_secs.is_some() + && request.renew_before_expiration_secs.is_some() + && request.renew_timeout_secs.is_some() + && request.acme_url.is_some(), + "the stored certbot config is unreadable ({err:#}), so it can only be \ + replaced as a whole: resend with renew_interval_secs, \ + renew_before_expiration_secs, renew_timeout_secs and acme_url all set" + ); + warn!("certbot config is unreadable ({err:#}); replacing it wholesale"); + GlobalCertbotConfig::default() + } + }; + + // Update only the fields that are specified + if let Some(secs) = request.renew_interval_secs { + config.renew_interval = Duration::from_secs(secs); + } + if let Some(secs) = request.renew_before_expiration_secs { + config.renew_before_expiration = Duration::from_secs(secs); + } + if let Some(secs) = request.renew_timeout_secs { + config.renew_timeout = Duration::from_secs(secs); + } + if let Some(url) = request.acme_url { + config.acme_url = url; + } + Ok(config) +} + +#[cfg(test)] +mod certbot_config_tests { + use super::*; + + fn stored() -> GlobalCertbotConfig { + GlobalCertbotConfig { + renew_interval: Duration::from_secs(3600), + acme_url: "https://acme-staging.example/directory".to_string(), + ..Default::default() + } + } + + #[test] + fn a_partial_update_keeps_the_fields_it_does_not_mention() { + let merged = merge_certbot_config( + Ok(stored()), + SetCertbotConfigRequest { + renew_timeout_secs: Some(60), + ..Default::default() + }, + ) + .expect("a readable record merges"); + assert_eq!(merged.renew_timeout, Duration::from_secs(60)); + assert_eq!(merged.acme_url, stored().acme_url); + } + + #[test] + fn a_partial_update_cannot_repair_an_unreadable_record() { + // Falling back to the defaults here would reset `acme_url` to empty, + // silently moving issuance to Let's Encrypt production. + let err = merge_certbot_config( + Err(anyhow::anyhow!("corrupt record")), + SetCertbotConfigRequest { + renew_interval_secs: Some(60), + ..Default::default() + }, + ) + .expect_err("a partial update must not inherit from an unreadable record"); + assert!(err.to_string().contains("acme_url"), "{err:#}"); + } + + #[test] + fn a_complete_request_replaces_an_unreadable_record() { + // The only repair path: no field is inherited, so nothing is guessed. + let merged = merge_certbot_config( + Err(anyhow::anyhow!("corrupt record")), + SetCertbotConfigRequest { + renew_interval_secs: Some(60), + renew_before_expiration_secs: Some(86400), + renew_timeout_secs: Some(30), + acme_url: Some("https://acme-staging.example/directory".to_string()), + }, + ) + .expect("a complete request replaces the record"); + assert_eq!(merged.renew_interval, Duration::from_secs(60)); + assert_eq!(merged.acme_url, "https://acme-staging.example/directory"); + } +} + #[cfg(test)] mod zt_domain_tests { use super::validate_zt_domain; diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 4312e22c1..4140b7d3d 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -1691,6 +1691,26 @@ mod corruption_tests { assert!(kv.get_certbot_config().is_err()); } + #[test] + fn a_corrupt_certbot_config_can_still_be_replaced_by_an_operator() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let kv = test_kv(dir.path()); + put_raw(&kv, keys::GLOBAL_CERTBOT_CONFIG, b"not-messagepack"); + + // The key is a singleton with no delete RPC, so overwriting it is the + // only repair path there is; the write must not inherit the read's + // fail-closed behaviour. (Which fields an operator has to supply to be + // allowed to overwrite is decided one layer up, in `admin_service`.) + kv.set_certbot_config(&GlobalCertbotConfig { + acme_url: "https://acme-staging.example/directory".to_string(), + ..Default::default() + }) + .expect("save should succeed"); + + let repaired = kv.get_certbot_config().expect("record should be readable"); + assert_eq!(repaired.acme_url, "https://acme-staging.example/directory"); + } + #[test] fn corrupt_global_records_fail_closed() { let dir = tempfile::tempdir().expect("failed to create temp dir"); From 58ce20fa9943a1b2a0ac425f270bc6be049a5721 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 11 Aug 2026 03:10:49 -0700 Subject: [PATCH 10/13] fix(gateway): fail the boot on a storage fault instead of quarantining MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Quarantining an unreadable data dir is right when the *contents* cannot be read: a torn WAL tail is the normal artifact of a crash, every record is replicated, and a gateway that will not start serves nothing. It is wrong when the *storage* is at fault. A full disk, an exhausted fd table, a data volume that has not finished mounting — these say nothing about the contents, so moving the directory aside discards intact state. Worse, the condition survives a restart, so each boot attempt quarantines again and buries the real data under a pile of `.corrupt.*` directories. For a single-node deployment holding the only copy of the ACME account and DNS credentials, that is the difference between a restart and a rebuild. wavekv reports both classes through `anyhow`, so they are told apart by what is in the error chain. Unreadable content arrives as a decode failure, a checksum or header `bail!`, or a read that ran off the end of a truncated file — the last of which is an `io::Error`, but only ever `UnexpectedEof` or `InvalidData`. Any other `io::Error` is the storage layer, and now fails the boot with that error as the cause, which is also what puts the actual fault in front of the operator instead of a misleading "started empty" line. --- dstack/gateway/src/kv/mod.rs | 70 ++++++++++++++++++++++++++++++++++-- 1 file changed, 68 insertions(+), 2 deletions(-) diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 4140b7d3d..27a5dbbf8 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -607,15 +607,47 @@ pub struct KvStore { my_node_id: NodeId, } +/// Whether opening the persistent store failed because the storage is +/// unavailable, rather than because the stored bytes are unreadable. +/// +/// wavekv reports both through `anyhow`, so they have to be told apart by what +/// is in the error chain. Unreadable content arrives as a decode failure, a +/// checksum or header `bail!`, or a read that ran off the end of a truncated +/// file — the last of which is an `io::Error`, but only ever `UnexpectedEof` or +/// `InvalidData`. Every other `io::Error` is the storage layer talking: no +/// space left, permission denied, too many open files, the data volume not +/// mounted yet. +fn is_storage_failure(err: &anyhow::Error) -> bool { + err.chain() + .filter_map(|cause| cause.downcast_ref::()) + .any(|io| { + !matches!( + io.kind(), + std::io::ErrorKind::UnexpectedEof | std::io::ErrorKind::InvalidData + ) + }) +} + impl KvStore { /// Create a new sync store. /// - /// If the on-disk WAL/snapshot cannot be read, the data directory is moved - /// aside and the store starts empty rather than refusing to boot: the + /// If the on-disk WAL/snapshot cannot be *read*, the data directory is + /// moved aside and the store starts empty rather than refusing to boot: the /// persistent state is replicated on every peer, a torn WAL tail is the /// normal artifact of a crash, and a gateway that cannot start serves no /// traffic at all. Nothing is deleted — the unreadable directory is kept /// under `.corrupt.` for inspection. + /// + /// A failure of the *storage* is a different matter and fails the boot. A + /// full disk, an exhausted fd table or a volume that has not finished + /// mounting all say nothing about the contents, so moving the directory + /// aside would discard intact state — and, because the condition persists + /// across restarts, would do it again on every attempt, burying the real + /// data under a pile of `.corrupt.*` directories. Failing here instead + /// leaves the state alone and puts the actual cause in front of the + /// operator, which for a single-node deployment holding the only copy of + /// the ACME account and DNS credentials is the difference between a restart + /// and a rebuild. pub fn new( my_node_id: NodeId, peer_ids: Vec, @@ -624,6 +656,15 @@ impl KvStore { let data_dir = data_dir.as_ref(); let persistent = match Node::new_with_persistence(my_node_id, peer_ids.clone(), data_dir) { Ok(node) => node, + Err(err) if is_storage_failure(&err) => { + return Err(err).with_context(|| { + format!( + "cannot open the WaveKV data dir {}; refusing to start rather than \ + quarantine a directory whose contents are most likely intact", + data_dir.display() + ) + }); + } Err(err) => { // Keep the original open error in the context: if moving the // directory aside also fails, the reason the open failed is the @@ -1760,6 +1801,31 @@ mod corruption_tests { assert_eq!(kv.get_instance_handshakes("cvm").len(), 2); } + #[test] + fn a_storage_failure_fails_the_boot_instead_of_quarantining_intact_state() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + // Stand in for the storage being unusable — a full disk, an exhausted + // fd table, a volume that has not finished mounting. None of these say + // anything about the contents, and the condition survives a restart, so + // quarantining here would discard intact state once per boot attempt. + let data_dir = dir.path().join("kv"); + std::fs::write(&data_dir, b"not a directory").expect("failed to create blocker"); + + let Err(err) = KvStore::new(1, vec![], &data_dir) else { + panic!("startup must fail when the storage is unusable"); + }; + assert!( + format!("{err:#}").contains("refusing to start"), + "wrong failure: {err:#}" + ); + let quarantined = std::fs::read_dir(dir.path()) + .expect("failed to read temp dir") + .filter_map(|entry| entry.ok()) + .filter(|entry| entry.file_name().to_string_lossy().contains(".corrupt.")) + .count(); + assert_eq!(quarantined, 0, "quarantined a directory it could not read"); + } + #[test] fn an_unreadable_data_dir_is_quarantined_instead_of_blocking_startup() { let dir = tempfile::tempdir().expect("failed to create temp dir"); From 67f8509e919e0591345a4a467dcf29c5067f7d31 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 11 Aug 2026 03:10:56 -0700 Subject: [PATCH 11/13] fix(gateway): report a stuck KV record once, not on every reload MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A record is refused for what it contains, so nothing about the next reload makes it acceptable — it stays refused until someone rewrites it. Logging the whole refused set every round turns one stuck record into an unbounded stream of identical `error!` lines, emitted at whatever rate peer syncs happen to wake the watch task, which buries the first occurrence exactly when it matters. The reload now keeps the reason last logged per instance and reports a record when it starts being refused or its reason changes, and again at `info!` when it becomes usable. The log carries transitions instead of a level, and a refusal is still never silent. --- dstack/gateway/src/main_service.rs | 40 +++++++++++++++++++++++- dstack/gateway/src/main_service/tests.rs | 24 ++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index bfd8e77e8..e783da440 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -114,6 +114,9 @@ pub(crate) struct ProxyState { kv_store: Arc, handshake_cache: Arc, admin_shutdown: Option, + /// Reason last logged for each KV instance record this node refuses, so a + /// record that stays bad is reported once rather than on every reload. + reported_rejections: BTreeMap, } /// Options for creating a Proxy instance @@ -248,6 +251,7 @@ impl ProxyInner { kv_store: kv_store.clone(), handshake_cache: handshake_cache.clone(), admin_shutdown: None, + reported_rejections: BTreeMap::new(), }); let auth_client = AuthClient::new(config.auth.clone()); // Bootstrap WaveKV first if sync is enabled, so certbot can load certs from peers @@ -589,6 +593,40 @@ fn report_rejected_instances(rejected: &[import::RejectedInstance]) { } } +/// Report refused records, but only what has changed since the last reload. +/// +/// A record is refused because of what it contains, so nothing about the next +/// reload will make it acceptable — it stays refused until someone rewrites it. +/// Logging the whole set every round turns one stuck record into an unbounded +/// stream of identical `error!` lines, at whatever rate peer syncs happen to +/// wake the watch task, which buries the very first occurrence. Report a record +/// when it starts being refused or its reason changes, and again when it +/// recovers, so the log carries transitions instead of a level. +fn report_new_rejections( + reported: &mut BTreeMap, + rejected: &[import::RejectedInstance], +) { + let mut current = BTreeMap::new(); + for import::RejectedInstance { + instance_id, + reason, + .. + } in rejected + { + let reason = format!("{reason:#}"); + if reported.get(instance_id) != Some(&reason) { + error!("ignoring KV instance record {instance_id}: {reason}"); + } + current.insert(instance_id.clone(), reason); + } + for instance_id in reported.keys() { + if !current.contains_key(instance_id) { + info!("KV instance record {instance_id} is usable again"); + } + } + *reported = current; +} + fn build_state_from_kv_store(config: &Config, instances: LoadedInstances) -> ProxyStateMut { let mut state = ProxyStateMut::default(); @@ -907,7 +945,6 @@ const LOCAL_REGISTRATION_GRACE: Duration = Duration::from_secs(60); fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()> { let accepted = import::accept_instances(&proxy.config.wg, store.load_all_instances()); - report_rejected_instances(&accepted.rejected); // An unreadable record is not a deletion. Its instance keeps whatever the // data plane already holds, so it must be exempt from the removal pass // below; a record that lost an IP or key conflict is not exempt, because @@ -919,6 +956,7 @@ fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()> .collect(); let instances = accepted.instances; let mut state = proxy.lock(); + report_new_rejections(&mut state.reported_rejections, &accepted.rejected); let mut wg_changed = false; // Instances deleted (or recycled) on another node must stop being routable diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs index a1528d5cb..896b5d7ca 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -661,3 +661,27 @@ async fn a_future_dated_registration_cannot_park_itself_in_the_data_plane() { assert!(!state.lock().state.instances.contains_key("zombie")); } + +#[tokio::test] +async fn a_stuck_bad_record_is_reported_once_and_again_when_it_recovers() { + let state = create_test_state().await; + sync_from_peer(&state, "peer-instance", "10.0.0.40", "not-a-key"); + + // A record is refused for what it contains, so it stays refused until + // someone rewrites it. `reported_rejections` is what keeps the reload from + // re-emitting the same `error!` on every round: a reason already in the map + // is not logged again, so a second identical reload must leave it untouched. + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + let first = state.lock().reported_rejections.clone(); + assert_eq!(first.len(), 1); + assert!(first["peer-instance"].contains("public key")); + + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + assert_eq!(state.lock().reported_rejections, first); + + // Recovery is a transition too, and clearing the entry is what lets a + // later relapse be reported instead of silently swallowed. + sync_from_peer(&state, "peer-instance", "10.0.0.40", &test_pubkey("good")); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + assert!(state.lock().reported_rejections.is_empty()); +} From 98e1bef43bcd06067326abe22b8c62fd1f6f7ad0 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 11 Aug 2026 03:31:12 -0700 Subject: [PATCH 12/13] refactor(gateway): keep one epoch-seconds helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `now_secs()` was already copied into `admin_service` and `distributed_certbot`, and the future-dated-observation fix earlier in this series added a third to `kv`. Three copies of four lines is not itself a problem; three copies that quietly differ would be, and they already had started to — two saturate with `unwrap_or_default()`, the new one with `unwrap_or(0)`. They agree today by luck. `main_service` also held `encode_ts`/`decode_ts`, the same conversion for an arbitrary `SystemTime`, so `now_secs()` is `encode_ts(now)` and belongs beside them rather than in whichever module needed it next. All three now live in `crate::time`, and the call sites that had the expression inlined — the peer last_seen write, the two cert-lock acquisitions, the proxy cert load, the node last_seen refresh, the debug service's reg_time — use them. No behaviour changes: every one of those saturated the same way already. Left alone: the call sites that write `duration_since(UNIX_EPOCH)?` and propagate. A clock behind the epoch is a real fault, and whether to report it or carry on is the caller's decision, not something to bury in a shared helper. --- dstack/gateway/src/admin_service.rs | 8 +---- dstack/gateway/src/debug_service.rs | 6 +--- dstack/gateway/src/distributed_certbot.rs | 10 ++---- dstack/gateway/src/kv/import.rs | 3 +- dstack/gateway/src/kv/mod.rs | 24 +++------------ dstack/gateway/src/main.rs | 1 + dstack/gateway/src/main_service.rs | 21 ++----------- dstack/gateway/src/main_service/tests.rs | 7 ----- dstack/gateway/src/time.rs | 37 +++++++++++++++++++++++ 9 files changed, 52 insertions(+), 65 deletions(-) create mode 100644 dstack/gateway/src/time.rs diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index 625af036e..bd8033119 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -35,6 +35,7 @@ use crate::{ main_service::Proxy, models::PortPolicyView, proxy::{stats::accel_status, NUM_CONNECTIONS}, + time::now_secs, }; pub struct AdminRpcHandler { @@ -751,13 +752,6 @@ impl RpcCall for AdminRpcHandler { // ==================== Helper Functions ==================== -fn now_secs() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs() -} - fn generate_cred_id() -> String { use std::time::SystemTime; let ts = SystemTime::now() diff --git a/dstack/gateway/src/debug_service.rs b/dstack/gateway/src/debug_service.rs index 703ecb0d8..761d9ad19 100644 --- a/dstack/gateway/src/debug_service.rs +++ b/dstack/gateway/src/debug_service.rs @@ -118,11 +118,7 @@ impl DebugRpc for DebugRpcHandler { .instances .values() .map(|inst| { - let reg_time = inst - .reg_time - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0); + let reg_time = crate::time::encode_ts(inst.reg_time); ProxyStateInstance { instance_id: inst.id.clone(), app_id: inst.app_id.clone(), diff --git a/dstack/gateway/src/distributed_certbot.rs b/dstack/gateway/src/distributed_certbot.rs index 715dcc3bf..a651b36de 100644 --- a/dstack/gateway/src/distributed_certbot.rs +++ b/dstack/gateway/src/distributed_certbot.rs @@ -8,7 +8,7 @@ //! with dynamic DNS credential configuration and attestation storage. use std::sync::Arc; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use std::time::Duration; use anyhow::{bail, Context, Result}; use certbot::{AcmeClient, Dns01Client}; @@ -23,6 +23,7 @@ use crate::kv::{ AcmeAttestation, CertAttestation, CertCredentials, CertData, DnsCredential, DnsProvider, KvStore, ZtDomainConfig, }; +use crate::time::now_secs; /// Lock timeout for certificate renewal (10 minutes) const RENEW_LOCK_TIMEOUT_SECS: u64 = 600; @@ -727,13 +728,6 @@ fn dns_credential_for(kv_store: &KvStore, config: &ZtDomainConfig) -> Result u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs() -} - fn get_cert_expiry(cert_pem: &str) -> Option { use x509_parser::prelude::*; let pem = Pem::iter_from_buffer(cert_pem.as_bytes()).next()?.ok()?; diff --git a/dstack/gateway/src/kv/import.rs b/dstack/gateway/src/kv/import.rs index 6f898c8e1..28ea9f3f5 100644 --- a/dstack/gateway/src/kv/import.rs +++ b/dstack/gateway/src/kv/import.rs @@ -30,8 +30,9 @@ use std::net::Ipv4Addr; use anyhow::{ensure, Result}; use base64::{engine::general_purpose::STANDARD, Engine}; -use super::{now_secs, InstanceData, LoadedInstances, MAX_CLOCK_DRIFT_SECS}; +use super::{InstanceData, LoadedInstances, MAX_CLOCK_DRIFT_SECS}; use crate::config::WgConfig; +use crate::time::now_secs; /// A WireGuard public key is 32 raw bytes, base64-encoded by `wg`. const WG_PUBLIC_KEY_BYTES: usize = 32; diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 27a5dbbf8..1cc9a396a 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -44,6 +44,8 @@ use std::{ }; use anyhow::{Context, Result}; + +use crate::time::now_secs; use serde::{Deserialize, Serialize}; use tokio::sync::watch; use wavekv::{node::NodeState, types::NodeId, Node}; @@ -440,13 +442,6 @@ pub fn gunzip_bounded(data: &[u8], limit: usize) -> Result> { /// the recycle timeout. pub const MAX_CLOCK_DRIFT_SECS: u64 = 300; -fn now_secs() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0) -} - /// Drop observations timestamped beyond [`MAX_CLOCK_DRIFT_SECS`] into the /// future, logging once per call with the number dropped. fn drop_future_observations( @@ -961,10 +956,7 @@ impl KvStore { } pub fn update_peer_last_seen(&self, peer_id: NodeId) { - let ts = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0); + let ts = now_secs(); let key = keys::last_seen_node(peer_id, self.my_node_id); if let Err(e) = self.ephemeral.write().put_encoded(key, &ts) { warn!("failed to update peer {peer_id} last_seen: {e}"); @@ -1250,10 +1242,7 @@ impl KvStore { /// Try to acquire certificate renew lock /// Returns true if lock acquired, false if already locked by another node pub fn try_acquire_cert_lock(&self, domain: &str, lock_timeout_secs: u64) -> bool { - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(); + let now = now_secs(); if let Some(existing) = self.get_cert_lock(domain) { // Check if lock is still valid (not expired) @@ -1291,10 +1280,7 @@ impl KvStore { /// replication latency; it is not mutual exclusion. A crashed holder is /// covered by the timeout. pub fn try_acquire_rotation_lock(&self, lock_timeout_secs: u64) -> Option { - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(); + let now = now_secs(); if let Some(existing) = self.get_rotation_lock() { // Check if lock is still valid (not expired) diff --git a/dstack/gateway/src/main.rs b/dstack/gateway/src/main.rs index e7b2eb1a2..0507dea80 100644 --- a/dstack/gateway/src/main.rs +++ b/dstack/gateway/src/main.rs @@ -32,6 +32,7 @@ mod main_service; mod models; mod pp; mod proxy; +mod time; mod web_routes; #[global_allocator] diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index e783da440..5c827675b 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -45,6 +45,7 @@ use crate::{ }, models::{InstanceInfo, PortPolicyView, WgConf, WgPeer}, proxy::{create_acceptor_with_cert_resolver, AddressGroup, AddressInfo, AppAddressResolver}, + time::{decode_ts, encode_ts, now_secs}, }; mod auth_client; @@ -302,10 +303,7 @@ impl ProxyInner { })?; let key_pem = std::fs::read_to_string(cert_key) .with_context(|| format!("failed to read proxy cert_key {}", cert_key.display()))?; - let now = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs(); + let now = now_secs(); let cert_data = CertData { cert_pem, key_pem, @@ -1527,10 +1525,7 @@ impl ProxyState { } // Update this node's last_seen in KvStore - let now = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0); + let now = now_secs(); if let Err(err) = self .kv_store .sync_node_last_seen(self.config.sync.node_id, now) @@ -1596,16 +1591,6 @@ impl ProxyState { } } -fn decode_ts(ts: u64) -> SystemTime { - UNIX_EPOCH - .checked_add(Duration::from_secs(ts)) - .unwrap_or(UNIX_EPOCH) -} - -pub(crate) fn encode_ts(ts: SystemTime) -> u64 { - ts.duration_since(UNIX_EPOCH).unwrap_or_default().as_secs() -} - pub struct RpcHandler { remote_app_id: Option>, remote_app_info: Option, diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs index 896b5d7ca..75096d4d9 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -475,13 +475,6 @@ fn sync_from_peer_at( .unwrap(); } -fn now_secs() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap() - .as_secs() -} - #[tokio::test] async fn a_poisoned_peer_record_costs_only_its_own_instance() { let state = create_test_state().await; diff --git a/dstack/gateway/src/time.rs b/dstack/gateway/src/time.rs new file mode 100644 index 000000000..71e29a376 --- /dev/null +++ b/dstack/gateway/src/time.rs @@ -0,0 +1,37 @@ +// SPDX-FileCopyrightText: © 2024-2025 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Epoch-seconds conversions shared across the gateway. +//! +//! Timestamps cross the KV boundary as `u64` seconds since the Unix epoch and +//! are held in memory as `SystemTime`, so both directions are needed in several +//! modules. Keeping one pair of them means the saturating behaviour — a time +//! before the epoch reads as 0, a count of seconds `SystemTime` cannot +//! represent clamps to the epoch — is decided once instead of at each call +//! site. +//! +//! Call sites that would rather fail than saturate keep their own +//! `duration_since(UNIX_EPOCH)?`: a clock behind the epoch is a real fault, and +//! whether to report it or carry on is the caller's decision, not this +//! module's. + +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +/// Seconds since the Unix epoch for `ts`, or 0 if `ts` predates the epoch. +pub(crate) fn encode_ts(ts: SystemTime) -> u64 { + ts.duration_since(UNIX_EPOCH).unwrap_or_default().as_secs() +} + +/// The instant `ts` seconds after the Unix epoch, clamped to the epoch when +/// that instant is not representable. +pub(crate) fn decode_ts(ts: u64) -> SystemTime { + UNIX_EPOCH + .checked_add(Duration::from_secs(ts)) + .unwrap_or(UNIX_EPOCH) +} + +/// The local wall clock, as seconds since the Unix epoch. +pub(crate) fn now_secs() -> u64 { + encode_ts(SystemTime::now()) +} From 76fa87168bd33f436062563363d25ce65db92b44 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 11 Aug 2026 23:43:15 -0700 Subject: [PATCH 13/13] fix(gateway): reject self-referential WireGuard peers --- dstack/gateway/src/kv/import.rs | 19 +++++++++++++++++++ dstack/gateway/src/main_service.rs | 7 +++++++ 2 files changed, 26 insertions(+) diff --git a/dstack/gateway/src/kv/import.rs b/dstack/gateway/src/kv/import.rs index 28ea9f3f5..ff96e3b54 100644 --- a/dstack/gateway/src/kv/import.rs +++ b/dstack/gateway/src/kv/import.rs @@ -140,6 +140,10 @@ fn validate_instance( validate_id("instance_id", instance_id)?; validate_id("app_id", &data.app_id)?; validate_wg_public_key(&data.public_key)?; + ensure!( + data.public_key != wg.public_key, + "public key belongs to this gateway" + ); // The routable network, not this node's allocation share: in a cluster // every node carries peers for the CVMs registered on the other nodes, and // those hold addresses from the other nodes' shares by design. @@ -339,6 +343,21 @@ mod tests { } } + #[test] + fn rejects_the_gateways_own_public_key() { + let wg = wg_config(); + let accepted = accept_instances_at( + &wg, + loaded(vec![( + "self-peer", + instance("10.0.0.20", &wg.public_key, 100), + )]), + NOW, + ); + assert!(accepted.instances.is_empty()); + assert_eq!(accepted.rejected.len(), 1); + } + #[test] fn duplicate_ip_and_key_claims_resolve_to_the_older_registration() { let accepted = accept(vec![ diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index 5c827675b..0a89fa3b4 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -1260,6 +1260,13 @@ impl ProxyState { error!("excluding instance {} from wg config: {err:#}", info.id); continue; } + if info.public_key == self.config.wg.public_key { + error!( + "excluding instance {} from wg config: public key belongs to this gateway", + info.id + ); + continue; + } if !self.config.wg.is_routable_client_ip(info.ip) { error!( "excluding instance {} from wg config: ip {} is outside the wg network",