Merge "Store apk root hash to the instance disk"
diff --git a/apex/Android.bp b/apex/Android.bp
index ccf34fd..0e2d2d4 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -18,7 +18,6 @@
     arch: {
         arm64: {
             binaries: [
-                "authfs", // TODO(victorhsieh): move to microdroid once we can run the test in VM.
                 "crosvm",
                 "virtualizationservice",
             ],
@@ -32,7 +31,6 @@
         },
         x86_64: {
             binaries: [
-                "authfs", // TODO(victorhsieh): move to microdroid once we can run the test in VM.
                 "crosvm",
                 "virtualizationservice",
             ],
diff --git a/compos/compos_key_cmd/compos_key_cmd.cpp b/compos/compos_key_cmd/compos_key_cmd.cpp
index d847b6d..874a208 100644
--- a/compos/compos_key_cmd/compos_key_cmd.cpp
+++ b/compos/compos_key_cmd/compos_key_cmd.cpp
@@ -112,6 +112,16 @@
         return ScopedAStatus::ok();
     }
 
+    ::ndk::ScopedAStatus onPayloadReady(int32_t in_cid) override {
+        LOG(INFO) << "Payload is ready! cid = " << in_cid;
+        return ScopedAStatus::ok();
+    }
+
+    ::ndk::ScopedAStatus onPayloadFinished(int32_t in_cid, int32_t in_exit_code) override {
+        LOG(INFO) << "Payload finished! cid = " << in_cid << ", exit_code = " << in_exit_code;
+        return ScopedAStatus::ok();
+    }
+
     ::ndk::ScopedAStatus onDied(int32_t in_cid) override {
         LOG(WARNING) << "VM died! cid = " << in_cid;
         {
diff --git a/demo/java/com/android/microdroid/demo/MainActivity.java b/demo/java/com/android/microdroid/demo/MainActivity.java
index 8974475..6a46f73 100644
--- a/demo/java/com/android/microdroid/demo/MainActivity.java
+++ b/demo/java/com/android/microdroid/demo/MainActivity.java
@@ -176,6 +176,24 @@
                         }
 
                         @Override
+                        public void onPayloadReady(VirtualMachine vm) {
+                            // This check doesn't 100% prevent race condition, but is fine for demo.
+                            if (!mService.isShutdown()) {
+                                mPayloadOutput.postValue("(Payload is ready)");
+                            }
+                        }
+
+                        @Override
+                        public void onPayloadFinished(VirtualMachine vm, int exitCode) {
+                            // This check doesn't 100% prevent race condition, but is fine for demo.
+                            if (!mService.isShutdown()) {
+                                mPayloadOutput.postValue(
+                                        String.format(
+                                                "(Payload finished. exit code: %d)", exitCode));
+                            }
+                        }
+
+                        @Override
                         public void onDied(VirtualMachine vm) {
                             mService.shutdownNow();
                             mStatus.postValue(VirtualMachine.Status.STOPPED);
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 959e355..f4ac467 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -303,6 +303,24 @@
                         }
 
                         @Override
+                        public void onPayloadReady(int cid) {
+                            final VirtualMachineCallback cb = mCallback;
+                            if (cb == null) {
+                                return;
+                            }
+                            cb.onPayloadReady(VirtualMachine.this);
+                        }
+
+                        @Override
+                        public void onPayloadFinished(int cid, int exitCode) {
+                            final VirtualMachineCallback cb = mCallback;
+                            if (cb == null) {
+                                return;
+                            }
+                            cb.onPayloadFinished(VirtualMachine.this, exitCode);
+                        }
+
+                        @Override
                         public void onDied(int cid) {
                             final VirtualMachineCallback cb = mCallback;
                             if (cb == null) {
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java b/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
index 89bb260..988acd7 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
@@ -31,6 +31,12 @@
     /** Called when the payload starts in the VM. */
     void onPayloadStarted(@NonNull VirtualMachine vm, @Nullable ParcelFileDescriptor stream);
 
+    /** Called when the payload in the VM is ready to serve. */
+    void onPayloadReady(@NonNull VirtualMachine vm);
+
+    /** Called when the payload has finished in the VM. */
+    void onPayloadFinished(@NonNull VirtualMachine vm, int exitCode);
+
     /** Called when the VM died. */
     void onDied(@NonNull VirtualMachine vm);
 }
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 06dd3c6..278a14d 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -162,17 +162,23 @@
     info!("executing main task {:?}...", task);
     let mut child = build_command(task)?.spawn()?;
 
+    let local_cid = get_local_cid()?;
     info!("notifying payload started");
-    service.notifyPayloadStarted(get_local_cid()? as i32)?;
+    service.notifyPayloadStarted(local_cid as i32)?;
 
-    match child.wait()?.code() {
-        Some(0) => {
+    if let Some(code) = child.wait()?.code() {
+        info!("notifying payload finished");
+        service.notifyPayloadFinished(local_cid as i32, code)?;
+
+        if code == 0 {
             info!("task successfully finished");
-            Ok(())
+        } else {
+            error!("task exited with exit code: {}", code);
         }
-        Some(code) => bail!("task exited with exit code: {}", code),
-        None => bail!("task terminated by signal"),
+    } else {
+        error!("task terminated by signal");
     }
+    Ok(())
 }
 
 fn build_command(task: &Task) -> Result<Command> {
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
index c7a1471..15354a3 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
@@ -32,6 +32,16 @@
     void onPayloadStarted(int cid, in @nullable ParcelFileDescriptor stream);
 
     /**
+     * Called when the payload in the VM is ready to serve.
+     */
+    void onPayloadReady(int cid);
+
+    /**
+     * 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..10b14e0 100644
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
@@ -18,8 +18,20 @@
 /** {@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 is ready to serve.
+     * TODO(b/191845268): remove cid parameter
+     */
+    void notifyPayloadReady(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..6b60da3 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -640,6 +640,26 @@
         }
     }
 
+    /// Call all registered callbacks to notify that the payload is ready to serve.
+    pub fn notify_payload_ready(&self, cid: Cid) {
+        let callbacks = &*self.0.lock().unwrap();
+        for callback in callbacks {
+            if let Err(e) = callback.onPayloadReady(cid as i32) {
+                error!("Error notifying payload ready event from VM CID {}: {}", cid, e);
+            }
+        }
+    }
+
+    /// 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 +808,36 @@
             ))
         }
     }
