Merge "Mount system and vendor over dm-verity"
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/Android.bp b/microdroid/Android.bp
index 17cb615..1638f6f 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -52,7 +52,6 @@
         "microdroid_build_prop",
         "microdroid_init_rc",
         "microdroid_launcher",
-        "microdroid_manager",
 
         "ueventd.rc",
         "libbinder",
@@ -112,6 +111,7 @@
                 "apkdmverity",
                 "authfs",
                 "authfs_service",
+                "microdroid_manager",
                 "zipfuse",
 
                 // TODO(b/184872979): Needed by authfs. Remove once the Rust API is created.
diff --git a/microdroid_manager/Android.bp b/microdroid_manager/Android.bp
index 5fae7b1..0b1bc33 100644
--- a/microdroid_manager/Android.bp
+++ b/microdroid_manager/Android.bp
@@ -15,6 +15,8 @@
         "libapkverify",
         "libbinder_rpc_unstable_bindgen",
         "libbinder_rs",
+        "libbyteorder",
+        "libidsig",
         "libkernlog",
         "liblibc",
         "liblog_rust",
@@ -22,15 +24,23 @@
         "libmicrodroid_payload_config",
         "libnix",
         "libprotobuf",
+        "libring",
         "librustutils",
         "libserde",
         "libserde_json",
+        "libuuid",
         "libvsock",
+        "librand",
     ],
     shared_libs: [
         "libbinder_rpc_unstable",
     ],
     init_rc: ["microdroid_manager.rc"],
