Merge "Give a better name to STOP_REASON_ERROR"
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index ee32509..efbb179 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -74,8 +74,8 @@
 impl<'a> MemorySlices<'a> {
     fn new(
         fdt: usize,
-        payload: usize,
-        payload_size: usize,
+        kernel: usize,
+        kernel_size: usize,
         memory: &mut MemoryTracker,
     ) -> Result<Self, RebootReason> {
         // SAFETY - SIZE_2MB is non-zero.
@@ -120,18 +120,36 @@
             RebootReason::InvalidFdt
         })?;
 
-        let payload_size = NonZeroUsize::new(payload_size).ok_or_else(|| {
-            error!("Invalid payload size: {payload_size:#x}");
-            RebootReason::InvalidPayload
+        let kernel_range = fdt::kernel_range(fdt).map_err(|e| {
+            error!("Error while attempting to read the kernel range from the DT: {e}");
+            RebootReason::InvalidFdt
         })?;
 
-        let payload_range = memory.alloc(payload, payload_size).map_err(|e| {
-            error!("Failed to obtain the payload range: {e}");
-            RebootReason::InternalError
-        })?;
+        let kernel_range = if let Some(r) = kernel_range {
+            memory.alloc_range(&r).map_err(|e| {
+                error!("Failed to obtain the kernel range with DT range: {e}");
+                RebootReason::InternalError
+            })?
+        } else if cfg!(feature = "legacy") {
+            warn!("Failed to find the kernel range in the DT; falling back to legacy ABI");
+
+            let kernel_size = NonZeroUsize::new(kernel_size).ok_or_else(|| {
+                error!("Invalid kernel size: {kernel_size:#x}");
+                RebootReason::InvalidPayload
+            })?;
+
+            memory.alloc(kernel, kernel_size).map_err(|e| {
+                error!("Failed to obtain the kernel range with legacy range: {e}");
+                RebootReason::InternalError
+            })?
+        } else {
+            error!("Failed to locate the kernel from the DT");
+            return Err(RebootReason::InvalidPayload);
+        };
+
         // SAFETY - The tracker validated the range to be in main memory, mapped, and not overlap.
         let kernel =
-            unsafe { slice::from_raw_parts(payload_range.start as *const u8, payload_range.len()) };
+            unsafe { slice::from_raw_parts(kernel_range.start as *const u8, kernel_range.len()) };
 
         let ramdisk_range = fdt::initrd_range(fdt).map_err(|e| {
             error!("An error occurred while locating the ramdisk in the device tree: {e}");
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index 5b9efd2..dcd17b7 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -17,6 +17,24 @@
 use core::ffi::CStr;
 use core::ops::Range;
 
+/// Extract from /config the address range containing the pre-loaded kernel.
+pub fn kernel_range(fdt: &libfdt::Fdt) -> libfdt::Result<Option<Range<usize>>> {
+    let config = CStr::from_bytes_with_nul(b"/config\0").unwrap();
+    let addr = CStr::from_bytes_with_nul(b"kernel-address\0").unwrap();
+    let size = CStr::from_bytes_with_nul(b"kernel-size\0").unwrap();
+
+    if let Some(config) = fdt.node(config)? {
+        if let (Some(addr), Some(size)) = (config.getprop_u32(addr)?, config.getprop_u32(size)?) {
+            let addr = addr as usize;
+            let size = size as usize;
+
+            return Ok(Some(addr..(addr + size)));
+        }
+    }
+
+    Ok(None)
+}
+
 /// Extract from /chosen the address range containing the pre-loaded ramdisk.
 pub fn initrd_range(fdt: &libfdt::Fdt) -> libfdt::Result<Option<Range<usize>>> {
     let start = CStr::from_bytes_with_nul(b"linux,initrd-start\0").unwrap();
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index cf7e90a..d453e26 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -55,6 +55,6 @@
         error!("Failed to verify the payload: {e}");
         RebootReason::PayloadVerificationError
     })?;
-    info!("Payload verified. Starting payload...");
+    info!("Starting payload...");
     Ok(())
 }
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index bf2d411..5ce8f60 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -408,27 +408,33 @@
 
     @Test
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-6"})
-    public void protectedVmWithValidKernelImageRunsPvmfw() throws Exception {
+    public void protectedVmRunsPvmfw() throws Exception {
         // Arrange
         boolean protectedVm = true;
         assumeTrue(
                 "Skip if protected VMs are not supported",
                 getAndroidDevice().supportsMicrodroid(protectedVm));
-        File key = findTestFile("test.com.android.virt.pem");
+        final String configPath = "assets/vm_config_apex.json";
 
         // Act
-        // TODO(b/256148034): Do not resign kernel image
-        VmInfo vmInfo =
-                runMicrodroidWithResignedImages(key, /*keyOverrides=*/ Map.of(), protectedVm);
+        mMicrodroidDevice =
+                MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
+                        .debugLevel("full")
+                        .memoryMib(minMemorySize())
+                        .numCpus(NUM_VCPUS)
+                        .protectedVm(protectedVm)
+                        .build(getAndroidDevice());
 
         // Assert
-        vmInfo.mProcess.waitFor(5L, TimeUnit.SECONDS);
+        mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
         String consoleLog = getDevice().pullFileContents(CONSOLE_PATH);
-        assertWithMessage("pvmfw should start").that(consoleLog).contains("pVM firmware");
-        assertWithMessage("pvmfw should start payload")
+        assertWithMessage("Failed to verify that pvmfw started")
                 .that(consoleLog)
-                .contains("Payload verified. Starting payload...");
-        vmInfo.mProcess.destroy();
+                .contains("pVM firmware");
+        assertWithMessage("pvmfw failed to start kernel")
+                .that(consoleLog)
+                .contains("Starting payload...");
+        // TODO(b/260994818): Investigate the feasibility of checking DeathReason.
     }
 
     @Test