[service-vm] Persist the service VM instance image in VS

This cl persists the instance image of the service VM across
different VMs and manages it within the target
virtualizationservice.

Bug: 278858244
Test: Runs the ServiceVmClientApp in VM
Test: atest MicrodroidHostTests
Change-Id: Ic0a2205bae236a933d3ddd807bd124ebaaa18f8d
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 97151d7..f5f2718 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -380,8 +380,8 @@
 
         // Check if partition images are labeled incorrectly. This is to prevent random images
         // which are not protected by the Android Verified Boot (e.g. bits downloaded by apps) from
-        // being loaded in a pVM. This applies to everything in the raw config, and everything but
-        // the non-executable, generated partitions in the app config.
+        // being loaded in a pVM. This applies to everything but the instance image in the raw config,
+        // and everything but the non-executable, generated partitions in the app config.
         config
             .disks
             .iter()
@@ -390,7 +390,7 @@
                 if is_app_config {
                     !is_safe_app_partition(&partition.label)
                 } else {
-                    true // all partitions are checked
+                    !is_safe_raw_partition(&partition.label)
                 }
             })
             .try_for_each(check_label_for_partition)
@@ -769,6 +769,11 @@
         || label.starts_with("extra-idsig-")
 }
 
+/// Returns whether a partition with the given label is safe for a raw config VM.
+fn is_safe_raw_partition(label: &str) -> bool {
+    label == "vm-instance"
+}
+
 /// Check that a file SELinux label is acceptable.
 ///
 /// We only want to allow code in a VM to be sourced from places that apps, and the
@@ -1214,22 +1219,7 @@
     }
 
     fn requestCertificate(&self, csr: &[u8]) -> binder::Result<Vec<u8>> {
-        let cid = self.cid;
-        let Some(vm) = self.state.lock().unwrap().get_vm(cid) else {
-            error!("requestCertificate is called from an unknown CID {cid}");
-            return Err(anyhow!("cannot find a VM with CID {}", cid))
-                .or_service_specific_exception(-1);
-        };
-        let instance_img_path = vm.temporary_directory.join("rkpvm_instance.img");
-        let instance_img = OpenOptions::new()
-            .create(true)
-            .read(true)
-            .write(true)
-            .open(instance_img_path)
-            .context("Failed to create rkpvm_instance.img file")
-            .with_log()
-            .or_service_specific_exception(-1)?;
-        GLOBAL_SERVICE.requestCertificate(csr, &ParcelFileDescriptor::new(instance_img))
+        GLOBAL_SERVICE.requestCertificate(csr)
     }
 }
 
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
index 4c7164a..9d698ea 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
@@ -55,11 +55,9 @@
      * Requests a certificate using the provided certificate signing request (CSR).
      *
      * @param csr the certificate signing request.
-     * @param instanceImgFd The file descriptor of the instance image. The file should be open for
-     *         both reading and writing.
      * @return the X.509 encoded certificate.
      */
