diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 31fa8ef89..e9886cd60 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -345,25 +345,45 @@ struct GatewayKeyStore { } impl GatewayKeyStore { - fn load() -> Option { - let content = fs::read_to_string(GATEWAY_CACHE_PATH).ok()?; + fn load_from(path: &Path) -> Option { + let content = fs::read_to_string(path).ok()?; serde_json::from_str(&content).ok() } - fn save(&self) -> Result<()> { + fn load() -> Option { + 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> { - 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)); + } +}