Merge "[rialto] Enable Rialto to run in protected VM"
diff --git a/rialto/Android.bp b/rialto/Android.bp
index c2a19f3..5034bf4 100644
--- a/rialto/Android.bp
+++ b/rialto/Android.bp
@@ -11,7 +11,9 @@
     rustlibs: [
         "libaarch64_paging",
         "libbuddy_system_allocator",
+        "libhyp",
         "liblog_rust_nostd",
+        "libsmccc",
         "libvmbase",
     ],
     apex_available: ["com.android.virt"],
@@ -36,7 +38,7 @@
 }
 
 raw_binary {
-    name: "rialto",
+    name: "rialto_unsigned",
     src: ":rialto_elf",
     enabled: false,
     target: {
@@ -46,6 +48,42 @@
     },
 }
 
+// python -c "import hashlib; print(hashlib.sha256(b'rialto_salt').hexdigest())"
+rialto_salt = "ea9d8c3ae1785396884d0c16c7652921874e2b8703f336ff23760f2049ee9e29"
+
+filegroup {
+    name: "rialto_sign_key",
+    srcs: [":avb_testkey_rsa4096"],
+}
+
+avb_add_hash_footer {
+    name: "rialto_signed",
+    src: ":empty_file",
+    filename: "rialto",
+    partition_name: "boot",
+    private_key: ":rialto_sign_key",
+    salt: rialto_salt,
+    enabled: false,
+    arch: {
+        arm64: {
+            src: ":rialto_unsigned",
+            enabled: true,
+        },
+    },
+}
+
+prebuilt_etc {
+    name: "rialto_bin",
+    filename: "rialto.bin",
+    target: {
+        android_arm64: {
+            src: ":rialto_signed",
+        },
+    },
+    src: ":empty_file",
+    installable: false,
+}
+
 rust_test {
     name: "rialto_test",
     crate_name: "rialto_test",
@@ -62,7 +100,8 @@
         "libvmclient",
     ],
     data: [
-        ":rialto",
+        ":rialto_bin",
+        ":rialto_unsigned",
     ],
     test_suites: ["general-tests"],
     enabled: false,
diff --git a/rialto/src/error.rs b/rialto/src/error.rs
new file mode 100644
index 0000000..8f34676
--- /dev/null
+++ b/rialto/src/error.rs
@@ -0,0 +1,55 @@
+// Copyright 2022, 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.
+
+//! This module contains the error thrown by Rialto.
+
+use aarch64_paging::MapError;
+use core::{fmt, result};
+use hyp::mmio_guard::Error as MmioError;
+
+pub type Result<T> = result::Result<T, Error>;
+
+#[derive(Clone, Debug)]
+pub enum Error {
+    /// MMIO guard failed.
+    MmioGuard(MmioError),
+    /// Failed when attempting to map some range in the page table.
+    PageTableMapping(MapError),
+    /// Failed to initialize the logger.
+    LoggerInit,
+}
+
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::MmioGuard(e) => write!(f, "MMIO guard failed: {e}."),
+            Self::PageTableMapping(e) => {
+                write!(f, "Failed when attempting to map some range in the page table: {e}.")
+            }
+            Self::LoggerInit => write!(f, "Failed to initialize the logger."),
+        }
+    }
+}
+
+impl From<MmioError> for Error {
+    fn from(e: MmioError) -> Self {
+        Self::MmioGuard(e)
+    }
+}
+
+impl From<MapError> for Error {
+    fn from(e: MapError) -> Self {
+        Self::PageTableMapping(e)
+    }
+}
diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index 59ee0b6..76f5495 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -17,18 +17,20 @@
 #![no_main]
 #![no_std]
 
+mod error;
 mod exceptions;
 
 extern crate alloc;
 
+use crate::error::{Error, Result};
 use aarch64_paging::{
     idmap::IdMap,
     paging::{Attributes, MemoryRegion},
-    MapError,
 };
 use buddy_system_allocator::LockedHeap;
-use log::{debug, info};
-use vmbase::main;
+use hyp::mmio_guard;
+use log::{debug, error, info};
+use vmbase::{main, power::reboot};
 
 const SZ_1K: usize = 1024;
 const SZ_64K: usize = 64 * SZ_1K;
@@ -81,7 +83,7 @@
     info!("Initialized heap.");
 }
 
