From 07a3287bdd9be7e121cfe8d5fe26eb367a1bd83e Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 23 Sep 2026 22:41:40 -0700 Subject: [PATCH] fix(vmm): validate an update before writing any of it --- dstack/vmm/src/main_service.rs | 225 +++++++++++++++++++++------------ 1 file changed, 141 insertions(+), 84 deletions(-) diff --git a/dstack/vmm/src/main_service.rs b/dstack/vmm/src/main_service.rs index b2c34fe43..ce29fc748 100644 --- a/dstack/vmm/src/main_service.rs +++ b/dstack/vmm/src/main_service.rs @@ -872,7 +872,6 @@ impl RpcHandler { &self, vm_id: &str, manifest: &mut Manifest, - vm_work_dir: &VmWorkDir, vcpu: Option, memory: Option, disk_size: Option, @@ -902,48 +901,59 @@ impl RpcHandler { if disk_size < manifest.disk_size { bail!("Cannot shrink disk size"); } - if disk_size > manifest.disk_size { - let hda_path = vm_work_dir.hda_path(); - if hda_path.exists() { - info!("Resizing disk to {}GB", disk_size); - let new_size_str = format!("{}G", disk_size); - // Grow the disk the same way it was created, so a VM - // deployed with preallocation keeps its space reserved. - let mut args = vec!["resize".to_string()]; - if !manifest.disk_prealloc.is_off() { - args.push(format!( - "--preallocation={}", - manifest.disk_prealloc.as_str() - )); - } - args.push(hda_path.display().to_string()); - args.push(new_size_str); - // Off the async executor, for the same reason the launch - // path does it: growing a preallocated disk reserves the - // added space before it returns, and `full` writes it. - let output = tokio::task::spawn_blocking(move || { - std::process::Command::new("qemu-img").args(&args).output() - }) - .await - .context("disk resize task failed")? - .context("Failed to resize disk")?; - if !output.status.success() { - bail!( - "Failed to resize disk: {}", - String::from_utf8_lossy(&output.stderr) - ); - } - } else { - // A never-started stopped VM has no data disk yet. Its - // first launch creates hda.img from manifest.disk_size. - info!("Recording {}GB disk size for uninitialized VM", disk_size); - } - manifest.disk_size = disk_size; - } + manifest.disk_size = disk_size; } Ok(true) } + + /// Grows the data disk to `manifest.disk_size` if that exceeds `old_size`. + async fn grow_disk_if_needed( + vm_work_dir: &VmWorkDir, + manifest: &Manifest, + old_size: u32, + ) -> Result<()> { + let disk_size = manifest.disk_size; + if disk_size <= old_size { + return Ok(()); + } + let hda_path = vm_work_dir.hda_path(); + if !hda_path.exists() { + // A never-started stopped VM has no data disk yet. Its + // first launch creates hda.img from manifest.disk_size. + info!("Recording {}GB disk size for uninitialized VM", disk_size); + return Ok(()); + } + info!("Resizing disk to {}GB", disk_size); + let new_size_str = format!("{}G", disk_size); + // Grow the disk the same way it was created, so a VM + // deployed with preallocation keeps its space reserved. + let mut args = vec!["resize".to_string()]; + if !manifest.disk_prealloc.is_off() { + args.push(format!( + "--preallocation={}", + manifest.disk_prealloc.as_str() + )); + } + args.push(hda_path.display().to_string()); + args.push(new_size_str); + // Off the async executor, for the same reason the launch + // path does it: growing a preallocated disk reserves the + // added space before it returns, and `full` writes it. + let output = tokio::task::spawn_blocking(move || { + std::process::Command::new("qemu-img").args(&args).output() + }) + .await + .context("disk resize task failed")? + .context("Failed to resize disk")?; + if !output.status.success() { + bail!( + "Failed to resize disk: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + Ok(()) + } } impl VmmRpc for RpcHandler { @@ -1058,45 +1068,26 @@ impl VmmRpc for RpcHandler { // Again under the lock: removal can have claimed the VM while this // waited for it. self.app.refuse_if_removing(&request.id)?; + let vm_work_dir = self.app.work_dir(&request.id)?; + let mut manifest = vm_work_dir.manifest().context("Failed to read manifest")?; + let old_disk_size = manifest.disk_size; + // Validate the whole request against an in-memory manifest before + // touching the host, so a rejected update changes nothing. let new_id = if !request.compose_file.is_empty() { // check the compose file is valid let _app_compose: AppCompose = serde_json::from_str(&request.compose_file).context("Invalid compose file")?; - let compose_file_path = self.compose_file_path(&request.id)?; - if !compose_file_path.exists() { + if !self.compose_file_path(&request.id)?.exists() { bail!("The instance {} not found", request.id); } - // Read the manifest here rather than reusing the one below, so a - // rejected update leaves the stored compose file untouched. - let manifest = self - .app - .work_dir(&request.id)? - .manifest() - .context("Failed to read manifest")?; validate_disk_prealloc_against_compose(manifest.disk_prealloc, &request.compose_file)?; - fs::write(compose_file_path, &request.compose_file) - .context("Failed to write compose file")?; - app_id_of(&request.compose_file) } else { Default::default() }; - if !request.encrypted_env.is_empty() { - let encrypted_env_path = self.encrypted_env_path(&request.id)?; - fs::write(encrypted_env_path, &request.encrypted_env) - .context("Failed to write encrypted env")?; - } - if !request.user_config.is_empty() { - let user_config_path = self.user_config_path(&request.id)?; - fs::write(user_config_path, &request.user_config) - .context("Failed to write user config")?; - } - let vm_work_dir = self.app.work_dir(&request.id)?; - let mut manifest = vm_work_dir.manifest().context("Failed to read manifest")?; self.apply_resource_updates( &request.id, &mut manifest, - &vm_work_dir, request.vcpu, request.memory, request.disk_size, @@ -1125,7 +1116,7 @@ impl VmmRpc for RpcHandler { manifest.gateway_urls = request.gateway_urls.clone(); } if request.update_networking { - let networks = if request.networks.is_empty() { + manifest.networks = if request.networks.is_empty() { validate_default_network(&self.app.config.cvm)?; vec![] } else { @@ -1138,22 +1129,6 @@ impl VmmRpc for RpcHandler { let networks = networks_from_proto(&request.networks, &cvm)?; resolve_requested_networks(&networks, &cvm, manifest.vcpu)? }; - // Under the launch lock this whole call holds. Reading "not - // running" outside it and acting on the answer inside is the exact - // race the lock exists to close: a launch can start, prepare its - // interfaces and deploy QEMU in between, and the release would - // then delete the interfaces of a VM that is running -- silently, - // since QEMU stays up and the supervisor still reports it healthy. - let is_running = self - .app - .supervisor - .info(&request.id) - .await? - .is_some_and(|info| info.state.status.is_running()); - if !is_running && self.app.release_vm_interfaces(&request.id).await { - vm_work_dir.clear_runtime_networks()?; - } - manifest.networks = networks; } // Both only when this request moved one of the two halves, and after // both, since either half can move and the other still has to agree @@ -1168,12 +1143,49 @@ impl VmmRpc for RpcHandler { &resolved_nic_modes(&manifest.networks, &self.app.config.cvm, manifest.vcpu), )?; } - let compose_file = fs::read_to_string(vm_work_dir.app_compose_path()) - .context("failed to read app compose for swtpm decision")?; + let compose_file = if request.compose_file.is_empty() { + fs::read_to_string(vm_work_dir.app_compose_path()) + .context("failed to read app compose for swtpm decision")? + } else { + request.compose_file.clone() + }; manifest.swtpm = needs_swtpm( key_provider_from_compose(&compose_file)?, manifest.simulated_tee, ); + + Self::grow_disk_if_needed(&vm_work_dir, &manifest, old_disk_size).await?; + if request.update_networking { + // Under the launch lock this whole call holds. Reading "not + // running" outside it and acting on the answer inside is the exact + // race the lock exists to close: a launch can start, prepare its + // interfaces and deploy QEMU in between, and the release would + // then delete the interfaces of a VM that is running -- silently, + // since QEMU stays up and the supervisor still reports it healthy. + let is_running = self + .app + .supervisor + .info(&request.id) + .await? + .is_some_and(|info| info.state.status.is_running()); + if !is_running && self.app.release_vm_interfaces(&request.id).await { + vm_work_dir.clear_runtime_networks()?; + } + } + if !request.compose_file.is_empty() { + fs::write(self.compose_file_path(&request.id)?, &request.compose_file) + .context("Failed to write compose file")?; + } + if !request.encrypted_env.is_empty() { + let encrypted_env_path = self.encrypted_env_path(&request.id)?; + fs::write(encrypted_env_path, &request.encrypted_env) + .context("Failed to write encrypted env")?; + } + if !request.user_config.is_empty() { + let user_config_path = self.user_config_path(&request.id)?; + fs::write(user_config_path, &request.user_config) + .context("Failed to write user config")?; + } vm_work_dir .put_manifest(&manifest) .context("Failed to put manifest")?; @@ -1233,16 +1245,17 @@ impl VmmRpc for RpcHandler { self.app.refuse_if_removing(&request.id)?; let vm_work_dir = self.app.work_dir(&request.id)?; let mut manifest = vm_work_dir.manifest().context("failed to read manifest")?; + let old_disk_size = manifest.disk_size; self.apply_resource_updates( &request.id, &mut manifest, - &vm_work_dir, request.vcpu, request.memory, request.disk_size, request.image.as_deref(), ) .await?; + Self::grow_disk_if_needed(&vm_work_dir, &manifest, old_disk_size).await?; vm_work_dir .put_manifest(&manifest) .context("failed to update manifest")?; @@ -2895,6 +2908,50 @@ mod tests { assert!(port_map_from_proto(&[port(8080), port(8081)], &pm_cfg, &held).is_err()); } + #[tokio::test] + async fn a_rejected_update_writes_nothing() { + let tmp = tempfile::tempdir().unwrap(); + let mut config: crate::config::Config = + Figment::from(crate::config::load_config_figment(None)) + .extract() + .unwrap(); + config.run_path = tmp.path().to_path_buf(); + let manifest = + create_manifest_from_vm_config(test_vm_configuration(), &config.cvm).unwrap(); + let app = App::new( + config, + supervisor_client::SupervisorClient::new("http://127.0.0.1:0"), + ); + app.work_dir(&manifest.id) + .unwrap() + .put_manifest(&manifest) + .unwrap(); + let compose_path = app.compose_file_path(&manifest.id).unwrap(); + fs::create_dir_all(compose_path.parent().unwrap()).unwrap(); + fs::write(&compose_path, "{}").unwrap(); + + // Port mapping is disabled by default, so the new port is refused. + let err = RpcHandler { app } + .update_vm(UpdateVmRequest { + id: manifest.id.clone(), + compose_file: r#"{"manifest_version":"2","name":"new","runner":"docker-compose"}"# + .into(), + update_ports: true, + ports: vec![rpc::PortMapping { + protocol: "tcp".into(), + host_port: 8080, + vm_port: 8080, + host_address: String::new(), + nic_index: None, + }], + ..Default::default() + }) + .await + .unwrap_err(); + assert!(err.to_string().contains("disabled"), "{err:#}"); + assert_eq!(fs::read_to_string(&compose_path).unwrap(), "{}"); + } + #[test] fn resolve_volumes_attaches_duplicate_root_once() -> Result<()> { let tmp = tempfile::tempdir()?;