[rialto] Enable Rialto to run in protected VM

Prior to this cl, Rialto couldn't run in protected VM because MMIO
access is blocked by default in protected mode, this cl enables
the MMIO guard. We should also unmap the MMIO range used in Rialto
in the future.

Bug: 272226230
Test: atest rialto_test
Change-Id: Iaf4a9c74ddbd068ee10cc981c261d38a02c63993
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(),