Merge "lib/hyp: Introduce MMIO_GUARD_GRANULE_SIZE"
diff --git a/apex/Android.bp b/apex/Android.bp
index 1c4d357..2dc90d6 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -11,8 +11,14 @@
     name: "virt_apex",
     module_type: "apex",
     config_namespace: "ANDROID",
-    bool_variables: ["avf_enabled"],
-    properties: ["defaults"],
+    bool_variables: [
+        "avf_enabled",
+        "avf_kernel_modules_enabled",
+    ],
+    properties: [
+        "defaults",
+        "prebuilts",
+    ],
 }
 
 virt_apex {
@@ -24,6 +30,9 @@
                 defaults: ["com.android.virt_avf_disabled"],
             },
         },
+        avf_kernel_modules_enabled: {
+            prebuilts: ["microdroid_kernel_with_modules"],
+        },
     },
 }
 
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index 232485a..b03addf 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -130,6 +130,7 @@
             cpuTopology: cpu_topology,
             taskProfiles: parameters.task_profiles.clone(),
             gdbPort: 0, // Don't start gdb-server
+            customKernelImage: None,
         });
 
         // Let logs go to logcat.
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 1092476..f3045b7 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -461,3 +461,34 @@
         },
     },
 }
+
+avb_add_hash_footer {
+    name: "microdroid_kernel_with_modules_signed",
+    src: ":empty_file",
+    filename: "microdroid_kernel_with_modules",
+    partition_name: "boot",
+    private_key: ":microdroid_sign_key",
+    salt: bootloader_salt,
+    enabled: false,
+    arch: {
+        arm64: {
+            src: ":microdroid_kernel_with_modules_prebuilts-6.1-arm64",
+            enabled: true,
+        },
+    },
+    include_descriptors_from_images: [
+        ":microdroid_initrd_normal_hashdesc",
+        ":microdroid_initrd_debug_hashdesc",
+    ],
+}
+
+prebuilt_etc {
+    name: "microdroid_kernel_with_modules",
+    src: ":empty_file",
+    relative_install_path: "fs",
+    arch: {
+        arm64: {
+            src: ":microdroid_kernel_with_modules_signed",
+        },
+    },
+}
diff --git a/microdroid/kernel/with-modules/Android.bp b/microdroid/kernel/with-modules/Android.bp
new file mode 100644
index 0000000..f1ec06e
--- /dev/null
+++ b/microdroid/kernel/with-modules/Android.bp
@@ -0,0 +1,32 @@
+// Copyright (C) 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.
+
+package {
+    default_applicable_licenses: ["microdroid_kernel_with_modules_prebuilts_6.1_arm64_license"],
+}
+
+// See: http://go/android-license-faq
+license {
+    name: "microdroid_kernel_with_modules_prebuilts_6.1_arm64_license",
+    visibility: [":__subpackages__"],
+    license_kinds: [
+        "SPDX-license-identifier-GPL-2.0-only",
+    ],
+    // large-scale-change unable to identify any license_text files
+}
+
+filegroup {
+    name: "microdroid_kernel_with_modules_prebuilts-6.1-arm64",
+    srcs: ["kernel_with_modules-6.1"],
+}
diff --git a/microdroid/kernel/with-modules/README.md b/microdroid/kernel/with-modules/README.md
new file mode 100644
index 0000000..46f6a59
--- /dev/null
+++ b/microdroid/kernel/with-modules/README.md
@@ -0,0 +1,10 @@
+# Microdroid kernel with modules
+
+This directory contains prebuilts of the Microdroid kernel with support for
+loading vendor modules. Only arm64 architecture is supported.
+
+NOTE: the prebuilt was generated on a local machine, and can only be used on
+development devices.
+
+NOTE: this feature is WIP and kernel is subject to change.
+
diff --git a/microdroid/kernel/with-modules/kernel_with_modules-6.1 b/microdroid/kernel/with-modules/kernel_with_modules-6.1
new file mode 100755
index 0000000..b732ce3
--- /dev/null
+++ b/microdroid/kernel/with-modules/kernel_with_modules-6.1
Binary files differ
diff --git a/pvmfw/src/exceptions.rs b/pvmfw/src/exceptions.rs
index a10c300..4d376cd 100644
--- a/pvmfw/src/exceptions.rs
+++ b/pvmfw/src/exceptions.rs
@@ -14,11 +14,11 @@
 
 //! Exception handlers.
 