-fn init_kernel_pgt(pgt: &mut IdMap) -> Result<(), MapError> {
+fn init_kernel_pgt(pgt: &mut IdMap) -> Result<()> {
     // The first 1 GiB of address space is used by crosvm for MMIO.
     let reg_dev = MemoryRegion::new(0, SZ_1G);
     // SAFETY: Taking addresses of kernel image sections to set up page table
@@ -106,15 +108,39 @@
     Ok(())
 }
 
-/// Entry point for Rialto.
-pub fn main(_a0: u64, _a1: u64, _a2: u64, _a3: u64) {
-    vmbase::logger::init(log::LevelFilter::Debug).unwrap();
+fn try_init_logger() -> Result<()> {
+    match mmio_guard::init() {
+        // pKVM blocks MMIO by default, we need to enable MMIO guard to support logging.
+        Ok(()) => mmio_guard::map(vmbase::console::BASE_ADDRESS)?,
+        // MMIO guard enroll is not supported in unprotected VM.
+        Err(mmio_guard::Error::EnrollFailed(smccc::Error::NotSupported)) => {}
+        Err(e) => return Err(e.into()),
+    };
+    vmbase::logger::init(log::LevelFilter::Debug).map_err(|_| Error::LoggerInit)
+}
 
+fn try_main() -> Result<()> {
     info!("Welcome to Rialto!");
     init_heap();
 
     let mut pgt = IdMap::new(PT_ASID, PT_ROOT_LEVEL);
-    init_kernel_pgt(&mut pgt).unwrap();
+    init_kernel_pgt(&mut pgt)?;
+    Ok(())
+}
+
+/// Entry point for Rialto.
+pub fn main(_a0: u64, _a1: u64, _a2: u64, _a3: u64) {
+    if try_init_logger().is_err() {
+        // Don't log anything if the logger initialization fails.
+        reboot();
+    }
+    match try_main() {
+        Ok(()) => info!("Rialto ends successfully."),
+        Err(e) => {
+            error!("Rialto failed with {e}");
+            reboot()
+        }
+    }
 }
 
 extern "C" {
diff --git a/rialto/tests/test.rs b/rialto/tests/test.rs
index be5f118..7048b44 100644
--- a/rialto/tests/test.rs
+++ b/rialto/tests/test.rs
@@ -16,7 +16,8 @@
 
 use android_system_virtualizationservice::{
     aidl::android::system::virtualizationservice::{
-        CpuTopology::CpuTopology, VirtualMachineConfig::VirtualMachineConfig,
+        CpuTopology::CpuTopology, DiskImage::DiskImage, Partition::Partition,
+        PartitionType::PartitionType, VirtualMachineConfig::VirtualMachineConfig,
         VirtualMachineRawConfig::VirtualMachineRawConfig,
     },
     binder::{ParcelFileDescriptor, ProcessState},
@@ -31,11 +32,28 @@
 use std::time::Duration;
 use vmclient::{DeathReason, VmInstance};
 
-const RIALTO_PATH: &str = "/data/local/tmp/rialto_test/arm64/rialto.bin";
+const SIGNED_RIALTO_PATH: &str = "/data/local/tmp/rialto_test/arm64/rialto.bin";
+const UNSIGNED_RIALTO_PATH: &str = "/data/local/tmp/rialto_test/arm64/rialto_unsigned.bin";
+const INSTANCE_IMG_PATH: &str = "/data/local/tmp/rialto_test/arm64/instance.img";
+const INSTANCE_IMG_SIZE: i64 = 1024 * 1024; // 1MB
 
-/// Runs the Rialto VM as a non-protected VM via VirtualizationService.
 #[test]
-fn test_boots() -> Result<(), Error> {
+fn boot_rialto_in_protected_vm_successfully() -> Result<(), Error> {
+    boot_rialto_successfully(
+        SIGNED_RIALTO_PATH,
+        true, // protected_vm
+    )
+}
+
+#[test]
+fn boot_rialto_in_unprotected_vm_successfully() -> Result<(), Error> {
+    boot_rialto_successfully(
+        UNSIGNED_RIALTO_PATH,
+        false, // protected_vm
+    )
+}
+
+fn boot_rialto_successfully(rialto_path: &str, protected_vm: bool) -> Result<(), Error> {
     android_logger::init_once(
         android_logger::Config::default().with_tag("rialto").with_min_level(log::Level::Debug),
     );
@@ -52,18 +70,44 @@
         vmclient::VirtualizationService::new().context("Failed to spawn VirtualizationService")?;
     let service = virtmgr.connect().context("Failed to connect to VirtualizationService")?;
 
-    let rialto = File::open(RIALTO_PATH).context("Failed to open Rialto kernel binary")?;
+    let rialto = File::open(rialto_path).context("Failed to open Rialto kernel binary")?;
     let console = android_log_fd()?;
     let log = android_log_fd()?;
 
+    let disks = if protected_vm {
+        let instance_img = File::options()
+            .create(true)
+            .read(true)
+            .write(true)
+            .truncate(true)
+            .open(INSTANCE_IMG_PATH)?;
+        let instance_img = ParcelFileDescriptor::new(instance_img);
+
+        service
+            .initializeWritablePartition(
+                &instance_img,
+                INSTANCE_IMG_SIZE,
+                PartitionType::ANDROID_VM_INSTANCE,
+            )
+            .context("Failed to initialize instange.img")?;
+        let writable_partitions = vec![Partition {
+            label: "vm-instance".to_owned(),
+            image: Some(instance_img),
+            writable: true,
+        }];
+        vec![DiskImage { image: None, partitions: writable_partitions, writable: true }]
+    } else {
+        vec![]
+    };
+
     let config = VirtualMachineConfig::RawConfig(VirtualMachineRawConfig {
         name: String::from("RialtoTest"),
         kernel: None,
         initrd: None,
         params: None,
         bootloader: Some(ParcelFileDescriptor::new(rialto)),
-        disks: vec![],
-        protectedVm: false,
+        disks,
+        protectedVm: protected_vm,
         memoryMib: 300,
         cpuTopology: CpuTopology::ONE_CPU,
         platformVersion: "~1.0".to_string(),