Merge "Config file requires custom VM permission"
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index 1da17b9..d8ff8c6 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -353,7 +353,7 @@
             }
 
             if (!mProtectedVmSet) {
-                throw new IllegalStateException("protectedVm must be set explicitly");
+                throw new IllegalStateException("setProtectedVm(t/f) must be called explicitly");
             }
 
             if (mProtectedVm
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
new file mode 100644
index 0000000..95dc0b0
--- /dev/null
+++ b/pvmfw/src/entry.rs
@@ -0,0 +1,81 @@
+// 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.
+
+//! Low-level entry and exit points of pvmfw.
+
+use crate::helpers::FDT_MAX_SIZE;
+use crate::jump_to_payload;
+use crate::mmio_guard;
+use core::slice;
+use log::{debug, LevelFilter};
+use vmbase::{console, logger, main, power::reboot};
+
+#[derive(Debug, Clone)]
+enum RebootReason {
+    /// An unexpected internal error happened.
+    InternalError,
+}
+
+main!(start);
+
+/// Entry point for pVM firmware.
+pub fn start(fdt_address: u64, payload_start: u64, payload_size: u64, _arg3: u64) {
+    // Limitations in this function:
+    // - can't access non-pvmfw memory (only statically-mapped memory)
+    // - can't access MMIO (therefore, no logging)
+
+    match main_wrapper(fdt_address as usize, payload_start as usize, payload_size as usize) {
+        Ok(_) => jump_to_payload(fdt_address, payload_start),
+        Err(_) => reboot(),
+    }
+
+    // if we reach this point and return, vmbase::entry::rust_entry() will call power::shutdown().
+}
+
+/// Sets up the environment for main() and wraps its result for start().
+///
+/// Provide the abstractions necessary for start() to abort the pVM boot and for main() to run with
+/// the assumption that its environment has been properly configured.
+fn main_wrapper(fdt: usize, payload: usize, payload_size: usize) -> Result<(), RebootReason> {
+    // Limitations in this function:
+    // - only access MMIO once (and while) it has been mapped and configured
+    // - only perform logging once the logger has been initialized
+    // - only access non-pvmfw memory once (and while) it has been mapped
+    logger::init(LevelFilter::Info).map_err(|_| RebootReason::InternalError)?;
+
+    // TODO: Check that the FDT is fully contained in RAM.
+    // SAFETY - We trust the VMM, for now.
+    let fdt = unsafe { slice::from_raw_parts_mut(fdt as *mut u8, FDT_MAX_SIZE) };
+    // TODO: Check that the payload is fully contained in RAM and doesn't overlap with the FDT.
+    // SAFETY - We trust the VMM, for now.
+    let payload = unsafe { slice::from_raw_parts(payload as *const u8, payload_size) };
+
+    // Use debug!() to avoid printing to the UART if we failed to configure it as only local
+    // builds that have tweaked the logger::init() call will actually attempt to log the message.
+
+    mmio_guard::init().map_err(|e| {
+        debug!("{e}");
+        RebootReason::InternalError
+    })?;
+
+    mmio_guard::map(console::BASE_ADDRESS).map_err(|e| {
+        debug!("Failed to configure the UART: {e}");
+        RebootReason::InternalError
+    })?;
+
+    // This wrapper allows main() to be blissfully ignorant of platform details.
+    crate::main(fdt, payload);
+
+    Ok(())
+}
diff --git a/pvmfw/src/helpers.rs b/pvmfw/src/helpers.rs
index ac3a48e..cade62e 100644
--- a/pvmfw/src/helpers.rs
+++ b/pvmfw/src/helpers.rs
@@ -14,6 +14,7 @@
 
 //! Miscellaneous helper functions.
 
+pub const FDT_MAX_SIZE: usize = 2 << 20;
 pub const SIZE_4KB: usize = 4 << 10;
 
 /// Computes the address of the page containing a given address.
@@ -21,15 +22,6 @@
     addr & !(page_size - 1)
 }
 