-use crate::memory::{MemoryTrackerError, MEMORY};
+use crate::memory::MEMORY;
 use core::fmt;
 use vmbase::console;
 use vmbase::logger;
-use vmbase::memory::page_4kb_of;
+use vmbase::memory::{page_4kb_of, MemoryTrackerError};
 use vmbase::read_sysreg;
 use vmbase::{eprintln, power::reboot};
 
diff --git a/pvmfw/src/memory.rs b/pvmfw/src/memory.rs
index fa2d56b..c34afc9 100644
--- a/pvmfw/src/memory.rs
+++ b/pvmfw/src/memory.rs
@@ -22,7 +22,6 @@
 use alloc::boxed::Box;
 use buddy_system_allocator::LockedFrameAllocator;
 use core::alloc::Layout;
-use core::fmt;
 use core::iter::once;
 use core::num::NonZeroUsize;
 use core::ops::Range;
@@ -38,7 +37,7 @@
     dsb, layout,
     memory::{
         flush_dirty_range, mark_dirty_block, mmio_guard_unmap_page, page_4kb_of, set_dbm_enabled,
-        verify_lazy_mapped_block, MemorySharer, PageTable, SIZE_2MB, SIZE_4KB,
+        verify_lazy_mapped_block, MemorySharer, MemoryTrackerError, PageTable, SIZE_2MB, SIZE_4KB,
     },
     util::{align_up, RangeExt as _},
 };
@@ -74,66 +73,6 @@
     payload_range: MemoryRange,
 }
 
-/// Errors for MemoryTracker operations.
-#[derive(Debug, Clone)]
-pub enum MemoryTrackerError {
-    /// Tried to modify the memory base address.
-    DifferentBaseAddress,
-    /// Tried to shrink to a larger memory size.
-    SizeTooLarge,
-    /// Tracked regions would not fit in memory size.
-    SizeTooSmall,
-    /// Reached limit number of tracked regions.
-    Full,
-    /// Region is out of the tracked memory address space.
-    OutOfRange,
-    /// New region overlaps with tracked regions.
-    Overlaps,
-    /// Region couldn't be mapped.
-    FailedToMap,
-    /// Region couldn't be unmapped.
-    FailedToUnmap,
-    /// Error from the interaction with the hypervisor.
-    Hypervisor(hyp::Error),
-    /// Failure to set `SHARED_MEMORY`.
-    SharedMemorySetFailure,
-    /// Failure to set `SHARED_POOL`.
-    SharedPoolSetFailure,
-    /// Invalid page table entry.
-    InvalidPte,
-    /// Failed to flush memory region.
-    FlushRegionFailed,
-    /// Failed to set PTE dirty state.
-    SetPteDirtyFailed,
-}
-
-impl fmt::Display for MemoryTrackerError {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self {
-            Self::DifferentBaseAddress => write!(f, "Received different base address"),
-            Self::SizeTooLarge => write!(f, "Tried to shrink to a larger memory size"),
-            Self::SizeTooSmall => write!(f, "Tracked regions would not fit in memory size"),
-            Self::Full => write!(f, "Reached limit number of tracked regions"),
-            Self::OutOfRange => write!(f, "Region is out of the tracked memory address space"),
-            Self::Overlaps => write!(f, "New region overlaps with tracked regions"),
-            Self::FailedToMap => write!(f, "Failed to map the new region"),
-            Self::FailedToUnmap => write!(f, "Failed to unmap the new region"),
-            Self::Hypervisor(e) => e.fmt(f),
-            Self::SharedMemorySetFailure => write!(f, "Failed to set SHARED_MEMORY"),
-            Self::SharedPoolSetFailure => write!(f, "Failed to set SHARED_POOL"),
-            Self::InvalidPte => write!(f, "Page table entry is not valid"),
-            Self::FlushRegionFailed => write!(f, "Failed to flush memory region"),
-            Self::SetPteDirtyFailed => write!(f, "Failed to set PTE dirty state"),
-        }
-    }
-}
-
-impl From<hyp::Error> for MemoryTrackerError {
-    fn from(e: hyp::Error) -> Self {
-        Self::Hypervisor(e)
-    }
-}
-
 type Result<T> = result::Result<T, MemoryTrackerError>;
 
 static SHARED_POOL: OnceBox<LockedFrameAllocator<32>> = OnceBox::new();
diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index bc5ab2c..9736aa4 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -69,15 +69,19 @@
     Ok(())
 }
 
-fn try_init_logger() -> Result<()> {
-    match get_hypervisor().mmio_guard_init() {
+fn try_init_logger() -> Result<bool> {
+    let mmio_guard_supported = match get_hypervisor().mmio_guard_init() {
         // pKVM blocks MMIO by default, we need to enable MMIO guard to support logging.
-        Ok(()) => get_hypervisor().mmio_guard_map(vmbase::console::BASE_ADDRESS)?,
+        Ok(()) => {
+            get_hypervisor().mmio_guard_map(vmbase::console::BASE_ADDRESS)?;
+            true
+        }
         // MMIO guard enroll is not supported in unprotected VM.
-        Err(hyp::Error::MmioGuardNotsupported) => {}
+        Err(hyp::Error::MmioGuardNotsupported) => false,
         Err(e) => return Err(e.into()),
     };
-    vmbase::logger::init(log::LevelFilter::Debug).map_err(|_| Error::LoggerInit)
+    vmbase::logger::init(log::LevelFilter::Debug).map_err(|_| Error::LoggerInit)?;
+    Ok(mmio_guard_supported)
 }
 
 /// # Safety
@@ -96,19 +100,39 @@
     Ok(())
 }
 
+fn try_unshare_all_memory(mmio_guard_supported: bool) -> Result<()> {
+    if !mmio_guard_supported {
+        return Ok(());
+    }
+    info!("Starting unsharing memory...");
+
+    // TODO(b/284462758): Unshare all the memory here.
+
+    // No logging after unmapping UART.
+    get_hypervisor().mmio_guard_unmap(vmbase::console::BASE_ADDRESS)?;
+    Ok(())
+}
+
+fn unshare_all_memory(mmio_guard_supported: bool) {
+    if let Err(e) = try_unshare_all_memory(mmio_guard_supported) {
+        error!("Failed to unshare the memory: {e}");
+    }
+}
+
 /// Entry point for Rialto.
 pub fn main(fdt_addr: u64, _a1: u64, _a2: u64, _a3: u64) {
     init_heap();
-    if try_init_logger().is_err() {
+    let Ok(mmio_guard_supported) = try_init_logger() else {
         // Don't log anything if the logger initialization fails.
         reboot();
-    }
+    };
     // SAFETY: `fdt_addr` is supposed to be a valid pointer and points to
     // a valid `Fdt`.
     match unsafe { try_main(fdt_addr as usize) } {
-        Ok(()) => info!("Rialto ends successfully."),
+        Ok(()) => unshare_all_memory(mmio_guard_supported),
         Err(e) => {
             error!("Rialto failed with {e}");
+            unshare_all_memory(mmio_guard_supported);
             reboot()
         }
     }
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index f57cb59..86c8596 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -309,10 +309,12 @@
                 // VirtualMachineAppConfig:
                 // - controlling CPUs;
                 // - specifying a config file in the APK;
-                // - gdbPort is set, meaning that crosvm will start a gdb server.
+                // - gdbPort is set, meaning that crosvm will start a gdb server;
+                // - using anything other than the default kernel.
                 !config.taskProfiles.is_empty()
                     || matches!(config.payload, Payload::ConfigPath(_))
                     || config.gdbPort > 0
+                    || config.customKernelImage.as_ref().is_some()
             }
         };
         if is_custom {
@@ -593,6 +595,10 @@
     let vm_config_file = File::open(vm_config_path)?;
     let mut vm_config = VmConfig::load(&vm_config_file)?.to_parcelable()?;
 
+    if let Some(file) = config.customKernelImage.as_ref() {
+        vm_config.kernel = Some(ParcelFileDescriptor::new(clone_file(file)?))
+    }
+
     if config.memoryMib > 0 {
         vm_config.memoryMib = config.memoryMib;
     }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
index c467c2f..5e05bb9 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
@@ -73,6 +73,10 @@
     /**
      * Port at which crosvm will start a gdb server to debug guest kernel.
      * If set to zero, then gdb server won't be started.
+     *
+     * Note: Specifying a value here requires android.permission.USE_CUSTOM_VIRTUAL_MACHINE.
+     *
+     * TODO(b/286225150): move to a separate struct
      */
     int gdbPort = 0;
 
@@ -92,6 +96,17 @@
      * List of task profile names to apply for the VM
      *
      * Note: Specifying a value here requires android.permission.USE_CUSTOM_VIRTUAL_MACHINE.
+     *
+     * TODO(b/286225150): move to a separate struct
      */
     String[] taskProfiles;
+
+    /**
+     * If specified, boot Microdroid VM with the given kernel.
+     *
+     * Note: Specifying a value here requires android.permission.USE_CUSTOM_VIRTUAL_MACHINE.
+     *
+     * TODO(b/286225150): move to a separate struct
+     */
+    @nullable ParcelFileDescriptor customKernelImage;
 }
