From f58d42f78e648732853d4840702cb2e6163f2e19 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 12 Aug 2026 01:10:47 -0700 Subject: [PATCH 1/3] feat(gateway): add admin CVM removal API --- dstack/gateway/rpc/proto/gateway_rpc.proto | 8 +++++++ dstack/gateway/src/admin_service.rs | 18 ++++++++++++++- dstack/gateway/src/main_service.rs | 20 +++++++++++++++++ dstack/gateway/src/main_service/tests.rs | 26 ++++++++++++++++++++++ 4 files changed, 71 insertions(+), 1 deletion(-) diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index 68202032d..0ab66dfbe 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 (google.protobuf.Empty) {} // ==================== DNS Credential Management ==================== // List all DNS credentials @@ -499,6 +502,11 @@ service Admin { // ==================== DNS Credential Messages ==================== +// Emergency operator request to remove one CVM's instance record. +message RemoveCvmRequest { + string instance_id = 1; +} + // DNS credential information message DnsCredentialInfo { string id = 1; diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index bd8033119..6995f3827 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -17,7 +17,7 @@ use dstack_gateway_rpc::{ HandshakeEntry, HostInfo, LastSeenEntry, ListCertAttestationsRequest, ListCertAttestationsResponse, ListDnsCredentialsResponse, ListZtDomainsResponse, NodeStatusEntry, PeerSyncStatus as ProtoPeerSyncStatus, PortAttrs as RpcPortAttrs, - PortPolicy as RpcPortPolicy, RenewCertResponse, RenewZtDomainCertRequest, + PortPolicy as RpcPortPolicy, RemoveCvmRequest, RenewCertResponse, RenewZtDomainCertRequest, RenewZtDomainCertResponse, RotateAcmeCredentialsResponse, SetCertbotConfigRequest, SetDefaultDnsCredentialRequest, SetInstancePortPolicyRequest, SetNodeStatusRequest, SetNodeUrlRequest, StatusResponse, StoreSyncStatus, UpdateDnsCredentialRequest, @@ -305,6 +305,22 @@ impl AdminRpc for AdminRpcHandler { Ok(GetNodeStatusesResponse { statuses: entries }) } + async fn remove_cvm(self, request: RemoveCvmRequest) -> Result<()> { + let instance_id = request.instance_id.trim(); + ensure!(!instance_id.is_empty(), "instance_id is required"); + ensure!( + instance_id == request.instance_id, + "instance_id must not have leading or trailing whitespace" + ); + + let removed_locally = self.state.remove_cvm(instance_id)?; + warn!( + "Admin removed CVM {instance_id} from WaveKV and the local data plane \ + (present locally: {removed_locally})" + ); + Ok(()) + } + // ==================== DNS Credential Management ==================== async fn list_dns_credentials(self) -> Result { diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index 658f98e37..37a7054d3 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -129,6 +129,26 @@ pub struct ProxyOptions { } 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(); + state + .kv_store + .sync_delete_instance(instance_id) + .with_context(|| format!("failed to delete CVM {instance_id} from WaveKV"))?; + + let removed = state.forget_instance(instance_id).is_some(); + if removed { + state.reconfigure()?; + } + Ok(removed) + } + 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..393d4d5d1 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -623,6 +623,32 @@ 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(); + + assert!(state.proxy.remove_cvm("peer-instance").unwrap()); + 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. + assert!(!state.proxy.remove_cvm("peer-instance").unwrap()); +} + #[tokio::test] async fn an_instance_that_lost_an_ip_conflict_stops_being_routable() { let state = create_test_state().await; From 51207f78299899fed72982d18daff79faeaf3304 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 12 Aug 2026 01:36:14 -0700 Subject: [PATCH 2/3] fix(gateway): reconfigure unconditionally in RemoveCvm and report outcome Address review findings on the RemoveCvm admin RPC: - Reconfigure WireGuard unconditionally. The tombstone write and the in-memory removal are not repeated on a retry, so gating reconfigure on them left a failed reconfigure with no retry path and the removed CVM's WireGuard peer stuck on the interface. - Return record_existed/removed_locally to the operator. A mistyped instance_id still writes a tombstone, so it previously reported an indistinguishable success; now both fields come back false. - Move RemoveCvmRequest out of the DNS credential section in the proto. - Start the admin removal log message with lowercase per code style. --- dstack/gateway/rpc/proto/gateway_rpc.proto | 18 ++++++++++--- dstack/gateway/src/admin_service.rs | 25 ++++++++++------- dstack/gateway/src/kv/mod.rs | 9 ++++--- dstack/gateway/src/main_service.rs | 31 +++++++++++++++++----- dstack/gateway/src/main_service/tests.rs | 21 ++++++++++++--- 5 files changed, 78 insertions(+), 26 deletions(-) diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index 0ab66dfbe..d410aa451 100644 --- a/dstack/gateway/rpc/proto/gateway_rpc.proto +++ b/dstack/gateway/rpc/proto/gateway_rpc.proto @@ -436,7 +436,7 @@ service Admin { 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 (google.protobuf.Empty) {} + rpc RemoveCvm(RemoveCvmRequest) returns (RemoveCvmResponse) {} // ==================== DNS Credential Management ==================== // List all DNS credentials @@ -500,13 +500,25 @@ service Admin { rpc GetInstancePortPolicy(GetInstancePortPolicyRequest) returns (GetInstancePortPolicyResponse) {} } -// ==================== DNS Credential Messages ==================== - // 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 message DnsCredentialInfo { string id = 1; diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index 6995f3827..605f571f7 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, RemoveCvmRequest, 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,7 +306,7 @@ impl AdminRpc for AdminRpcHandler { Ok(GetNodeStatusesResponse { statuses: entries }) } - async fn remove_cvm(self, request: RemoveCvmRequest) -> Result<()> { + async fn remove_cvm(self, request: RemoveCvmRequest) -> Result { let instance_id = request.instance_id.trim(); ensure!(!instance_id.is_empty(), "instance_id is required"); ensure!( @@ -313,12 +314,16 @@ impl AdminRpc for AdminRpcHandler { "instance_id must not have leading or trailing whitespace" ); - let removed_locally = self.state.remove_cvm(instance_id)?; + let removal = self.state.remove_cvm(instance_id)?; warn!( - "Admin removed CVM {instance_id} from WaveKV and the local data plane \ - (present locally: {removed_locally})" + "admin removed CVM {instance_id} from WaveKV and the local data plane \ + (record existed: {}, present locally: {})", + removal.record_existed, removal.removed_locally ); - Ok(()) + Ok(RemoveCvmResponse { + record_existed: removal.record_existed, + removed_locally: removal.removed_locally, + }) } // ==================== DNS Credential Management ==================== 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 37a7054d3..15f1e48af 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -128,6 +128,18 @@ 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. /// @@ -135,18 +147,23 @@ impl Proxy { /// 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 { + pub fn remove_cvm(&self, instance_id: &str) -> Result { let mut state = self.lock(); - state + let record_existed = state .kv_store .sync_delete_instance(instance_id) .with_context(|| format!("failed to delete CVM {instance_id} from WaveKV"))?; - let removed = state.forget_instance(instance_id).is_some(); - if removed { - state.reconfigure()?; - } - Ok(removed) + 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 { diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs index 393d4d5d1..a87d18394 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -639,14 +639,29 @@ async fn an_operator_can_remove_a_cvm_whose_kv_record_is_unreadable() { ) .unwrap(); - assert!(state.proxy.remove_cvm("peer-instance").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. - assert!(!state.proxy.remove_cvm("peer-instance").unwrap()); + // 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] From 4ec05631e71bd7a51e8de900a10734d082a8a630 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 12 Aug 2026 01:45:09 -0700 Subject: [PATCH 3/3] fix(gateway): validate RemoveCvm instance_id like the KV import boundary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The handler only rejected leading/trailing whitespace, so an ID with internal whitespace, control characters (log injection via the warn! audit line), or unbounded length was accepted and written into a KV tombstone key. Reuse import::validate_id — every identifier a legitimate gateway writes satisfies it, so this rejects only typos. --- dstack/gateway/src/admin_service.rs | 11 +++++------ dstack/gateway/src/kv/import.rs | 19 ++++++++++++++++++- 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index 605f571f7..d5c1f6710 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -307,12 +307,11 @@ impl AdminRpc for AdminRpcHandler { } async fn remove_cvm(self, request: RemoveCvmRequest) -> Result { - let instance_id = request.instance_id.trim(); - ensure!(!instance_id.is_empty(), "instance_id is required"); - ensure!( - instance_id == request.instance_id, - "instance_id must not have leading or trailing whitespace" - ); + 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!( 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![