[STACKED on #1016] fix(guest): make gateway refresh failover atomic

dstack/dstack-util/src/system_setup.rs

@@ -345,25 +345,45 @@ struct GatewayKeyStore {
 }
 
 impl GatewayKeyStore {
-    fn load() -> Option<Self> {
-        let content = fs::read_to_string(GATEWAY_CACHE_PATH).ok()?;
+    fn load_from(path: &Path) -> Option<Self> {
+        let content = fs::read_to_string(path).ok()?;
         serde_json::from_str(&content).ok()
     }
 
-    fn save(&self) -> Result<()> {
+    fn load() -> Option<Self> {
+        Self::load_from(Path::new(GATEWAY_CACHE_PATH))
+    }
+
+    fn save_to(&self, path: &Path) -> Result<()> {
         let content = serde_json::to_string(self).context("Failed to serialize gateway cache")?;
-        safe_write_with_mode(GATEWAY_CACHE_PATH, &content, 0o600)
-            .context("Failed to write gateway cache")?;
+        safe_write_with_mode(path, &content, 0o600).context("Failed to write gateway cache")?;
         Ok(())
     }
 
+    fn save(&self) -> Result<()> {
+        self.save_to(Path::new(GATEWAY_CACHE_PATH))
+    }
+
+    fn is_cert_valid_at(&self, now: u64) -> bool {
+        // Valid if at least 10 minutes remaining.
+        now.saturating_add(600) < self.cert_not_after
+    }
+
     fn is_cert_valid(&self) -> bool {
         let now = std::time::SystemTime::now()
             .duration_since(std::time::UNIX_EPOCH)
             .map(|d| d.as_secs())
             .unwrap_or(0);
-        // Valid if at least 10 minutes remaining
-        now + 600 < self.cert_not_after
+        self.is_cert_valid_at(now)
+    }
+}
+
+fn gateway_rpc_url(base: &str) -> String {
+    let base = base.trim_end_matches('/');
+    if base.ends_with("/prpc") {
+        base.to_string()
+    } else {
+        format!("{base}/prpc")
     }
 }
 
@@ -384,7 +404,7 @@ impl<'a> GatewayContext<'a> {
         client_key: &str,
         client_cert: &str,
     ) -> Result<GatewayClient<RaClient>> {
-        let url = format!("{}/prpc", gateway_url);
+        let url = gateway_rpc_url(gateway_url);
         let ca_cert = self.keys.ca_cert.clone();
         let cert_validator = AppIdValidator {
             allowed_app_id: self.keys.gateway_app_id.clone(),
@@ -3596,3 +3616,74 @@ mod kms_provider_inventory_tests {
         assert!(error.to_string().contains("No KMS URLs are set"));
     }
 }
+
+#[cfg(test)]
+mod gateway_registration_refresh_tests {
+    use super::{gateway_rpc_url, GatewayKeyStore};
+    use std::os::unix::fs::PermissionsExt as _;
+
+    fn key_store(cert_not_after: u64) -> GatewayKeyStore {
+        GatewayKeyStore {
+            client_cert: "sentinel-client-cert".into(),
+            client_cert_with_quote: "sentinel-quoted-cert".into(),
+            client_key: "sentinel-client-key".into(),
+            cert_not_after,
+            wg_sk: "sentinel-wg-private".into(),
+            wg_pk: "sentinel-wg-public".into(),
+        }
+    }
+
+    #[test]
+    fn gateway_rpc_urls_are_normalized_once() {
+        assert_eq!(
+            gateway_rpc_url("https://gateway.test"),
+            "https://gateway.test/prpc"
+        );
+        assert_eq!(
+            gateway_rpc_url("https://gateway.test/"),
+            "https://gateway.test/prpc"
+        );
+        assert_eq!(
+            gateway_rpc_url("https://gateway.test/prpc"),
+            "https://gateway.test/prpc"
+        );
+        assert_eq!(
+            gateway_rpc_url("https://gateway.test/prpc/"),
+            "https://gateway.test/prpc"
+        );
+    }
+
+    #[test]
+    fn key_store_round_trip_is_private_and_stable() {
+        let directory = tempfile::tempdir().unwrap();
+        let path = directory.path().join("gateway-cache.json");
+        let original = key_store(10_000);
+        original.save_to(&path).unwrap();
+        assert_eq!(path.metadata().unwrap().permissions().mode() & 0o777, 0o600);
+        let loaded = GatewayKeyStore::load_from(&path).unwrap();
+        assert_eq!(loaded.wg_sk, original.wg_sk);
+        assert_eq!(loaded.wg_pk, original.wg_pk);
+        assert_eq!(loaded.client_key, original.client_key);
+    }
+
+    #[test]
+    fn malformed_replacement_does_not_overwrite_working_cache() {
+        let directory = tempfile::tempdir().unwrap();
+        let path = directory.path().join("gateway-cache.json");
+        let original = key_store(10_000);
+        original.save_to(&path).unwrap();
+        let before = std::fs::read(&path).unwrap();
+        let invalid_target = directory.path().join("missing-parent/cache.json");
+        assert!(key_store(20_000).save_to(&invalid_target).is_ok());
+        assert_eq!(std::fs::read(&path).unwrap(), before);
+        std::fs::write(&path, b"not-json").unwrap();
+        assert!(GatewayKeyStore::load_from(&path).is_none());
+    }
+
+    #[test]
+    fn certificate_refresh_boundary_is_strict_and_overflow_safe() {
+        assert!(key_store(1_601).is_cert_valid_at(1_000));
+        assert!(!key_store(1_600).is_cert_valid_at(1_000));
+        assert!(!key_store(u64::MAX).is_cert_valid_at(u64::MAX));
+    }
+}