diff --git a/vm/src/main.rs b/vm/src/main.rs
index 1d9f50b..bc3f4da 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -107,6 +107,10 @@
         /// Note: this is only supported on Android kernels android14-5.15 and higher.
         #[clap(long)]
         gdb: Option<NonZeroU16>,
+
+        /// Path to custom kernel image to use when booting Microdroid.
+        #[clap(long)]
+        kernel: Option<PathBuf>,
     },
     /// Run a virtual machine with Microdroid inside
     RunMicrodroid {
@@ -163,6 +167,10 @@
         /// Note: this is only supported on Android kernels android14-5.15 and higher.
         #[clap(long)]
         gdb: Option<NonZeroU16>,
+
+        /// Path to custom kernel image to use when booting Microdroid.
+        #[clap(long)]
+        kernel: Option<PathBuf>,
     },
     /// Run a virtual machine
     Run {
@@ -277,6 +285,7 @@
             task_profiles,
             extra_idsigs,
             gdb,
+            kernel,
         } => command_run_app(
             name,
             get_service()?.as_ref(),
@@ -296,6 +305,7 @@
             task_profiles,
             &extra_idsigs,
             gdb,
+            kernel.as_deref(),
         ),
         Opt::RunMicrodroid {
             name,
@@ -310,6 +320,7 @@
             cpu_topology,
             task_profiles,
             gdb,
+            kernel,
         } => command_run_microdroid(
             name,
             get_service()?.as_ref(),
@@ -324,6 +335,7 @@
             cpu_topology,
             task_profiles,
             gdb,
+            kernel.as_deref(),
         ),
         Opt::Run { name, config, cpu_topology, task_profiles, console, log, gdb } => {
             command_run(
diff --git a/vm/src/run.rs b/vm/src/run.rs
index 36edc64..54c1de4 100644
--- a/vm/src/run.rs
+++ b/vm/src/run.rs
@@ -60,6 +60,7 @@
     task_profiles: Vec<String>,
     extra_idsigs: &[PathBuf],
     gdb: Option<NonZeroU16>,
+    kernel: Option<&Path>,
 ) -> Result<(), Error> {
     let apk_file = File::open(apk).context("Failed to open APK file")?;
 
@@ -115,6 +116,8 @@
         None
     };
 
+    let kernel = kernel.map(|p| open_parcel_file(p, false)).transpose()?;
+
     let extra_idsig_files: Result<Vec<File>, _> = extra_idsigs.iter().map(File::open).collect();
     let extra_idsig_fds = extra_idsig_files?.into_iter().map(ParcelFileDescriptor::new).collect();
 
@@ -147,6 +150,7 @@
         cpuTopology: cpu_topology,
         taskProfiles: task_profiles,
         gdbPort: gdb.map(u16::from).unwrap_or(0) as i32, // 0 means no gdb
+        customKernelImage: kernel,
     });
     run(service, &config, &payload_config_str, console_path, log_path)
 }