-/// Validates a page size and computes the address of the page containing a given address.
-pub const fn checked_page_of(addr: usize, page_size: usize) -> Option<usize> {
-    if page_size.is_power_of_two() {
-        Some(page_of(addr, page_size))
-    } else {
-        None
-    }
-}
-
 /// Computes the address of the 4KiB page containing a given address.
 pub const fn page_4kb_of(addr: usize) -> usize {
     page_of(addr, SIZE_4KB)
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index 55604e5..d1951b3 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -17,60 +17,24 @@
 #![no_main]
 #![no_std]
 
+mod entry;
 mod exceptions;
 mod helpers;
+mod mmio_guard;
 mod smccc;
 
-use core::fmt;
-use helpers::checked_page_of;
-use log::{debug, error, info, LevelFilter};
-use vmbase::{console, logger, main, power::reboot};
+use log::{debug, info};
 
-#[derive(Debug, Clone)]
-enum Error {
-    /// Failed to configure the UART; no logs available.
-    FailedUartSetup,
-}
-
-impl fmt::Display for Error {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        let msg = match self {
-            Self::FailedUartSetup => "Failed to configure the UART",
-        };
-        write!(f, "{}", msg)
-    }
-}
-
-fn main(fdt_address: u64, payload_start: u64, payload_size: u64, arg3: u64) -> Result<(), Error> {
-    // We need to inform the hypervisor that the MMIO page containing the UART may be shared back.
-    let mmio_granule = smccc::mmio_guard_info().map_err(|_| Error::FailedUartSetup)?;
-    let uart_page = checked_page_of(console::BASE_ADDRESS, mmio_granule as usize)
-        .ok_or(Error::FailedUartSetup)?;
-    smccc::mmio_guard_map(uart_page as u64).map_err(|_| Error::FailedUartSetup)?;
-
+fn main(fdt: &mut [u8], payload: &[u8]) {
     info!("pVM firmware");
     debug!(
-        "fdt_address={:#018x}, payload_start={:#018x}, payload_size={:#018x}, x3={:#018x}",
-        fdt_address, payload_start, payload_size, arg3,
+        "fdt_address={:#018x}, payload_start={:#018x}, payload_size={:#018x}",
+        fdt.as_ptr() as usize,
+        payload.as_ptr() as usize,
+        payload.len(),
     );
 
     info!("Starting payload...");
-
-    Ok(())
-}
-
-main!(main_wrapper);
-
-/// Entry point for pVM firmware.
-pub fn main_wrapper(fdt_address: u64, payload_start: u64, payload_size: u64, arg3: u64) {
-    if logger::init(LevelFilter::Debug).is_err() {
-    } else if let Err(e) = main(fdt_address, payload_start, payload_size, arg3) {
-        error!("Boot rejected: {e}");
-    } else {
-        jump_to_payload(fdt_address, payload_start);
-    }
-
-    reboot()
 }
 
 fn jump_to_payload(fdt_address: u64, payload_start: u64) {
diff --git a/pvmfw/src/mmio_guard.rs b/pvmfw/src/mmio_guard.rs
new file mode 100644
index 0000000..4cde737
--- /dev/null
+++ b/pvmfw/src/mmio_guard.rs
@@ -0,0 +1,53 @@
+// 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.
+
+//! Safe MMIO_GUARD support.
+
+use crate::helpers;
+use crate::smccc;
+use core::{fmt, result};
+
+#[derive(Debug, Clone)]
+pub enum Error {
+    /// Failed to obtain the MMIO_GUARD granule size.
+    InfoFailed(smccc::Error),
+    /// Failed to MMIO_GUARD_MAP a page.
+    MapFailed(smccc::Error),
+    /// The MMIO_GUARD granule used by the hypervisor is not supported.
+    UnsupportedGranule(usize),
+}
+
+type Result<T> = result::Result<T, Error>;
+
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::InfoFailed(e) => write!(f, "Failed to get the MMIO_GUARD granule: {e}"),
+            Self::MapFailed(e) => write!(f, "Failed to MMIO_GUARD map: {e}"),
+            Self::UnsupportedGranule(g) => write!(f, "Unsupported MMIO_GUARD granule: {g}"),
+        }
+    }
+}
+
+pub fn init() -> Result<()> {
+    let mmio_granule = smccc::mmio_guard_info().map_err(Error::InfoFailed)? as usize;
+    if mmio_granule != helpers::SIZE_4KB {
+        return Err(Error::UnsupportedGranule(mmio_granule));
+    }
+    Ok(())
+}
+
+pub fn map(addr: usize) -> Result<()> {
+    smccc::mmio_guard_map(helpers::page_4kb_of(addr) as u64).map_err(Error::MapFailed)
+}
diff --git a/vm/vm_shell.sh b/vm/vm_shell.sh
index ec9243b..c0dd38f 100755
--- a/vm/vm_shell.sh
+++ b/vm/vm_shell.sh
@@ -27,7 +27,7 @@
     adb forward tcp:8000 vsock:${cid}:5555
     adb connect localhost:8000
     adb -s localhost:8000 root
-    sleep 2
+    adb -s localhost:8000 wait-for-device
     adb -s localhost:8000 shell
     exit 0
 }
@@ -40,12 +40,7 @@
     exit 1
 fi
 
-if [ -n "${selected_cid}" ]; then
-    if [[ ! " ${available_cids[*]} " =~ " ${selected_cid} " ]]; then
-        echo VM of CID $selected_cid does not exist. Available CIDs: ${available_cids}
-        exit 1
-    fi
-else
+if [ ! -n "${selected_cid}" ]; then
     PS3="Select CID of VM to adb-shell into: "
     select cid in ${available_cids}
     do
@@ -54,4 +49,9 @@
     done
 fi
 
+if [[ ! " ${available_cids[*]} " =~ " ${selected_cid} " ]]; then
+    echo VM of CID $selected_cid does not exist. Available CIDs: ${available_cids}
+    exit 1
+fi
+
 connect_vm ${selected_cid}