+    multilib: {
+        lib32: {
+            enabled: false,
+        },
+    },
 }
 
 rust_binary {
diff --git a/microdroid_manager/src/instance.rs b/microdroid_manager/src/instance.rs
new file mode 100644
index 0000000..3ef3b63
--- /dev/null
+++ b/microdroid_manager/src/instance.rs
@@ -0,0 +1,309 @@
+// Copyright 2021, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Provides routines to read/write on the instance disk.
+//!
+//! Instance disk is a disk where the identity of a VM instance is recorded. The identity usually
+//! includes certificates of the VM payload that is trusted, but not limited to it. Instance disk
+//! is empty when a VM is first booted. The identity data is filled in during the first boot, and
+//! then encrypted and signed. Subsequent boots decrypts and authenticates the data and uses the
+//! identity data to further verify the payload (e.g. against the certificate).
+//!
+//! Instance disk consists of a disk header and one or more partitions each of which consists of a
+//! header and payload. Each header (both the disk header and a partition header) is 512 bytes
+//! long. Payload is just next to the header and its size can be arbitrary. Headers are located at
+//! 512 bytes boundaries. So, when the size of a payload is not multiple of 512, there exists a gap
+//! between the end of the payload and the start of the next partition (if there is any).
+//!
+//! Each partition is identified by a UUID. A partition is created for a program loader that
+//! participates in the boot chain of the VM. Each program loader is expected to locate the
+//! partition that corresponds to the loader using the UUID that is assigned to the loader.
+//!
+//! The payload of a partition is encrypted/signed by a key that is unique to the loader and to the
+//! VM as well. Failing to decrypt/authenticate a partition by a loader stops the boot process.
+
+use anyhow::{anyhow, bail, Context, Result};
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+use ring::aead::{Aad, Algorithm, LessSafeKey, Nonce, UnboundKey, AES_256_GCM};
+use ring::hkdf::{Salt, HKDF_SHA256};
+use std::fs::{File, OpenOptions};
+use std::io::{Read, Seek, SeekFrom, Write};
+use uuid::Uuid;
+
+/// Path to the instance disk inside the VM
+const INSTANCE_IMAGE_PATH: &str = "/dev/block/by-name/vm-instance";
+
+/// Magic string in the instance disk header
+const DISK_HEADER_MAGIC: &str = "Android-VM-instance";
+
+/// Version of the instance disk format
+const DISK_HEADER_VERSION: u16 = 1;
+
+/// Size of the headers in the instance disk
+const DISK_HEADER_SIZE: u64 = 512;
+const PARTITION_HEADER_SIZE: u64 = 512;
+
+/// UUID of the partition that microdroid manager uses
+const MICRODROID_PARTITION_UUID: &str = "cf9afe9a-0662-11ec-a329-c32663a09d75";
+
+/// Encryption algorithm used to cipher payload
+static ENCRYPT_ALG: &Algorithm = &AES_256_GCM;
+
+/// Handle to the instance disk
+pub struct InstanceDisk {
+    file: File,
+}
+
+/// Information from a partition header
+struct PartitionHeader {
+    uuid: Uuid,
+    payload_size: u64, // in bytes
+}
+
+/// Offset of a partition in the instance disk
+type PartitionOffset = u64;
+
+impl InstanceDisk {
+    /// Creates handle to instance disk
+    pub fn new() -> Result<Self> {
+        let mut file = OpenOptions::new()
+            .read(true)
+            .write(true)
+            .open(INSTANCE_IMAGE_PATH)
+            .with_context(|| format!("Failed to open {}", INSTANCE_IMAGE_PATH))?;
+
+        // Check if this file is a valid instance disk by examining the header (the first block)
+        let mut magic = [0; DISK_HEADER_MAGIC.len()];
+        file.read_exact(&mut magic)?;
+        if magic != DISK_HEADER_MAGIC.as_bytes() {
+            bail!("invalid magic: {:?}", magic);
+        }
+
+        let version = file.read_u16::<LittleEndian>()?;
+        if version == 0 {
+            bail!("invalid version: {}", version);
+        }
+        if version > DISK_HEADER_VERSION {
+            bail!("unsupported version: {}", version);
+        }
+
+        Ok(Self { file })
+    }
+
+    /// Reads the identity data that was written by microdroid manager. The returned data is
+    /// plaintext, although it is stored encrypted. In case when the partition for microdroid
+    /// manager doesn't exist, which can happen if it's the first boot, `Ok(None)` is returned.
+    pub fn read_microdroid_data(&mut self) -> Result<Option<Box<[u8]>>> {
+        let (header, offset) = self.locate_microdroid_header()?;
+        if header.is_none() {
+            return Ok(None);
+        }
+        let header = header.unwrap();
+        let payload_offset = offset + PARTITION_HEADER_SIZE;
+        self.file.seek(SeekFrom::Start(payload_offset))?;
+
+        // Read the 12-bytes nonce (unencrypted)
+        let mut nonce = [0; 12];
+        self.file.read_exact(&mut nonce)?;
+        let nonce = Nonce::assume_unique_for_key(nonce);
+
+        // Read the encrypted payload
+        let payload_size = header.payload_size - 12; // we already have read the nonce
+        let mut data = vec![0; payload_size as usize];
+        self.file.read_exact(&mut data)?;
+
+        // Read the header as well because it's part of the signed data (though not encrypted).
+        let mut header = [0; PARTITION_HEADER_SIZE as usize];
+        self.file.seek(SeekFrom::Start(offset))?;
+        self.file.read_exact(&mut header)?;
+
+        // Decrypt and authenticate the data (along with the header). The data is decrypted in
+        // place. `open_in_place` returns slice to the decrypted part in the buffer.
+        let plaintext_len = get_key().open_in_place(nonce, Aad::from(&header), &mut data)?.len();
+        // Truncate to remove the tag
+        data.truncate(plaintext_len);
+
+        Ok(Some(data.into_boxed_slice()))
+    }
+
+    /// Writes identity data to the partition for microdroid manager. The partition is appended
+    /// if it doesn't exist. The data is stored encrypted.
+    pub fn write_microdroid_data(&mut self, data: &[u8]) -> Result<()> {
+        let (header, offset) = self.locate_microdroid_header()?;
+
+        // By encrypting and signing the data, tag will be appended. The tag also becomes part of
+        // the encrypted payload which will be written. In addition, a 12-bytes nonce will be
+        // prepended (non-encrypted).
+        let payload_size = (data.len() + ENCRYPT_ALG.tag_len() + 12) as u64;
+
+        // If the partition exists, make sure we don't change the partition size. If not (i.e.
+        // partition is not found), write the header at the empty place.
+        if let Some(header) = header {
+            if header.payload_size != payload_size {
+                bail!("Can't change payload size from {} to {}", header.payload_size, payload_size);
+            }
+        } else {
+            let uuid = Uuid::parse_str(MICRODROID_PARTITION_UUID)?;
+            self.write_header_at(offset, &uuid, payload_size)?;
+        }
+
+        // Read the header as it is used as additionally authenticated data (AAD).
+        let mut header = [0; PARTITION_HEADER_SIZE as usize];
+        self.file.seek(SeekFrom::Start(offset))?;
+        self.file.read_exact(&mut header)?;
+
+        // Generate a nonce randomly and recorde it on the disk first.
+        let nonce = Nonce::assume_unique_for_key(rand::random::<[u8; 12]>());
+        self.file.seek(SeekFrom::Start(offset + PARTITION_HEADER_SIZE))?;
+        self.file.write_all(nonce.as_ref())?;
+
+        // Then encrypt and sign the data. The non-encrypted input data is copied to a vector
+        // because it is encrypted in place, and also the tag is appended.
+        let mut data = data.to_vec();
+        get_key().seal_in_place_append_tag(nonce, Aad::from(&header), &mut data)?;
+
+        // Persist the encrypted payload data
+        self.file.write_all(&data)?;
+        self.file.flush()?;
+
+        Ok(())
+    }
+
+    /// Read header at `header_offset` and parse it into a `PartitionHeader`.
+    fn read_header_at(&mut self, header_offset: u64) -> Result<PartitionHeader> {
+        assert!(
+            header_offset % PARTITION_HEADER_SIZE == 0,
+            "header offset {} is not aligned to 512 bytes",
+            header_offset
+        );
+
+        let mut uuid = [0; 16];
+        self.file.seek(SeekFrom::Start(header_offset))?;
+        self.file.read_exact(&mut uuid)?;
+        let uuid = Uuid::from_bytes(uuid);
+        let payload_size = self.file.read_u64::<LittleEndian>()?;
+
+        Ok(PartitionHeader { uuid, payload_size })
+    }
+
+    /// Write header at `header_offset`
+    fn write_header_at(
+        &mut self,
+        header_offset: u64,
+        uuid: &Uuid,
+        payload_size: u64,
+    ) -> Result<()> {
+        self.file.seek(SeekFrom::Start(header_offset))?;
+        self.file.write_all(uuid.as_bytes())?;
+        self.file.write_u64::<LittleEndian>(payload_size)?;
+        Ok(())
+    }
+
+    /// Locate the header of the partition for microdroid manager. A pair of `PartitionHeader` and
+    /// the offset of the partition in the disk is returned. If the partition is not found,
+    /// `PartitionHeader` is `None` and the offset points to the empty partition that can be used
+    /// for the partition.
+    fn locate_microdroid_header(&mut self) -> Result<(Option<PartitionHeader>, PartitionOffset)> {
+        let microdroid_uuid = Uuid::parse_str(MICRODROID_PARTITION_UUID)?;
+
+        // the first partition header is located just after the disk header
+        let mut header_offset = DISK_HEADER_SIZE;
+        loop {
+            let header = self.read_header_at(header_offset)?;
+            if header.uuid == microdroid_uuid {
+                // found a matching header
+                return Ok((Some(header), header_offset));
+            } else if header.uuid == Uuid::nil() {
+                // found an empty space
+                return Ok((None, header_offset));
+            }
+            // Move to the next partition. Be careful about overflow.
+            let payload_size = round_to_multiple(header.payload_size, PARTITION_HEADER_SIZE)?;
+            let part_size = payload_size
+                .checked_add(PARTITION_HEADER_SIZE)
+                .ok_or_else(|| anyhow!("partition too large"))?;
+            header_offset = header_offset
+                .checked_add(part_size)
+                .ok_or_else(|| anyhow!("next partition at invalid offset"))?;
+        }
+    }
+}
+
+/// Round `n` up to the nearest multiple of `unit`
+fn round_to_multiple(n: u64, unit: u64) -> Result<u64> {
+    assert!((unit & (unit - 1)) == 0, "{} is not power of two", unit);
+    let ret = (n + unit - 1) & !(unit - 1);
+    if ret < n {
+        bail!("overflow")
+    }
+    Ok(ret)
+}
+
+struct ZeroOnDropKey(LessSafeKey);
+
+impl Drop for ZeroOnDropKey {
+    fn drop(&mut self) {
+        // Zeroize the key by overwriting it with a key constructed from zeros of same length
+        // This works because the raw key bytes are allocated inside the struct, not on the heap
+        let zero = [0; 32];
+        let zero_key = LessSafeKey::new(UnboundKey::new(ENCRYPT_ALG, &zero).unwrap());
+        unsafe {
+            ::std::ptr::write_volatile::<LessSafeKey>(&mut self.0, zero_key);
+        }
+    }
+}
+
+impl std::ops::Deref for ZeroOnDropKey {
+    type Target = LessSafeKey;
+    fn deref(&self) -> &LessSafeKey {
+        &self.0
+    }
+}
+
+/// Returns the key that is used to encrypt the microdroid manager partition. It is derived from
+/// the sealing CDI of the previous stage, which is Android Boot Loader (ABL).
+fn get_key() -> ZeroOnDropKey {
+    // Sealing CDI from the previous stage. For now, this is hardcoded.
+    // TODO(jiyong): actually read this from the previous stage
+    const SEALING_CDI: [u8; 32] = [10; 32];
+
+    // Derive a key from the Sealing CDI
+    // Step 1 is extraction: https://datatracker.ietf.org/doc/html/rfc5869#section-2.2 where a
+    // pseduo random key (PRK) is extracted from (Input Keying Material - IKM, which is secret) and
+    // optional salt.
+    let salt = Salt::new(HKDF_SHA256, &[]); // use 0 as salt
+    let prk = salt.extract(&SEALING_CDI); // Sealing CDI as IKM
+
+    // Step 2 is expansion: https://datatracker.ietf.org/doc/html/rfc5869#section-2.3 where the PRK
+    // (optionally with the `info` which gives contextual information) is expanded into the output
+    // keying material (OKM). Note that the process fails only when the size of OKM is longer than
+    // 255 * SHA256_HASH_SIZE (32), which isn't the case here.
+    let info = [b"microdroid_manager_key".as_ref()];
+    let okm = prk.expand(&info, HKDF_SHA256).unwrap(); // doesn't fail as explained above
+    let mut key = [0; 32];
+    okm.fill(&mut key).unwrap(); // doesn't fail as explained above
+
+    // The term LessSafe might be misleading here. LessSafe here just means that the API can
+    // possibly accept same nonces for different messages. However, since we encrypt/decrypt only a
+    // single message (the microdroid_manager partition payload) with a randomly generated nonce,
+    // this is safe enough.
+    let ret = ZeroOnDropKey(LessSafeKey::new(UnboundKey::new(ENCRYPT_ALG, &key).unwrap()));
+
+    // Don't forget to zeroize the raw key array as well
+    unsafe {
+        ::std::ptr::write_volatile::<[u8; 32]>(&mut key, [0; 32]);
+    }
+
+    ret
+}
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index ee0e797..278a14d 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -14,14 +14,17 @@
 
 //! Microdroid Manager
 
+mod instance;
 mod ioutil;
 mod metadata;
 
+use crate::instance::InstanceDisk;
 use anyhow::{anyhow, bail, Context, Result};
 use apkverify::verify;
 use binder::unstable_api::{new_spibinder, AIBinder};
 use binder::{FromIBinder, Strong};
-use log::{error, info, warn};
+use idsig::V4Signature;
+use log::{debug, error, info, warn};
 use microdroid_payload_config::{Task, TaskType, VmPayloadConfig};
 use nix::ioctl_read_bad;
 use rustutils::system_properties::PropertyWatcher;
@@ -94,6 +97,18 @@
         return Err(err);
     }
 
