Merge pull request #1035 from Dstack-TEE/fix/gateway-kv-robustness

fix(gateway): harden the KV→data-plane boundary against bad replicated state

Author: kvinwang

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<DnsCredentialInfo> {
         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<GetDefaultDnsCredentialResponse> {
         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<CertbotConfigResponse> {
-        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<Proxy> 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<GlobalCertbotConfig>,
+    request: SetCertbotConfigRequest,
+) -> Result<GlobalCertbotConfig> {
+    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;

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<()> {

dstack/gateway/src/debug_service.rs

@@ -85,6 +85,7 @@ impl DebugRpc for DebugRpcHandler {
         // Get all instances
         let instances: Vec<InstanceEntry> = 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(),