-    byte[] requestCertificate(in byte[] csr, in ParcelFileDescriptor instanceImgFd);
+    byte[] requestCertificate(in byte[] csr);
 
     /**
      * Get a list of assignable devices.
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index b2513d9..4c97ad4 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -156,14 +156,10 @@
         Ok(cids)
     }
 
-    fn requestCertificate(
-        &self,
-        csr: &[u8],
-        instance_img_fd: &ParcelFileDescriptor,
-    ) -> binder::Result<Vec<u8>> {
+    fn requestCertificate(&self, csr: &[u8]) -> binder::Result<Vec<u8>> {
         check_manage_access()?;
         info!("Received csr. Getting certificate...");
-        request_certificate(csr, instance_img_fd)
+        request_certificate(csr)
             .context("Failed to get certificate")
             .with_log()
             .or_service_specific_exception(-1)
diff --git a/virtualizationservice/src/main.rs b/virtualizationservice/src/main.rs
index bf8b944..3af0d42 100644
--- a/virtualizationservice/src/main.rs
+++ b/virtualizationservice/src/main.rs
@@ -17,6 +17,7 @@
 mod aidl;
 mod atom;
 mod rkpvm;
+mod service_vm;
 
 use crate::aidl::{
     remove_temporary_dir, BINDER_SERVICE_IDENTIFIER, TEMPORARY_DIRECTORY,
@@ -55,6 +56,8 @@
 
     clear_temporary_files().expect("Failed to delete old temporary files");
 
+    ProcessState::start_thread_pool();
+
     let service = VirtualizationServiceInternal::init();
     let service = BnVirtualizationServiceInternal::new_binder(service, BinderFeatures::default());
     register_lazy_service(BINDER_SERVICE_IDENTIFIER, service.as_binder()).unwrap();
diff --git a/virtualizationservice/src/rkpvm.rs b/virtualizationservice/src/rkpvm.rs
index 63160f4..bb05edd 100644
--- a/virtualizationservice/src/rkpvm.rs
+++ b/virtualizationservice/src/rkpvm.rs
@@ -16,74 +16,13 @@
 //! The RKP VM will be recognized and attested by the RKP server periodically and
 //! serves as a trusted platform to attest a client VM.
 
-use android_system_virtualizationservice::{
-    aidl::android::system::virtualizationservice::{
-        CpuTopology::CpuTopology, DiskImage::DiskImage, Partition::Partition,
-        PartitionType::PartitionType, VirtualMachineConfig::VirtualMachineConfig,
-        VirtualMachineRawConfig::VirtualMachineRawConfig,
-    },
-    binder::{ParcelFileDescriptor, ProcessState},
-};
-use anyhow::{anyhow, Context, Result};
+use crate::service_vm;
+use anyhow::{anyhow, Result};
 use log::info;
-use std::fs::File;
 use std::time::Duration;
-use vmclient::VmInstance;
 
-const RIALTO_PATH: &str = "/apex/com.android.virt/etc/rialto.bin";
-
-pub(crate) fn request_certificate(
-    csr: &[u8],
-    instance_img_fd: &ParcelFileDescriptor,
-) -> Result<Vec<u8>> {
-    // We need to start the thread pool for Binder to work properly, especially link_to_death.
-    ProcessState::start_thread_pool();
-
-    let virtmgr = vmclient::VirtualizationService::new().context("Failed to spawn virtmgr")?;
-    let service = virtmgr.connect().context("virtmgr failed to connect")?;
-    info!("service_vm: Connected to VirtualizationService");
-    // TODO(b/272226230): Either turn rialto into the service VM or use an empty payload here.
-    // If using an empty payload, the service code will be part of pvmfw.
-    let rialto = File::open(RIALTO_PATH).context("Failed to open Rialto kernel binary")?;
-
-    // TODO(b/272226230): Initialize the partition from virtualization manager.
-    const INSTANCE_IMG_SIZE_BYTES: i64 = 1 << 20; // 1MB
-    service
-        .initializeWritablePartition(
-            instance_img_fd,
-            INSTANCE_IMG_SIZE_BYTES,
-            PartitionType::ANDROID_VM_INSTANCE,
-        )
-        .context("Failed to initialize instange.img")?;
-    let instance_img =
-        instance_img_fd.as_ref().try_clone().context("Failed to clone instance.img")?;
-    let instance_img = ParcelFileDescriptor::new(instance_img);
-    let writable_partitions = vec![Partition {
-        label: "vm-instance".to_owned(),
-        image: Some(instance_img),
-        writable: true,
-    }];
-    info!("service_vm: Finished initializing instance.img...");
-
-    let config = VirtualMachineConfig::RawConfig(VirtualMachineRawConfig {
-        name: String::from("Service VM"),
-        kernel: None,
-        initrd: None,
-        params: None,
-        bootloader: Some(ParcelFileDescriptor::new(rialto)),
-        disks: vec![DiskImage { image: None, partitions: writable_partitions, writable: true }],
-        protectedVm: true,
-        memoryMib: 300,
-        cpuTopology: CpuTopology::ONE_CPU,
-        platformVersion: "~1.0".to_string(),
-        gdbPort: 0, // No gdb
-        ..Default::default()
-    });
-    let vm = VmInstance::create(service.as_ref(), &config, None, None, None, None)
-        .context("Failed to create service VM")?;
-
-    info!("service_vm: Starting Service VM...");
-    vm.start().context("Failed to start service VM")?;
+pub(crate) fn request_certificate(csr: &[u8]) -> Result<Vec<u8>> {
+    let vm = service_vm::start()?;
 
     // TODO(b/274441673): The host can send the CSR to the RKP VM for attestation.
     // Wait for VM to finish.
diff --git a/virtualizationservice/src/service_vm.rs b/virtualizationservice/src/service_vm.rs
new file mode 100644
index 0000000..5a1744f
--- /dev/null
+++ b/virtualizationservice/src/service_vm.rs
@@ -0,0 +1,102 @@
+// Copyright 2023, 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.
+
+//! Service VM.
+
+use android_system_virtualizationservice::{
+    aidl::android::system::virtualizationservice::{
+        CpuTopology::CpuTopology, DiskImage::DiskImage,
+        IVirtualizationService::IVirtualizationService, Partition::Partition,
+        PartitionType::PartitionType, VirtualMachineConfig::VirtualMachineConfig,
+        VirtualMachineRawConfig::VirtualMachineRawConfig,
+    },
+    binder::ParcelFileDescriptor,
+};
+use anyhow::{Context, Result};
+use log::info;
+use std::fs::{File, OpenOptions};
+use std::path::Path;
+use vmclient::VmInstance;
+
+const VIRT_DATA_DIR: &str = "/data/misc/apexdata/com.android.virt";
+const RIALTO_PATH: &str = "/apex/com.android.virt/etc/rialto.bin";
+const INSTANCE_IMG_NAME: &str = "service_vm_instance.img";
+const INSTANCE_IMG_SIZE_BYTES: i64 = 1 << 20; // 1MB
+const MEMORY_MB: i32 = 300;
+
+/// Starts the service VM and returns its instance.
+/// The same instance image is used for different VMs.
+/// TODO(b/278858244): Allow only one service VM running at each time.
+pub fn start() -> Result<VmInstance> {
+    let virtmgr = vmclient::VirtualizationService::new().context("Failed to spawn VirtMgr")?;
+    let service = virtmgr.connect().context("Failed to connect to VirtMgr")?;
+    info!("Connected to VirtMgr for service VM");
+
+    let vm = vm_instance(service.as_ref())?;
+
+    vm.start().context("Failed to start service VM")?;
+    info!("Service VM started");
+    Ok(vm)
+}
+
+fn vm_instance(service: &dyn IVirtualizationService) -> Result<VmInstance> {
+    let instance_img = instance_img(service)?;
+    let writable_partitions = vec![Partition {
+        label: "vm-instance".to_owned(),
+        image: Some(instance_img),
+        writable: true,
+    }];
+    let rialto = File::open(RIALTO_PATH).context("Failed to open Rialto kernel binary")?;
+    let config = VirtualMachineConfig::RawConfig(VirtualMachineRawConfig {
+        name: String::from("Service VM"),
+        bootloader: Some(ParcelFileDescriptor::new(rialto)),
+        disks: vec![DiskImage { image: None, partitions: writable_partitions, writable: true }],
+        protectedVm: true,
+        memoryMib: MEMORY_MB,
+        cpuTopology: CpuTopology::ONE_CPU,
+        platformVersion: "~1.0".to_string(),
+        gdbPort: 0, // No gdb
+        ..Default::default()
+    });
+    let console_out = None;
+    let console_in = None;
+    let log = None;
+    let callback = None;
+    VmInstance::create(service, &config, console_out, console_in, log, callback)
+        .context("Failed to create service VM")
+}
+
+fn instance_img(service: &dyn IVirtualizationService) -> Result<ParcelFileDescriptor> {
+    let instance_img_path = Path::new(VIRT_DATA_DIR).join(INSTANCE_IMG_NAME);
+    if instance_img_path.exists() {
+        // TODO(b/298174584): Try to recover if the service VM is triggered by rkpd.
+        return Ok(OpenOptions::new()
+            .read(true)
+            .write(true)
+            .open(instance_img_path)
+            .map(ParcelFileDescriptor::new)?);
+    }
+    let instance_img = OpenOptions::new()
+        .create(true)
+        .read(true)
+        .write(true)
+        .open(instance_img_path)
+        .map(ParcelFileDescriptor::new)?;
+    service.initializeWritablePartition(
+        &instance_img,
+        INSTANCE_IMG_SIZE_BYTES,
+        PartitionType::ANDROID_VM_INSTANCE,
+    )?;
+    Ok(instance_img)
+}