+
+    fn notifyPayloadReady(&self, cid: i32) -> binder::Result<()> {
+        let cid = cid as Cid;
+        if let Some(vm) = self.state.lock().unwrap().get_vm(cid) {
+            info!("VM having CID {} payload is ready", cid);
+            vm.callbacks.notify_payload_ready(cid);
+            Ok(())
+        } else {
+            error!("notifyPayloadReady is called from an unknown cid {}", cid);
+            Err(new_binder_exception(
+                ExceptionCode::SERVICE_SPECIFIC,
+                format!("cannot find a VM with cid {}", cid),
+            ))
+        }
+    }
+
+    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 {
diff --git a/vm/src/run.rs b/vm/src/run.rs
index 3b34bac..4b656ae 100644
--- a/vm/src/run.rs
+++ b/vm/src/run.rs
@@ -189,6 +189,14 @@
         Ok(())
     }
 
+    fn onPayloadReady(&self, _cid: i32) -> BinderResult<()> {
+        Ok(())
+    }
+
+    fn onPayloadFinished(&self, _cid: i32, _exit_code: i32) -> BinderResult<()> {
+        Ok(())
+    }
+
     fn onDied(&self, _cid: i32) -> BinderResult<()> {
         // No need to explicitly report the event to the user (e.g. via println!) because this
         // callback is registered only when the vm tool is invoked as interactive mode (e.g. not