diff --git a/vm/src/main.rs b/vm/src/main.rs
index f47d9e8..061cfd7 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -17,9 +17,14 @@
 mod sync;
 
 use android_system_virtmanager::aidl::android::system::virtmanager::IVirtManager::IVirtManager;
+use android_system_virtmanager::aidl::android::system::virtmanager::IVirtualMachine::IVirtualMachine;
+use android_system_virtmanager::aidl::android::system::virtmanager::IVirtualMachineCallback::{
+    BnVirtualMachineCallback, IVirtualMachineCallback,
+};
 use android_system_virtmanager::binder::{
     get_interface, DeathRecipient, IBinder, ParcelFileDescriptor, ProcessState, Strong,
 };
+use android_system_virtmanager::binder::{Interface, Result as BinderResult};
 use anyhow::{Context, Error};
 use std::fs::File;
 use std::io;
@@ -91,14 +96,26 @@
         // Pass the VM reference back to Virt Manager and have it hold it in the background.
         virt_manager.debugHoldVmRef(&vm).context("Failed to pass VM to Virt Manager")
     } else {
-        // Wait until the VM dies. If we just returned immediately then the IVirtualMachine Binder
-        // object would be dropped and the VM would be killed.
-        wait_for_death(&mut vm.as_binder())?;
-        println!("VM died");
-        Ok(())
+        // Wait until the VM or VirtManager dies. If we just returned immediately then the
+        // IVirtualMachine Binder object would be dropped and the VM would be killed.
+        wait_for_vm(vm)
     }
 }
 
+/// Wait until the given VM or the VirtManager itself dies.
+fn wait_for_vm(vm: Strong<dyn IVirtualMachine>) -> Result<(), Error> {
+    let dead = AtomicFlag::default();
+    let callback =
+        BnVirtualMachineCallback::new_binder(VirtualMachineCallback { dead: dead.clone() });
+    vm.registerCallback(&callback)?;
+    let death_recipient = wait_for_death(&mut vm.as_binder(), dead.clone())?;
+    dead.wait();
+    // Ensure that death_recipient isn't dropped before we wait on the flag, as it is removed
+    // from the Binder when it's dropped.
+    drop(death_recipient);
+    Ok(())
+}
+
 /// Retrieve reference to a previously daemonized VM and stop it.
 fn command_stop(virt_manager: Strong<dyn IVirtManager>, cid: u32) -> Result<(), Error> {
     virt_manager
@@ -115,18 +132,31 @@
     Ok(())
 }
 
-/// Block until the given Binder object dies.
-fn wait_for_death(binder: &mut impl IBinder) -> Result<(), Error> {
-    let dead = AtomicFlag::default();
-    let mut death_recipient = {
-        let dead = dead.clone();
-        DeathRecipient::new(move || {
-            dead.raise();
-        })
-    };
+/// Raise the given flag when the given Binder object dies.
+///
+/// If the returned DeathRecipient is dropped then this will no longer do anything.
+fn wait_for_death(binder: &mut impl IBinder, dead: AtomicFlag) -> Result<DeathRecipient, Error> {
+    let mut death_recipient = DeathRecipient::new(move || {
+        println!("VirtManager died");
+        dead.raise();
+    });
     binder.link_to_death(&mut death_recipient)?;
-    dead.wait();
-    Ok(())
+    Ok(death_recipient)
+}
+
+#[derive(Debug)]
+struct VirtualMachineCallback {
+    dead: AtomicFlag,
+}
+
+impl Interface for VirtualMachineCallback {}
+
+impl IVirtualMachineCallback for VirtualMachineCallback {
+    fn onDied(&self, _cid: i32) -> BinderResult<()> {
+        println!("VM died");
+        self.dead.raise();
+        Ok(())
+    }
 }
 
 /// Safely duplicate the standard output file descriptor.