diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index c1e988d19..bd8033119 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,14 +24,18 @@ 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}, + time::now_secs, }; pub struct AdminRpcHandler { @@ -310,7 +314,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 +327,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 +387,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 +418,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 +442,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 +462,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 +617,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,22 +628,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={:?}", @@ -759,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() @@ -851,7 +837,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")?; } @@ -899,6 +885,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/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/debug_service.rs b/dstack/gateway/src/debug_service.rs index b00176a14..761d9ad19 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, @@ -117,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 a7dc07a95..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; @@ -100,7 +101,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 +207,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 +342,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 +427,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 +481,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 +529,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,22 +719,15 @@ 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") } } -fn now_secs() -> 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 new file mode 100644 index 000000000..ff96e3b54 --- /dev/null +++ b/dstack/gateway/src/kv/import.rs @@ -0,0 +1,504 @@ +// 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. +//! +//! 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::{ensure, Result}; +use base64::{engine::general_purpose::STANDARD, Engine}; + +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; + +/// 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 +/// 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. +pub struct AcceptedInstances { + pub instances: BTreeMap, + 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"); + ensure!( + !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}"))?; + 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. +/// +/// `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)?; + 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. + ensure!( + wg.is_routable_client_ip(data.ip), + "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(()) +} + +/// 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, 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, + } = 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) + .then_with(|| left_id.cmp(right_id)) + }); + + let mut instances = BTreeMap::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, now, &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; + } + 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; + 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: key(7), + 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 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 { + accept_instances_at(&wg_config(), loaded(records), NOW) + } + + #[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 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![ + ("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 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"), + "{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}" + ); + } + } + + #[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: + // 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_at( + &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(), + }, + 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 0fa35907e..1cc9a396a 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -29,15 +29,23 @@ //! - `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}; +use std::{ + collections::{BTreeMap, BTreeSet}, + net::Ipv4Addr, + path::Path, + time::Duration, +}; use anyhow::{Context, Result}; + +use crate::time::now_secs; use serde::{Deserialize, Serialize}; use tokio::sync::watch; use wavekv::{node::NodeState, types::NodeId, Node}; @@ -89,6 +97,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")] @@ -406,6 +430,40 @@ 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; + +/// 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 @@ -423,6 +481,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, @@ -432,6 +491,10 @@ trait GetPutCodec { &self, prefix: &str, ) -> impl Iterator; + fn iter_decoded_strict serde::Deserialize<'de>>( + &self, + prefix: &str, + ) -> impl Iterator)>; } impl GetPutCodec for NodeState { @@ -446,6 +509,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}"))?; @@ -483,6 +567,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). @@ -499,16 +602,84 @@ 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 + /// 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. + /// + /// 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, 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) 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 + // 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", + 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 @@ -563,16 +734,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 ==================== @@ -646,9 +829,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)| { @@ -656,14 +843,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() } @@ -675,9 +870,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)| { @@ -685,14 +881,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() } @@ -752,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}"); @@ -776,9 +977,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 @@ -803,9 +1010,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 @@ -817,19 +1027,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 @@ -865,7 +1082,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 @@ -876,13 +1098,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() } @@ -968,16 +1199,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") } @@ -989,9 +1213,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 @@ -1012,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) @@ -1053,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) @@ -1165,6 +1389,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!( @@ -1427,6 +1682,214 @@ 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 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"); + 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()); + } + + #[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); + } + + #[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"); + 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"); + 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()) + .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" + ); + } + + #[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)] mod peer_url_tests { use super::validate_peer_url; 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 de0d033a6..0a89fa3b4 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,11 +39,13 @@ 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, LoadedInstances, NodeData, NodeStatus, PortPolicy, + WaveKvSyncService, }, - models::{InstanceInfo, PortPolicyView, WgConf}, + models::{InstanceInfo, PortPolicyView, WgConf, WgPeer}, proxy::{create_acceptor_with_cert_resolver, AddressGroup, AddressInfo, AppAddressResolver}, + time::{decode_ts, encode_ts, now_secs}, }; mod auth_client; @@ -52,16 +53,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, @@ -124,6 +115,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 @@ -176,11 +170,12 @@ 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(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 @@ -257,6 +252,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 @@ -307,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, @@ -467,7 +460,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")? @@ -529,7 +524,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 +576,63 @@ 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: &[import::RejectedInstance]) { + for import::RejectedInstance { + instance_id, + reason, + .. + } in rejected + { + error!("ignoring KV instance record {instance_id}: {reason:#}"); + } +} + +/// 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(); + 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(), @@ -638,7 +685,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; @@ -876,11 +933,49 @@ 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 instances = store.load_all_instances(); + let accepted = import::accept_instances(&proxy.config.wg, store.load_all_instances()); + // 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(); + 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 + // here too, rather than lingering until this node's own recycle timeout. + 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()) + .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(), @@ -934,26 +1029,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 +1059,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 +1251,38 @@ 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 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", + 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()?) } @@ -1311,18 +1416,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) { @@ -1331,6 +1430,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(()) } @@ -1423,10 +1532,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) @@ -1492,16 +1598,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/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..75096d4d9 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,237 @@ 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) { + 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( + instance_id, + &InstanceData { + app_id: "peer-app".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, + }, + ) + .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"); +} + +#[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")); +} + +#[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")); +} + +#[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")); +} + +#[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()); +} 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)] 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()) +}