Add onPayloadFinished callback

onPayloadFinished will be called from the VM when the payload has just
finished.

Bug: 195381416
Test: atest MicrodroidHostTestCases
Change-Id: I9bb68c2192ee19a40634d10ba77af9d4c3ec30d0
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
index c7a1471..4dc4333 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
@@ -32,6 +32,11 @@
     void onPayloadStarted(int cid, in @nullable ParcelFileDescriptor stream);
 
     /**
+     * Called when the payload has finished in the VM. `exitCode` is the exit code of the payload.
+     */
+    void onPayloadFinished(int cid, int exitCode);
+
+    /**
      * Called when the VM dies.
      *
      * Note that this will not be called if the VirtualizationService itself dies, so you should
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
index b5cda7d..730f6c2 100644
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
@@ -18,8 +18,14 @@
 /** {@hide} */
 interface IVirtualMachineService {
     /**
-     * Notifies that the virtual machine is ready.
+     * Notifies that the payload has started.
      * TODO(b/191845268): remove cid parameter
      */
     void notifyPayloadStarted(int cid);
+
+    /**
+     * Notifies that the payload has finished.
+     * TODO(b/191845268): remove cid parameter
+     */
+    void notifyPayloadFinished(int cid, int exitCode);
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index c5d3548..562843b 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -640,6 +640,16 @@
         }
     }
 
+    /// Call all registered callbacks to notify that the payload has finished.
+    pub fn notify_payload_finished(&self, cid: Cid, exit_code: i32) {
+        let callbacks = &*self.0.lock().unwrap();
+        for callback in callbacks {
+            if let Err(e) = callback.onPayloadFinished(cid as i32, exit_code) {
+                error!("Error notifying payload finish event from VM CID {}: {}", cid, e);
+            }
+        }
+    }
+
     /// Call all registered callbacks to say that the VM has died.
     pub fn callback_on_died(&self, cid: Cid) {
         let callbacks = &*self.0.lock().unwrap();
@@ -788,6 +798,21 @@
             ))
         }
     }
+
+    fn notifyPayloadFinished(&self, cid: i32, exit_code: i32) -> binder::Result<()> {
+        let cid = cid as Cid;
+        if let Some(vm) = self.state.lock().unwrap().get_vm(cid) {
+            info!("VM having CID {} finished payload", cid);
+            vm.callbacks.notify_payload_finished(cid, exit_code);
+            Ok(())
+        } else {
+            error!("notifyPayloadFinished is called from an unknown cid {}", cid);
+            Err(new_binder_exception(
+                ExceptionCode::SERVICE_SPECIFIC,
+                format!("cannot find a VM with cid {}", cid),
+            ))
+        }
+    }
 }
 
 impl VirtualMachineService {