Merge "Change signing API to use pre-initialized key"
diff --git a/compos/compos_key_cmd/compos_key_cmd.cpp b/compos/compos_key_cmd/compos_key_cmd.cpp
index 7ad1591..e65324a 100644
--- a/compos/compos_key_cmd/compos_key_cmd.cpp
+++ b/compos/compos_key_cmd/compos_key_cmd.cpp
@@ -54,6 +54,7 @@
 using aidl::android::system::virtualizationservice::IVirtualizationService;
 using aidl::android::system::virtualizationservice::IVirtualMachine;
 using aidl::android::system::virtualizationservice::IVirtualMachineCallback;
+using aidl::android::system::virtualizationservice::PartitionType;
 using aidl::android::system::virtualizationservice::VirtualMachineConfig;
 using aidl::com::android::compos::CompOsKeyData;
 using aidl::com::android::compos::ICompOsService;
@@ -100,7 +101,7 @@
 class Callback : public BnVirtualMachineCallback {
 public:
     ::ndk::ScopedAStatus onPayloadStarted(
-            int32_t in_cid, const ::ndk::ScopedFileDescriptor& /*in_stdout*/) override {
+            int32_t in_cid, const ::ndk::ScopedFileDescriptor& /*in_stream*/) override {
         // TODO: Consider copying stdout somewhere useful?
         LOG(INFO) << "Payload started! cid = " << in_cid;
         {
@@ -446,7 +447,8 @@
         return ErrnoError() << "Failed to create image file";
     }
 
-    auto status = service->initializeWritablePartition(fd, 10 * 1024 * 1024);
+    auto status = service->initializeWritablePartition(fd, 10 * 1024 * 1024,
+                                                       PartitionType::ANDROID_VM_INSTANCE);
     if (!status.isOk()) {
         return Error() << "Failed to initialize partition: " << status.getDescription();
     }
diff --git a/demo/java/com/android/microdroid/demo/MainActivity.java b/demo/java/com/android/microdroid/demo/MainActivity.java
index 9374f5d..044a55d 100644
--- a/demo/java/com/android/microdroid/demo/MainActivity.java
+++ b/demo/java/com/android/microdroid/demo/MainActivity.java
@@ -24,6 +24,7 @@
 import android.system.virtualmachine.VirtualMachineConfig;
 import android.system.virtualmachine.VirtualMachineException;
 import android.system.virtualmachine.VirtualMachineManager;
+import android.util.Log;
 import android.view.View;
 import android.widget.Button;
 import android.widget.CheckBox;
@@ -50,6 +51,8 @@
  * the virtual machine to the UI.
  */
 public class MainActivity extends AppCompatActivity {
+    private static final String TAG = "MicrodroidDemo";
+
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
@@ -139,19 +142,24 @@
                         new VirtualMachineCallback() {
                             @Override
                             public void onPayloadStarted(
-                                    VirtualMachine vm, ParcelFileDescriptor out) {
+                                    VirtualMachine vm, ParcelFileDescriptor stream) {
+                                if (stream == null) {
+                                    mPayloadOutput.postValue("(no output available)");
+                                    return;
+                                }
                                 try {
                                     BufferedReader reader =
                                             new BufferedReader(
                                                     new InputStreamReader(
                                                             new FileInputStream(
-                                                                    out.getFileDescriptor())));
+                                                                    stream.getFileDescriptor())));
                                     String line;
                                     while ((line = reader.readLine()) != null) {
                                         mPayloadOutput.postValue(line);
                                     }
                                 } catch (IOException e) {
-                                    // Consume
+                                    Log.e(TAG, "IOException while reading payload: "
+                                            + e.getMessage());
                                 }
                             }
 
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 3152c3b..959e355 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -29,6 +29,7 @@
 import android.system.virtualizationservice.IVirtualMachine;
 import android.system.virtualizationservice.IVirtualMachineCallback;
 import android.system.virtualizationservice.IVirtualizationService;
+import android.system.virtualizationservice.PartitionType;
 import android.system.virtualizationservice.VirtualMachineAppConfig;
 
 import java.io.File;
@@ -165,7 +166,8 @@
         try {
             service.initializeWritablePartition(
                     ParcelFileDescriptor.open(vm.mInstanceFilePath, MODE_READ_WRITE),
-                    INSTANCE_FILE_SIZE);
+                    INSTANCE_FILE_SIZE,
+                    PartitionType.ANDROID_VM_INSTANCE);
         } catch (FileNotFoundException e) {
             throw new VirtualMachineException("instance image missing", e);
         } catch (RemoteException e) {
@@ -288,7 +290,6 @@
             android.system.virtualizationservice.VirtualMachineConfig vmConfigParcel =
                     android.system.virtualizationservice.VirtualMachineConfig.appConfig(appConfig);
 
-
             mVirtualMachine = service.startVm(vmConfigParcel, mConsoleWriter);
             mVirtualMachine.registerCallback(
                     new IVirtualMachineCallback.Stub() {
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java b/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
index 07af4a1..89bb260 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
@@ -17,6 +17,7 @@
 package android.system.virtualmachine;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.os.ParcelFileDescriptor;
 
 /**
@@ -28,7 +29,7 @@
 public interface VirtualMachineCallback {
 
     /** Called when the payload starts in the VM. */
-    void onPayloadStarted(@NonNull VirtualMachine vm, @NonNull ParcelFileDescriptor stdout);
+    void onPayloadStarted(@NonNull VirtualMachine vm, @Nullable ParcelFileDescriptor stream);
 
     /** Called when the VM died. */
     void onDied(@NonNull VirtualMachine vm);
diff --git a/microdroid_manager/Android.bp b/microdroid_manager/Android.bp
index a082beb..5fae7b1 100644
--- a/microdroid_manager/Android.bp
+++ b/microdroid_manager/Android.bp
@@ -9,19 +9,27 @@
     edition: "2018",
     prefer_rlib: true,
     rustlibs: [
+        "android.system.virtualizationservice-rust",
+        "android.system.virtualmachineservice-rust",
         "libanyhow",
         "libapkverify",
+        "libbinder_rpc_unstable_bindgen",
+        "libbinder_rs",
         "libkernlog",
         "liblibc",
         "liblog_rust",
         "libmicrodroid_metadata",
         "libmicrodroid_payload_config",
+        "libnix",
         "libprotobuf",
         "librustutils",
         "libserde",
         "libserde_json",
         "libvsock",
     ],
+    shared_libs: [
+        "libbinder_rpc_unstable",
+    ],
     init_rc: ["microdroid_manager.rc"],
 }
 
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index d7e256b..ee0e797 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -19,20 +19,70 @@
 
 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 microdroid_payload_config::{Task, TaskType, VmPayloadConfig};
+use nix::ioctl_read_bad;
 use rustutils::system_properties::PropertyWatcher;
-use std::fs::{self, File};
-use std::os::unix::io::{FromRawFd, IntoRawFd};
+use std::fs::{self, File, OpenOptions};
+use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
 use std::path::Path;
 use std::process::{Command, Stdio};
 use std::str;
 use std::time::Duration;
 use vsock::VsockStream;
 
+use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::IVirtualMachineService;
+
 const WAIT_TIMEOUT: Duration = Duration::from_secs(10);
 const DM_MOUNTED_APK_PATH: &str = "/dev/block/mapper/microdroid-apk";
 
+/// The CID representing the host VM
+const VMADDR_CID_HOST: u32 = 2;
+
+/// Port number that virtualizationservice listens on connections from the guest VMs for the
+/// VirtualMachineService binder service
+/// Sync with virtualizationservice/src/aidl.rs
+const PORT_VM_BINDER_SERVICE: u32 = 5000;
+
+fn get_vms_rpc_binder() -> Result<Strong<dyn IVirtualMachineService>> {
+    // SAFETY: AIBinder returned by RpcClient has correct reference count, and the ownership can be
+    // safely taken by new_spibinder.
+    let ibinder = unsafe {
+        new_spibinder(binder_rpc_unstable_bindgen::RpcClient(
+            VMADDR_CID_HOST,
+            PORT_VM_BINDER_SERVICE,
+        ) as *mut AIBinder)
+    };
+    if let Some(ibinder) = ibinder {
+        <dyn IVirtualMachineService>::try_from(ibinder).context("Cannot connect to RPC service")
+    } else {
+        bail!("Invalid raw AIBinder")
+    }
+}
+
+const IOCTL_VM_SOCKETS_GET_LOCAL_CID: usize = 0x7b9;
+ioctl_read_bad!(
+    /// Gets local cid from /dev/vsock
+    vm_sockets_get_local_cid,
+    IOCTL_VM_SOCKETS_GET_LOCAL_CID,
+    u32
+);
+
+// TODO: remove this after VS can check the peer addresses of binder clients
+fn get_local_cid() -> Result<u32> {
+    let f = OpenOptions::new()
+        .read(true)
+        .write(false)
+        .open("/dev/vsock")
+        .context("failed to open /dev/vsock")?;
+    let mut ret = 0;
+    // SAFETY: the kernel only modifies the given u32 integer.
+    unsafe { vm_sockets_get_local_cid(f.as_raw_fd(), &mut ret) }?;
+    Ok(ret)
+}
+
 fn main() -> Result<()> {
     kernlog::init()?;
     info!("started.");
@@ -44,6 +94,8 @@
         return Err(err);
     }
 
+    let service = get_vms_rpc_binder().expect("cannot connect to VirtualMachineService");
+
     if !metadata.payload_config_path.is_empty() {
         let config = load_config(Path::new(&metadata.payload_config_path))?;
 
@@ -54,7 +106,7 @@
 
         // TODO(jooyung): wait until sys.boot_completed?
         if let Some(main_task) = &config.task {
-            exec_task(main_task).map_err(|e| {
+            exec_task(main_task, &service).map_err(|e| {
                 error!("failed to execute task: {}", e);
                 e
             })?;
@@ -85,29 +137,13 @@
 
 /// Executes the given task. Stdout of the task is piped into the vsock stream to the
 /// virtualizationservice in the host side.
-fn exec_task(task: &Task) -> Result<()> {
-    const VMADDR_CID_HOST: u32 = 2;
-    const PORT_VIRT_SVC: u32 = 3000;
-    let stdout = match VsockStream::connect_with_cid_port(VMADDR_CID_HOST, PORT_VIRT_SVC) {
-        Ok(stream) => {
-            // SAFETY: the ownership of the underlying file descriptor is transferred from stream
-            // to the file object, and then into the Command object. When the command is finished,
-            // the file descriptor is closed.
-            let f = unsafe { File::from_raw_fd(stream.into_raw_fd()) };
-            Stdio::from(f)
-        }
-        Err(e) => {
-            error!("failed to connect to virtualization service: {}", e);
-            // Don't fail hard here. Even if we failed to connect to the virtualizationservice,
-            // we keep executing the task. This can happen if the owner of the VM doesn't register
-            // callback to accept the stream. Use /dev/null as the stdout so that the task can
-            // make progress without waiting for someone to consume the output.
-            Stdio::null()
-        }
-    };
+fn exec_task(task: &Task, service: &Strong<dyn IVirtualMachineService>) -> Result<()> {
     info!("executing main task {:?}...", task);
-    // TODO(jiyong): consider piping the stream into stdio (and probably stderr) as well.
-    let mut child = build_command(task)?.stdout(stdout).spawn()?;
+    let mut child = build_command(task)?.spawn()?;
+
+    info!("notifying payload started");
+    service.notifyPayloadStarted(get_local_cid()? as i32)?;
+
     match child.wait()?.code() {
         Some(0) => {
             info!("task successfully finished");
@@ -119,7 +155,10 @@
 }
 
 fn build_command(task: &Task) -> Result<Command> {
-    Ok(match task.type_ {
+    const VMADDR_CID_HOST: u32 = 2;
+    const PORT_VIRT_SVC: u32 = 3000;
+
+    let mut command = match task.type_ {
         TaskType::Executable => {
             let mut command = Command::new(&task.command);
             command.args(&task.args);
@@ -130,7 +169,30 @@
             command.arg(find_library_path(&task.command)?).args(&task.args);
             command
         }
-    })
+    };
+
+    match VsockStream::connect_with_cid_port(VMADDR_CID_HOST, PORT_VIRT_SVC) {
+        Ok(stream) => {
+            // SAFETY: the ownership of the underlying file descriptor is transferred from stream
+            // to the file object, and then into the Command object. When the command is finished,
+            // the file descriptor is closed.
+            let file = unsafe { File::from_raw_fd(stream.into_raw_fd()) };
+            command
+                .stdin(Stdio::from(file.try_clone()?))
+                .stdout(Stdio::from(file.try_clone()?))
+                .stderr(Stdio::from(file));
+        }
+        Err(e) => {
+            error!("failed to connect to virtualization service: {}", e);
+            // Don't fail hard here. Even if we failed to connect to the virtualizationservice,
+            // we keep executing the task. This can happen if the owner of the VM doesn't register
+            // callback to accept the stream. Use /dev/null as the stream so that the task can
+            // make progress without waiting for someone to consume the output.
+            command.stdin(Stdio::null()).stdout(Stdio::null()).stderr(Stdio::null());
+        }
+    }
+
+    Ok(command)
 }
 
 fn find_library_path(name: &str) -> Result<String> {
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index c3b36ee..8b9d0fa 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -21,9 +21,12 @@
     prefer_rlib: true,
     rustlibs: [
         "android.system.virtualizationservice-rust",
+        "android.system.virtualmachineservice-rust",
         "android.os.permissions_aidl-rust",
         "libandroid_logger",
         "libanyhow",
+        "libbinder_rpc_unstable_bindgen",
+        "libbinder_rs",
         "libcommand_fds",
         "libdisk",
         "libidsig",
@@ -39,6 +42,9 @@
         "libzip",
         "libvsock",
     ],
+    shared_libs: [
+        "libbinder_rpc_unstable",
+    ],
 }
 
 rust_binary {
diff --git a/virtualizationservice/aidl/Android.bp b/virtualizationservice/aidl/Android.bp
index 7d85bd3..974bdc6 100644
--- a/virtualizationservice/aidl/Android.bp
+++ b/virtualizationservice/aidl/Android.bp
@@ -4,7 +4,7 @@
 
 aidl_interface {
     name: "android.system.virtualizationservice",
-    srcs: ["**/*.aidl"],
+    srcs: ["android/system/virtualizationservice/**/*.aidl"],
     // This is never accessed directly. Apps are expected to use this indirectly via the Java
     // wrapper android.system.virtualmachine.
     unstable: true,
@@ -27,3 +27,15 @@
         },
     },
 }
+
+aidl_interface {
+    name: "android.system.virtualmachineservice",
+    srcs: ["android/system/virtualmachineservice/**/*.aidl"],
+    unstable: true,
+    backend: {
+        rust: {
+            enabled: true,
+            apex_available: ["com.android.virt"],
+        },
+    },
+}
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
index 7bb18a4..c7a1471 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
@@ -23,14 +23,13 @@
  */
 oneway interface IVirtualMachineCallback {
     /**
-     * Called when the payload starts in the VM. `stdout` is the stdout of the payload.
+     * Called when the payload starts in the VM. `stream` is the input/output port of the payload.
      *
      * <p>Note: when the virtual machine object is shared to multiple processes and they register
-     * this callback to the same virtual machine object, the processes will compete to read from the
-     * same payload stdout. As a result, each process might get only a part of the entire output
-     * stream. To avoid such a case, keep only one process to read from the stdout.
+     * this callback to the same virtual machine object, the processes will compete to access the
+     * same payload stream. Keep only one process to access the stream.
      */
-    void onPayloadStarted(int cid, in ParcelFileDescriptor stdout);
+    void onPayloadStarted(int cid, in @nullable ParcelFileDescriptor stream);
 
     /**
      * Called when the VM dies.
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
index d136465..7c4b897 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
@@ -16,6 +16,7 @@
 package android.system.virtualizationservice;
 
 import android.system.virtualizationservice.IVirtualMachine;
+import android.system.virtualizationservice.PartitionType;
 import android.system.virtualizationservice.VirtualMachineConfig;
 import android.system.virtualizationservice.VirtualMachineDebugInfo;
 
@@ -32,7 +33,8 @@
      *
      * The file must be open with both read and write permissions, and should be a new empty file.
      */
-    void initializeWritablePartition(in ParcelFileDescriptor imageFd, long size);
+    void initializeWritablePartition(
+            in ParcelFileDescriptor imageFd, long size, PartitionType type);
 
     /**
      * Create or update an idsig file that digests the given APK file. The idsig file follows the
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/PartitionType.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/PartitionType.aidl
new file mode 100644
index 0000000..f25e674
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/PartitionType.aidl
@@ -0,0 +1,32 @@
+/*
+ * 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.
+ */
+package android.system.virtualizationservice;
+
+/**
+ * Type of the writable partition that virtualizationservice creates via
+ * initializeWritablePartition.
+ */
+@Backing(type="int")
+enum PartitionType {
+    /**
+     * The partition is simply initialized as all zeros
+     */
+    RAW = 0,
+    /**
+     * The partition is initialized as an instance image which is formatted to hold per-VM secrets
+     */
+    ANDROID_VM_INSTANCE = 1,
+}
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
new file mode 100644
index 0000000..b5cda7d
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
@@ -0,0 +1,25 @@
+/*
+ * 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.
+ */
+package android.system.virtualmachineservice;
+
+/** {@hide} */
+interface IVirtualMachineService {
+    /**
+     * Notifies that the virtual machine is ready.
+     * TODO(b/191845268): remove cid parameter
+     */
+    void notifyPayloadStarted(int cid);
+}
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 23a9c03..c5d3548 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -20,8 +20,8 @@
 use crate::{Cid, FIRST_GUEST_CID};
 
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualizationService::IVirtualizationService;
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::DiskImage::DiskImage;
+use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualizationService::IVirtualizationService;
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualMachine::{
     BnVirtualMachine, IVirtualMachine,
 };
@@ -32,10 +32,15 @@
     VirtualMachineRawConfig::VirtualMachineRawConfig,
 };
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::VirtualMachineDebugInfo::VirtualMachineDebugInfo;
+use android_system_virtualizationservice::aidl::android::system::virtualizationservice::PartitionType::PartitionType;
 use android_system_virtualizationservice::binder::{
     self, BinderFeatures, ExceptionCode, Interface, ParcelFileDescriptor, Status, Strong, ThreadState,
 };
+use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::{
+    BnVirtualMachineService, IVirtualMachineService,
+};
 use anyhow::{bail, Context, Result};
+use ::binder::unstable_api::AsNative;
 use disk::QcowFile;
 use idsig::{V4Signature, HashAlgorithm};
 use log::{debug, error, warn, info};
@@ -43,12 +48,13 @@
 use std::convert::TryInto;
 use std::ffi::CString;
 use std::fs::{File, OpenOptions, create_dir};
+use std::io::{Error, ErrorKind, Write};
 use std::num::NonZeroU32;
 use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex, Weak};
 use vmconfig::VmConfig;
-use vsock::{VsockListener, SockAddr, VsockStream};
+use vsock::{SockAddr, VsockListener, VsockStream};
 use zip::ZipArchive;
 
 pub const BINDER_SERVICE_IDENTIFIER: &str = "android.system.virtualizationservice";
@@ -60,13 +66,24 @@
 const VMADDR_CID_HOST: u32 = 2;
 
 /// Port number that virtualizationservice listens on connections from the guest VMs for the
-/// payload output
-const PORT_VIRT_SERVICE: u32 = 3000;
+/// payload input and output
+const PORT_VIRT_STREAM_SERVICE: u32 = 3000;
+
+/// Port number that virtualizationservice listens on connections from the guest VMs for the
+/// VirtualMachineService binder service
+/// Sync with microdroid_manager/src/main.rs
+const PORT_VM_BINDER_SERVICE: u32 = 5000;
 
 /// The size of zero.img.
 /// Gaps in composite disk images are filled with a shared zero.img.
 const ZERO_FILLER_SIZE: u64 = 4096;
 
+/// Magic string for the instance image
+const ANDROID_VM_INSTANCE_MAGIC: &str = "Android-VM-instance";
+
+/// Version of the instance image format
+const ANDROID_VM_INSTANCE_VERSION: u16 = 1;
+
 /// Implementation of `IVirtualizationService`, the entry point of the AIDL service.
 #[derive(Debug, Default)]
 pub struct VirtualizationService {
@@ -190,6 +207,7 @@
         &self,
         image_fd: &ParcelFileDescriptor,
         size: i64,
+        partition_type: PartitionType,
     ) -> binder::Result<()> {
         check_manage_access()?;
         let size = size.try_into().map_err(|e| {
@@ -200,13 +218,28 @@
         })?;
         let image = clone_file(image_fd)?;
 
-        QcowFile::new(image, size).map_err(|e| {
+        let mut part = QcowFile::new(image, size).map_err(|e| {
             new_binder_exception(
                 ExceptionCode::SERVICE_SPECIFIC,
                 format!("Failed to create QCOW2 image: {}", e),
             )
         })?;
 
+        match partition_type {
+            PartitionType::RAW => Ok(()),
+            PartitionType::ANDROID_VM_INSTANCE => format_as_android_vm_instance(&mut part),
+            _ => Err(Error::new(
+                ErrorKind::Unsupported,
+                format!("Unsupported partition type {:?}", partition_type),
+            )),
+        }
+        .map_err(|e| {
+            new_binder_exception(
+                ExceptionCode::SERVICE_SPECIFIC,
+                format!("Failed to initialize partition as {:?}: {}", partition_type, e),
+            )
+        })?;
+
         Ok(())
     }
 
@@ -273,9 +306,33 @@
 impl VirtualizationService {
     pub fn init() -> VirtualizationService {
         let service = VirtualizationService::default();
+
+        // server for payload output
         let state = service.state.clone(); // reference to state (not the state itself) is copied
         std::thread::spawn(move || {
-            handle_connection_from_vm(state).unwrap();
+            handle_stream_connection_from_vm(state).unwrap();
+        });
+
+        // binder server for vm
+        let state = service.state.clone(); // reference to state (not the state itself) is copied
+        std::thread::spawn(move || {
+            let mut service = VirtualMachineService::new_binder(state).as_binder();
+            debug!("virtual machine service is starting as an RPC service.");
+            // SAFETY: Service ownership is transferring to the server and won't be valid afterward.
+            // Plus the binder objects are threadsafe.
+            let retval = unsafe {
+                binder_rpc_unstable_bindgen::RunRpcServer(
+                    service.as_native_mut() as *mut binder_rpc_unstable_bindgen::AIBinder,
+                    PORT_VM_BINDER_SERVICE,
+                )
+            };
+            if retval {
+                debug!("RPC server has shut down gracefully");
+            } else {
+                bail!("Premature termination of RPC server");
+            }
+
+            Ok(retval)
         });
         service
     }
@@ -283,8 +340,8 @@
 
 /// Waits for incoming connections from VM. If a new connection is made, notify the event to the
 /// client via the callback (if registered).
-fn handle_connection_from_vm(state: Arc<Mutex<State>>) -> Result<()> {
-    let listener = VsockListener::bind_with_cid_port(VMADDR_CID_HOST, PORT_VIRT_SERVICE)?;
+fn handle_stream_connection_from_vm(state: Arc<Mutex<State>>) -> Result<()> {
+    let listener = VsockListener::bind_with_cid_port(VMADDR_CID_HOST, PORT_VIRT_STREAM_SERVICE)?;
     for stream in listener.incoming() {
         let stream = match stream {
             Err(e) => {
@@ -296,13 +353,11 @@
         if let Ok(SockAddr::Vsock(addr)) = stream.peer_addr() {
             let cid = addr.cid();
             let port = addr.port();
-            info!("connected from cid={}, port={}", cid, port);
-            if cid < FIRST_GUEST_CID {
-                warn!("connection is not from a guest VM");
-                continue;
-            }
+            info!("payload stream connected from cid={}, port={}", cid, port);
             if let Some(vm) = state.lock().unwrap().get_vm(cid) {
-                vm.callbacks.notify_payload_started(cid, stream);
+                vm.stream.lock().unwrap().insert(stream);
+            } else {
+                error!("connection from cid={} is not from a guest VM", cid);
             }
         }
     }
@@ -320,6 +375,12 @@
     Ok(())
 }
 
+fn format_as_android_vm_instance(part: &mut dyn Write) -> std::io::Result<()> {
+    part.write_all(ANDROID_VM_INSTANCE_MAGIC.as_bytes())?;
+    part.write_all(&ANDROID_VM_INSTANCE_VERSION.to_le_bytes())?;
+    part.flush()
+}
+
 /// Given the configuration for a disk image, assembles the `DiskFile` to pass to crosvm.
 ///
 /// This may involve assembling a composite disk from a set of partition images.
@@ -569,11 +630,11 @@
 
 impl VirtualMachineCallbacks {
     /// Call all registered callbacks to notify that the payload has started.
-    pub fn notify_payload_started(&self, cid: Cid, stream: VsockStream) {
+    pub fn notify_payload_started(&self, cid: Cid, stream: Option<VsockStream>) {
         let callbacks = &*self.0.lock().unwrap();
-        let pfd = vsock_stream_to_pfd(stream);
+        let pfd = stream.map(vsock_stream_to_pfd);
         for callback in callbacks {
-            if let Err(e) = callback.onPayloadStarted(cid as i32, &pfd) {
+            if let Err(e) = callback.onPayloadStarted(cid as i32, pfd.as_ref()) {
                 error!("Error notifying payload start event from VM CID {}: {}", cid, e);
             }
         }
@@ -702,3 +763,38 @@
         }
     }
 }
+
+/// Implementation of `IVirtualMachineService`, the entry point of the AIDL service.
+#[derive(Debug, Default)]
+struct VirtualMachineService {
+    state: Arc<Mutex<State>>,
+}
+
+impl Interface for VirtualMachineService {}
+
+impl IVirtualMachineService for VirtualMachineService {
+    fn notifyPayloadStarted(&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 {} started payload", cid);
+            let stream = vm.stream.lock().unwrap().take();
+            vm.callbacks.notify_payload_started(cid, stream);
+            Ok(())
+        } else {
+            error!("notifyPayloadStarted is called from an unknown cid {}", cid);
+            Err(new_binder_exception(
+                ExceptionCode::SERVICE_SPECIFIC,
+                format!("cannot find a VM with cid {}", cid),
+            ))
+        }
+    }
+}
+
+impl VirtualMachineService {
+    fn new_binder(state: Arc<Mutex<State>>) -> Strong<dyn IVirtualMachineService> {
+        BnVirtualMachineService::new_binder(
+            VirtualMachineService { state },
+            BinderFeatures::default(),
+        )
+    }
+}
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index 5873cd9..5984ff0 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -26,8 +26,9 @@
 use std::path::PathBuf;
 use std::process::Command;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use std::thread;
+use vsock::VsockStream;
 
 const CROSVM_PATH: &str = "/apex/com.android.virt/bin/crosvm";
 
@@ -73,6 +74,8 @@
     running: AtomicBool,
     /// Callbacks to clients of the VM.
     pub callbacks: VirtualMachineCallbacks,
+    /// Input/output stream of the payload run in the VM.
+    pub stream: Mutex<Option<VsockStream>>,
 }
 
 impl VmInstance {
@@ -96,6 +99,7 @@
             requester_debug_pid,
             running: AtomicBool::new(true),
             callbacks: Default::default(),
+            stream: Mutex::new(None),
         }
     }
 
diff --git a/vm/src/create_partition.rs b/vm/src/create_partition.rs
index acebbf2..22f7bea 100644
--- a/vm/src/create_partition.rs
+++ b/vm/src/create_partition.rs
@@ -15,6 +15,7 @@
 //! Command to create an empty partition
 
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualizationService::IVirtualizationService;
+use android_system_virtualizationservice::aidl::android::system::virtualizationservice::PartitionType::PartitionType;
 use android_system_virtualizationservice::binder::{ParcelFileDescriptor, Strong};
 use anyhow::{Context, Error};
 use std::convert::TryInto;
@@ -26,6 +27,7 @@
     service: Strong<dyn IVirtualizationService>,
     image_path: &Path,
     size: u64,
+    partition_type: PartitionType,
 ) -> Result<(), Error> {
     let image = OpenOptions::new()
         .create_new(true)
@@ -34,7 +36,14 @@
         .open(image_path)
         .with_context(|| format!("Failed to create {:?}", image_path))?;
     service
-        .initializeWritablePartition(&ParcelFileDescriptor::new(image), size.try_into()?)
-        .context("Failed to initialize partition with size {}, size")?;
+        .initializeWritablePartition(
+            &ParcelFileDescriptor::new(image),
+            size.try_into()?,
+            partition_type,
+        )
+        .context(format!(
+            "Failed to initialize partition type: {:?}, size: {}",
+            partition_type, size
+        ))?;
     Ok(())
 }
diff --git a/vm/src/main.rs b/vm/src/main.rs
index 09f11d5..fe47d2c 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -19,6 +19,7 @@
 mod sync;
 
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualizationService::IVirtualizationService;
+use android_system_virtualizationservice::aidl::android::system::virtualizationservice::PartitionType::PartitionType;
 use android_system_virtualizationservice::binder::{wait_for_interface, ProcessState, Strong};
 use anyhow::{Context, Error};
 use create_partition::command_create_partition;
@@ -91,9 +92,21 @@
 
         /// The desired size of the partition, in bytes.
         size: u64,
+
+        /// Type of the partition
+        #[structopt(short="t", long="type", default_value="raw", parse(try_from_str=parse_partition_type))]
+        partition_type: PartitionType,
     },
 }
 
+fn parse_partition_type(s: &str) -> Result<PartitionType, String> {
+    match s {
+        "raw" => Ok(PartitionType::RAW),
+        "instance" => Ok(PartitionType::ANDROID_VM_INSTANCE),
+        _ => Err(format!("Invalid partition type {}", s)),
+    }
+}
+
 fn main() -> Result<(), Error> {
     env_logger::init();
     let opt = Opt::from_args();
@@ -122,7 +135,9 @@
         }
         Opt::Stop { cid } => command_stop(service, cid),
         Opt::List => command_list(service),
-        Opt::CreatePartition { path, size } => command_create_partition(service, &path, size),
+        Opt::CreatePartition { path, size, partition_type } => {
+            command_create_partition(service, &path, size, partition_type)
+        }
     }
 }
 
diff --git a/vm/src/run.rs b/vm/src/run.rs
index 8db43fb..3b34bac 100644
--- a/vm/src/run.rs
+++ b/vm/src/run.rs
@@ -22,6 +22,7 @@
     BnVirtualMachineCallback, IVirtualMachineCallback,
 };
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
+    PartitionType::PartitionType,
     VirtualMachineAppConfig::VirtualMachineAppConfig,
     VirtualMachineConfig::VirtualMachineConfig,
 };
@@ -60,7 +61,12 @@
 
     if !instance.exists() {
         const INSTANCE_FILE_SIZE: u64 = 10 * 1024 * 1024;
-        command_create_partition(service.clone(), instance, INSTANCE_FILE_SIZE)?;
+        command_create_partition(
+            service.clone(),
+            instance,
+            INSTANCE_FILE_SIZE,
+            PartitionType::ANDROID_VM_INSTANCE,
+        )?;
     }
 
     let config = VirtualMachineConfig::AppConfig(VirtualMachineAppConfig {
@@ -163,16 +169,22 @@
 impl Interface for VirtualMachineCallback {}
 
 impl IVirtualMachineCallback for VirtualMachineCallback {
-    fn onPayloadStarted(&self, _cid: i32, stdout: &ParcelFileDescriptor) -> BinderResult<()> {
-        // Show the stdout of the payload
-        let mut reader = BufReader::new(stdout.as_ref());
-        loop {
-            let mut s = String::new();
-            match reader.read_line(&mut s) {
-                Ok(0) => break,
-                Ok(_) => print!("{}", s),
-                Err(e) => eprintln!("error reading from virtual machine: {}", e),
-            };
+    fn onPayloadStarted(
+        &self,
+        _cid: i32,
+        stream: Option<&ParcelFileDescriptor>,
+    ) -> BinderResult<()> {
+        // Show the output of the payload
+        if let Some(stream) = stream {
+            let mut reader = BufReader::new(stream.as_ref());
+            loop {
+                let mut s = String::new();
+                match reader.read_line(&mut s) {
+                    Ok(0) => break,
+                    Ok(_) => print!("{}", s),
+                    Err(e) => eprintln!("error reading from virtual machine: {}", e),
+                };
+            }
         }
         Ok(())
     }