@@ -189,6 +193,7 @@
     cpu_topology: CpuTopology,
     task_profiles: Vec<String>,
     gdb: Option<NonZeroU16>,
+    kernel: Option<&Path>,
 ) -> Result<(), Error> {
     let apk = find_empty_payload_apk_path()?;
     println!("found path {}", apk.display());
@@ -220,6 +225,7 @@
         task_profiles,
         &extra_sig,
         gdb,
+        kernel,
     )
 }
 
diff --git a/vmbase/src/memory/error.rs b/vmbase/src/memory/error.rs
new file mode 100644
index 0000000..273db56
--- /dev/null
+++ b/vmbase/src/memory/error.rs
@@ -0,0 +1,77 @@
+// 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.
+
+//! Error relating to memory management.
+
+use core::fmt;
+
+/// Errors for MemoryTracker operations.
+#[derive(Debug, Clone)]
+pub enum MemoryTrackerError {
+    /// Tried to modify the memory base address.
+    DifferentBaseAddress,
+    /// Tried to shrink to a larger memory size.
+    SizeTooLarge,
+    /// Tracked regions would not fit in memory size.
+    SizeTooSmall,
+    /// Reached limit number of tracked regions.
+    Full,
+    /// Region is out of the tracked memory address space.
+    OutOfRange,
+    /// New region overlaps with tracked regions.
+    Overlaps,
+    /// Region couldn't be mapped.
+    FailedToMap,
+    /// Region couldn't be unmapped.
+    FailedToUnmap,
+    /// Error from the interaction with the hypervisor.
+    Hypervisor(hyp::Error),
+    /// Failure to set `SHARED_MEMORY`.
+    SharedMemorySetFailure,
+    /// Failure to set `SHARED_POOL`.
+    SharedPoolSetFailure,
+    /// Invalid page table entry.
+    InvalidPte,
+    /// Failed to flush memory region.
+    FlushRegionFailed,
+    /// Failed to set PTE dirty state.
+    SetPteDirtyFailed,
+}
+
+impl fmt::Display for MemoryTrackerError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::DifferentBaseAddress => write!(f, "Received different base address"),
+            Self::SizeTooLarge => write!(f, "Tried to shrink to a larger memory size"),
+            Self::SizeTooSmall => write!(f, "Tracked regions would not fit in memory size"),
+            Self::Full => write!(f, "Reached limit number of tracked regions"),
+            Self::OutOfRange => write!(f, "Region is out of the tracked memory address space"),
+            Self::Overlaps => write!(f, "New region overlaps with tracked regions"),
+            Self::FailedToMap => write!(f, "Failed to map the new region"),
+            Self::FailedToUnmap => write!(f, "Failed to unmap the new region"),
+            Self::Hypervisor(e) => e.fmt(f),
+            Self::SharedMemorySetFailure => write!(f, "Failed to set SHARED_MEMORY"),
+            Self::SharedPoolSetFailure => write!(f, "Failed to set SHARED_POOL"),
+            Self::InvalidPte => write!(f, "Page table entry is not valid"),
+            Self::FlushRegionFailed => write!(f, "Failed to flush memory region"),
+            Self::SetPteDirtyFailed => write!(f, "Failed to set PTE dirty state"),
+        }
+    }
+}
+
+impl From<hyp::Error> for MemoryTrackerError {
+    fn from(e: hyp::Error) -> Self {
+        Self::Hypervisor(e)
+    }
+}
diff --git a/vmbase/src/memory/mod.rs b/vmbase/src/memory/mod.rs
index c4a54f2..f919bcd 100644
--- a/vmbase/src/memory/mod.rs
+++ b/vmbase/src/memory/mod.rs
@@ -15,11 +15,13 @@
 //! Memory management.
 
 mod dbm;
+mod error;
 mod page_table;
 mod shared;
 mod util;
 
 pub use dbm::{flush_dirty_range, mark_dirty_block, set_dbm_enabled};
+pub use error::MemoryTrackerError;
 pub use page_table::PageTable;
 pub use shared::{mmio_guard_unmap_page, verify_lazy_mapped_block, MemorySharer};
 pub use util::{