diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index 68202032d..d410aa451 100644 --- a/dstack/gateway/rpc/proto/gateway_rpc.proto +++ b/dstack/gateway/rpc/proto/gateway_rpc.proto @@ -434,6 +434,9 @@ service Admin { rpc GetGlobalConnections(google.protobuf.Empty) returns (GlobalConnectionsStats) {} // Get all node statuses rpc GetNodeStatuses(google.protobuf.Empty) returns (GetNodeStatusesResponse) {} + // Remove a CVM from WaveKV and the local data plane. This is an idempotent + // operator recovery action and also works when the stored record is unreadable. + rpc RemoveCvm(RemoveCvmRequest) returns (RemoveCvmResponse) {} // ==================== DNS Credential Management ==================== // List all DNS credentials @@ -497,6 +500,23 @@ service Admin { rpc GetInstancePortPolicy(GetInstancePortPolicyRequest) returns (GetInstancePortPolicyResponse) {} } +// Emergency operator request to remove one CVM's instance record. +message RemoveCvmRequest { + string instance_id = 1; +} + +// Outcome of a RemoveCvm request. Both fields are false when the request +// names an instance this cluster has never seen (or a retry of a removal +// that already completed), so a mistyped instance_id is visible to the +// operator instead of silently reporting success. +message RemoveCvmResponse { + // Whether a live instance record existed in WaveKV before the tombstone + // was written. Also true for records that existed but were unreadable. + bool record_existed = 1; + // Whether the CVM was present in this node's local data plane. + bool removed_locally = 2; +} + // ==================== DNS Credential Messages ==================== // DNS credential information diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index bd8033119..d5c1f6710 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -17,11 +17,12 @@ use dstack_gateway_rpc::{ HandshakeEntry, HostInfo, LastSeenEntry, ListCertAttestationsRequest, ListCertAttestationsResponse, ListDnsCredentialsResponse, ListZtDomainsResponse, NodeStatusEntry, PeerSyncStatus as ProtoPeerSyncStatus, PortAttrs as RpcPortAttrs, - PortPolicy as RpcPortPolicy, RenewCertResponse, RenewZtDomainCertRequest, - RenewZtDomainCertResponse, RotateAcmeCredentialsResponse, SetCertbotConfigRequest, - SetDefaultDnsCredentialRequest, SetInstancePortPolicyRequest, SetNodeStatusRequest, - SetNodeUrlRequest, StatusResponse, StoreSyncStatus, UpdateDnsCredentialRequest, - WaveKvStatusResponse, ZtDomainCertStatus, ZtDomainConfig as ProtoZtDomainConfig, ZtDomainInfo, + PortPolicy as RpcPortPolicy, RemoveCvmRequest, RemoveCvmResponse, RenewCertResponse, + RenewZtDomainCertRequest, RenewZtDomainCertResponse, RotateAcmeCredentialsResponse, + SetCertbotConfigRequest, SetDefaultDnsCredentialRequest, SetInstancePortPolicyRequest, + SetNodeStatusRequest, SetNodeUrlRequest, StatusResponse, StoreSyncStatus, + UpdateDnsCredentialRequest, WaveKvStatusResponse, ZtDomainCertStatus, + ZtDomainConfig as ProtoZtDomainConfig, ZtDomainInfo, }; use ra_rpc::{CallContext, RpcCall}; use tracing::{info, warn}; @@ -305,6 +306,25 @@ impl AdminRpc for AdminRpcHandler { Ok(GetNodeStatusesResponse { statuses: entries }) } + async fn remove_cvm(self, request: RemoveCvmRequest) -> Result { + let instance_id = request.instance_id.as_str(); + // Same bound the KV import boundary puts on identifiers. Legitimate + // gateways never write an instance_id outside it, so this rejects only + // typos — and keeps the ID safe to embed in logs and KV keys. + crate::kv::import::validate_id("instance_id", instance_id)?; + + let removal = self.state.remove_cvm(instance_id)?; + warn!( + "admin removed CVM {instance_id} from WaveKV and the local data plane \ + (record existed: {}, present locally: {})", + removal.record_existed, removal.removed_locally + ); + Ok(RemoveCvmResponse { + record_existed: removal.record_existed, + removed_locally: removal.removed_locally, + }) + } + // ==================== DNS Credential Management ==================== async fn list_dns_credentials(self) -> Result { diff --git a/dstack/gateway/src/kv/import.rs b/dstack/gateway/src/kv/import.rs index ff96e3b54..64d0ae410 100644 --- a/dstack/gateway/src/kv/import.rs +++ b/dstack/gateway/src/kv/import.rs @@ -113,7 +113,11 @@ pub fn validate_wg_public_key(public_key: &str) -> Result<()> { Ok(()) } -fn validate_id(field: &str, value: &str) -> Result<()> { +/// Validate an identifier field: non-empty, bounded, and free of whitespace +/// and control characters. Every identifier a legitimate gateway writes into +/// a KV record satisfies this, so it is also the acceptance bound for +/// operator-supplied instance IDs (e.g. `Admin.RemoveCvm`). +pub(crate) fn validate_id(field: &str, value: &str) -> Result<()> { ensure!(!value.is_empty(), "{field} is empty"); ensure!( value.len() <= MAX_ID_LEN, @@ -322,6 +326,19 @@ mod tests { assert!(validate_wg_public_key(&key(3)).is_ok()); } + #[test] + fn rejects_ids_unfit_for_kv_keys_and_logs() { + let too_long = "a".repeat(MAX_ID_LEN + 1); + for bad in ["", " id", "id ", "in id", "in\nid", too_long.as_str()] { + assert!( + validate_id("instance_id", bad).is_err(), + "accepted id {bad:?}" + ); + } + validate_id("instance_id", "peer-instance").unwrap(); + validate_id("instance_id", &"a".repeat(MAX_ID_LEN)).unwrap(); + } + #[test] fn one_bad_record_does_not_drop_the_others() { let accepted = accept(vec![ diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index f10aa5b4a..1c6ce5705 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -727,8 +727,11 @@ impl KvStore { } /// Sync instance deletion to other nodes - pub fn sync_delete_instance(&self, instance_id: &str) -> Result<()> { - self.persistent.write().delete(keys::inst(instance_id))?; + /// + /// Returns whether a live record (including an undecodable one) existed + /// before the tombstone was written. + pub fn sync_delete_instance(&self, instance_id: &str) -> Result { + let previous = self.persistent.write().delete(keys::inst(instance_id))?; self.ephemeral .write() .delete(keys::conn(instance_id, self.my_node_id))?; @@ -736,7 +739,7 @@ impl KvStore { self.ephemeral .write() .delete(keys::handshake(instance_id, self.my_node_id))?; - Ok(()) + Ok(previous.is_some_and(|entry| !entry.is_deleted())) } /// Load all instances from the sync store. diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index 658f98e37..15f1e48af 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -128,7 +128,44 @@ pub struct ProxyOptions { pub tls_config: TlsConfig, } +/// Outcome of an operator-initiated CVM removal. +/// +/// Both fields are false when the instance was never known (or the removal +/// already completed), which lets the operator distinguish a mistyped +/// instance_id from an actual removal. +pub struct CvmRemoval { + /// A live instance record (decodable or not) existed in WaveKV. + pub record_existed: bool, + /// The CVM was present in this node's local data plane. + pub removed_locally: bool, +} + impl Proxy { + /// Remove one CVM by explicit operator request. + /// + /// The tombstone is written even when this node cannot decode the stored + /// record or no longer has the CVM in memory. This makes the operation an + /// idempotent recovery path for bad replicated instance records without + /// exposing arbitrary raw-KV deletion. + pub fn remove_cvm(&self, instance_id: &str) -> Result { + let mut state = self.lock(); + let record_existed = state + .kv_store + .sync_delete_instance(instance_id) + .with_context(|| format!("failed to delete CVM {instance_id} from WaveKV"))?; + + let removed_locally = state.forget_instance(instance_id).is_some(); + // Reconfigure unconditionally: the tombstone write and the in-memory + // removal are not repeated on a retry, so gating this on them would + // leave a failed reconfigure with no retry path and the removed CVM's + // WireGuard peer stuck on the interface. + state.reconfigure()?; + Ok(CvmRemoval { + record_existed, + removed_locally, + }) + } + pub async fn new(options: ProxyOptions) -> Result { let (port_policy_tx, port_policy_rx) = unbounded_channel(); let inner = ProxyInner::new(options, port_policy_tx).await?; diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs index 75096d4d9..a87d18394 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -623,6 +623,47 @@ async fn an_undecodable_record_keeps_the_instance_it_describes() { assert!(state.lock().state.instances.contains_key("peer-instance")); } +#[tokio::test] +async fn an_operator_can_remove_a_cvm_whose_kv_record_is_unreadable() { + 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(); + + state + .kv_store + .persistent() + .write() + .put( + crate::kv::keys::inst("peer-instance"), + b"not-messagepack".to_vec(), + ) + .unwrap(); + + let removal = state.proxy.remove_cvm("peer-instance").unwrap(); + assert!(removal.record_existed); + assert!(removal.removed_locally); + assert!(!state.lock().state.instances.contains_key("peer-instance")); + let loaded = state.kv_store.load_all_instances(); + assert!(!loaded.decoded.contains_key("peer-instance")); + assert!(!loaded.undecodable.contains("peer-instance")); + + // The recovery operation is safe to retry after a timeout or lost reply, + // and the retry tells the operator there was nothing left to remove. + let retry = state.proxy.remove_cvm("peer-instance").unwrap(); + assert!(!retry.record_existed); + assert!(!retry.removed_locally); +} + +#[tokio::test] +async fn removing_an_unknown_cvm_reports_that_nothing_existed() { + let state = create_test_state().await; + + // A mistyped instance_id must not be mistaken for a successful removal. + let removal = state.proxy.remove_cvm("no-such-instance").unwrap(); + assert!(!removal.record_existed); + assert!(!removal.removed_locally); +} + #[tokio::test] async fn an_instance_that_lost_an_ip_conflict_stops_being_routable() { let state = create_test_state().await;