+    let mut instance = InstanceDisk::new()?;
+    // TODO(jiyong): the data should have an internal structure
+    if let Some(data) = instance.read_microdroid_data().context("Failed to read identity data")? {
+        debug!("read apk root hash: {}", to_hex_string(&data));
+        //TODO(jiyong) apkdmverity should use this root hash instead of the one read from the idsig
+        //file, if the root hash is found in the instance image.
+    } else {
+        let data = get_apk_roothash()?;
+        debug!("write apk root hash: {}", to_hex_string(&data));
+        instance.write_microdroid_data(data.as_ref()).context("Failed to write identity data")?;
+    }
+
     let service = get_vms_rpc_binder().expect("cannot connect to VirtualMachineService");
 
     if !metadata.payload_config_path.is_empty() {
@@ -129,6 +144,12 @@
     Ok(())
 }
 
+fn get_apk_roothash() -> Result<Box<[u8]>> {
+    let mut idsig = File::open("/dev/block/by-name/microdroid-apk-idsig")?;
+    let idsig = V4Signature::from(&mut idsig)?;
+    Ok(idsig.hashing_info.raw_root_hash)
+}
+
 fn load_config(path: &Path) -> Result<VmPayloadConfig> {
     info!("loading config from {:?}...", path);
     let file = ioutil::wait_for_file(path, WAIT_TIMEOUT)?;
@@ -141,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> {
@@ -208,3 +235,7 @@
 
     Ok(path)
 }
+
+fn to_hex_string(buf: &[u8]) -> String {
+    buf.iter().map(|b| format!("{:02X}", b)).collect()
+}
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..ccb4085 100644
--- a/vm/src/run.rs
+++ b/vm/src/run.rs
@@ -189,6 +189,16 @@
         Ok(())
     }
 
+    fn onPayloadReady(&self, _cid: i32) -> BinderResult<()> {
+        eprintln!("payload is ready");
+        Ok(())
+    }
+
+    fn onPayloadFinished(&self, _cid: i32, exit_code: i32) -> BinderResult<()> {
+        eprintln!("payload finished with exit code {}", exit_code);
+        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