Merge changes I5f88ff52,I6de6e460,Id7f24fc6 into main

* changes:
  pvmfw: Parametrise GUEST_PAGE_SIZE
  pvmfw: Move DT sanitization to main()
  pvmfw: Read kernel/initrd ranges from untrusted DT
diff --git a/guest/pvmfw/avb/src/verify.rs b/guest/pvmfw/avb/src/verify.rs
index a073502..c85d886 100644
--- a/guest/pvmfw/avb/src/verify.rs
+++ b/guest/pvmfw/avb/src/verify.rs
@@ -45,6 +45,8 @@
     pub capabilities: Vec<Capability>,
     /// Rollback index of kernel.
     pub rollback_index: u64,
+    /// Page size of kernel, if present.
+    pub page_size: Option<usize>,
 }
 
 impl VerifiedBootData<'_> {
@@ -279,6 +281,7 @@
     let descriptors = vbmeta_image.descriptors()?;
     let hash_descriptors = HashDescriptors::get(&descriptors)?;
     let capabilities = verify_property_and_get_capabilities(&descriptors)?;
+    let page_size = None; // TODO(ptosi): Read from payload.
 
     if initrd.is_none() {
         hash_descriptors.verify_no_initrd()?;
@@ -289,6 +292,7 @@
             public_key: trusted_public_key,
             capabilities,
             rollback_index,
+            page_size,
         });
     }
 
@@ -309,5 +313,6 @@
         public_key: trusted_public_key,
         capabilities,
         rollback_index,
+        page_size,
     })
 }
diff --git a/guest/pvmfw/avb/tests/api_test.rs b/guest/pvmfw/avb/tests/api_test.rs
index 430c4b3..5bf777a 100644
--- a/guest/pvmfw/avb/tests/api_test.rs
+++ b/guest/pvmfw/avb/tests/api_test.rs
@@ -51,6 +51,7 @@
         &load_latest_initrd_normal()?,
         b"initrd_normal",
         DebugLevel::None,
+        None,
     )
 }
 
@@ -63,6 +64,7 @@
         salt,
         expected_rollback_index,
         vec![Capability::TrustySecurityVm],
+        None,
     )
 }
 
@@ -72,6 +74,7 @@
         &load_latest_initrd_debug()?,
         b"initrd_debug",
         DebugLevel::Full,
+        None,
     )
 }
 
@@ -93,6 +96,7 @@
         public_key: &public_key,
         capabilities: vec![],
         rollback_index: 0,
+        page_size: None,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
@@ -137,6 +141,7 @@
         public_key: &public_key,
         capabilities: vec![Capability::RemoteAttest],
         rollback_index: 0,
+        page_size: None,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
@@ -412,6 +417,7 @@
         public_key: &public_key,
         capabilities: vec![],
         rollback_index: 5,
+        page_size: None,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
     Ok(())
diff --git a/guest/pvmfw/avb/tests/utils.rs b/guest/pvmfw/avb/tests/utils.rs
index 61bfbf2..86efbba 100644
--- a/guest/pvmfw/avb/tests/utils.rs
+++ b/guest/pvmfw/avb/tests/utils.rs
@@ -114,6 +114,7 @@
     initrd: &[u8],
     initrd_salt: &[u8],
     expected_debug_level: DebugLevel,
+    page_size: Option<usize>,
 ) -> Result<()> {
     let public_key = load_trusted_public_key()?;
     let kernel = load_latest_signed_kernel()?;
@@ -133,6 +134,7 @@
         public_key: &public_key,
         capabilities,
         rollback_index: if cfg!(llpvm_changes) { 1 } else { 0 },
+        page_size,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
@@ -144,6 +146,7 @@
     salt: &[u8],
     expected_rollback_index: u64,
     capabilities: Vec<Capability>,
+    page_size: Option<usize>,
 ) -> Result<()> {
     let public_key = load_trusted_public_key()?;
     let verified_boot_data = verify_payload(
@@ -163,6 +166,7 @@
         public_key: &public_key,
         capabilities,
         rollback_index: expected_rollback_index,
+        page_size,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
diff --git a/guest/pvmfw/src/dice.rs b/guest/pvmfw/src/dice.rs
index b597309..6694881 100644
--- a/guest/pvmfw/src/dice.rs
+++ b/guest/pvmfw/src/dice.rs
@@ -200,6 +200,7 @@
         public_key: b"public key",
         capabilities: vec![],
         rollback_index: 42,
+        page_size: None,
     };
     const HASH: Hash = *b"sixtyfourbyteslongsentencearerarebutletsgiveitatrycantbethathard";
 
diff --git a/guest/pvmfw/src/entry.rs b/guest/pvmfw/src/entry.rs
index 343c2fc..e987dfd 100644
--- a/guest/pvmfw/src/entry.rs
+++ b/guest/pvmfw/src/entry.rs
@@ -122,13 +122,7 @@
 
     let config_entries = appended.get_entries();
 
-    let slices = memory::MemorySlices::new(
-        fdt,
-        payload,
-        payload_size,
-        config_entries.vm_dtbo,
-        config_entries.vm_ref_dt,
-    )?;
+    let slices = memory::MemorySlices::new(fdt, payload, payload_size)?;
 
     // This wrapper allows main() to be blissfully ignorant of platform details.
     let (next_bcc, debuggable_payload) = crate::main(
@@ -137,6 +131,8 @@
         slices.ramdisk,
         config_entries.bcc,
         config_entries.debug_policy,
+        config_entries.vm_dtbo,
+        config_entries.vm_ref_dt,
     )?;
     // Keep UART MMIO_GUARD-ed for debuggable payloads, to enable earlycon.
     let keep_uart = cfg!(debuggable_vms_improvements) && debuggable_payload;
diff --git a/guest/pvmfw/src/fdt.rs b/guest/pvmfw/src/fdt.rs
index 027f163..bfbd2e6 100644
--- a/guest/pvmfw/src/fdt.rs
+++ b/guest/pvmfw/src/fdt.rs
@@ -16,7 +16,6 @@
 
 use crate::bootargs::BootArgsIterator;
 use crate::device_assignment::{self, DeviceAssignmentInfo, VmDtbo};
-use crate::helpers::GUEST_PAGE_SIZE;
 use crate::Box;
 use crate::RebootReason;
 use alloc::collections::BTreeMap;
@@ -83,7 +82,7 @@
 
 /// Extract from /config the address range containing the pre-loaded kernel. Absence of /config is
 /// not an error.
-fn read_kernel_range_from(fdt: &Fdt) -> libfdt::Result<Option<Range<usize>>> {
+pub fn read_kernel_range_from(fdt: &Fdt) -> libfdt::Result<Option<Range<usize>>> {
     let addr = cstr!("kernel-address");
     let size = cstr!("kernel-size");
 
@@ -101,7 +100,7 @@
 
 /// Extract from /chosen the address range containing the pre-loaded ramdisk. Absence is not an
 /// error as there can be initrd-less VM.
-fn read_initrd_range_from(fdt: &Fdt) -> libfdt::Result<Option<Range<usize>>> {
+pub fn read_initrd_range_from(fdt: &Fdt) -> libfdt::Result<Option<Range<usize>>> {
     let start = cstr!("linux,initrd-start");
     let end = cstr!("linux,initrd-end");
 
@@ -147,7 +146,10 @@
 /// Reads and validates the memory range in the DT.
 ///
 /// Only one memory range is expected with the crosvm setup for now.
-fn read_and_validate_memory_range(fdt: &Fdt) -> Result<Range<usize>, RebootReason> {
+fn read_and_validate_memory_range(
+    fdt: &Fdt,
+    guest_page_size: usize,
+) -> Result<Range<usize>, RebootReason> {
     let mut memory = fdt.memory().map_err(|e| {
         error!("Failed to read memory range from DT: {e}");
         RebootReason::InvalidFdt
@@ -169,8 +171,8 @@
     }
 
     let size = range.len();
-    if size % GUEST_PAGE_SIZE != 0 {
-        error!("Memory size {:#x} is not a multiple of page size {:#x}", size, GUEST_PAGE_SIZE);
+    if size % guest_page_size != 0 {
+        error!("Memory size {:#x} is not a multiple of page size {:#x}", size, guest_page_size);
         return Err(RebootReason::InvalidFdt);
     }
 
@@ -854,16 +856,17 @@
 fn validate_swiotlb_info(
     swiotlb_info: &SwiotlbInfo,
     memory: &Range<usize>,
+    guest_page_size: usize,
 ) -> Result<(), RebootReason> {
     let size = swiotlb_info.size;
     let align = swiotlb_info.align;
 
-    if size == 0 || (size % GUEST_PAGE_SIZE) != 0 {
+    if size == 0 || (size % guest_page_size) != 0 {
         error!("Invalid swiotlb size {:#x}", size);
         return Err(RebootReason::InvalidFdt);
     }
 
-    if let Some(align) = align.filter(|&a| a % GUEST_PAGE_SIZE != 0) {
+    if let Some(align) = align.filter(|&a| a % guest_page_size != 0) {
         error!("Invalid swiotlb alignment {:#x}", align);
         return Err(RebootReason::InvalidFdt);
     }
@@ -989,7 +992,6 @@
 
 #[derive(Debug)]
 pub struct DeviceTreeInfo {
-    pub kernel_range: Option<Range<usize>>,
     pub initrd_range: Option<Range<usize>>,
     pub memory_range: Range<usize>,
     bootargs: Option<CString>,
@@ -1015,15 +1017,11 @@
 }
 
 pub fn sanitize_device_tree(
-    fdt: &mut [u8],
+    fdt: &mut Fdt,
     vm_dtbo: Option<&mut [u8]>,
     vm_ref_dt: Option<&[u8]>,
+    guest_page_size: usize,
 ) -> Result<DeviceTreeInfo, RebootReason> {
-    let fdt = Fdt::from_mut_slice(fdt).map_err(|e| {
-        error!("Failed to load FDT: {e}");
-        RebootReason::InvalidFdt
-    })?;
-
     let vm_dtbo = match vm_dtbo {
         Some(vm_dtbo) => Some(VmDtbo::from_mut_slice(vm_dtbo).map_err(|e| {
             error!("Failed to load VM DTBO: {e}");
@@ -1032,7 +1030,7 @@
         None => None,
     };
 
-    let info = parse_device_tree(fdt, vm_dtbo.as_deref())?;
+    let info = parse_device_tree(fdt, vm_dtbo.as_deref(), guest_page_size)?;
 
     fdt.clone_from(FDT_TEMPLATE).map_err(|e| {
         error!("Failed to instantiate FDT from the template DT: {e}");
@@ -1085,18 +1083,17 @@
     Ok(info)
 }
 
-fn parse_device_tree(fdt: &Fdt, vm_dtbo: Option<&VmDtbo>) -> Result<DeviceTreeInfo, RebootReason> {
-    let kernel_range = read_kernel_range_from(fdt).map_err(|e| {
-        error!("Failed to read kernel range from DT: {e}");
-        RebootReason::InvalidFdt
-    })?;
-
+fn parse_device_tree(
+    fdt: &Fdt,
+    vm_dtbo: Option<&VmDtbo>,
+    guest_page_size: usize,
+) -> Result<DeviceTreeInfo, RebootReason> {
     let initrd_range = read_initrd_range_from(fdt).map_err(|e| {
         error!("Failed to read initrd range from DT: {e}");
         RebootReason::InvalidFdt
     })?;
 
-    let memory_range = read_and_validate_memory_range(fdt)?;
+    let memory_range = read_and_validate_memory_range(fdt, guest_page_size)?;
 
     let bootargs = read_bootargs_from(fdt).map_err(|e| {
         error!("Failed to read bootargs from DT: {e}");
@@ -1149,7 +1146,7 @@
             error!("Swiotlb info missing from DT");
             RebootReason::InvalidFdt
         })?;
-    validate_swiotlb_info(&swiotlb_info, &memory_range)?;
+    validate_swiotlb_info(&swiotlb_info, &memory_range, guest_page_size)?;
 
     let device_assignment = match vm_dtbo {
         Some(vm_dtbo) => {
@@ -1194,7 +1191,6 @@
     })?;
 
     Ok(DeviceTreeInfo {
-        kernel_range,
         initrd_range,
         memory_range,
         bootargs,
diff --git a/guest/pvmfw/src/helpers.rs b/guest/pvmfw/src/helpers.rs
deleted file mode 100644
index 0552640..0000000
--- a/guest/pvmfw/src/helpers.rs
+++ /dev/null
@@ -1,19 +0,0 @@
-// 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.
-
-//! Miscellaneous helper functions.
-
-use vmbase::memory::SIZE_4KB;
-
-pub const GUEST_PAGE_SIZE: usize = SIZE_4KB;
diff --git a/guest/pvmfw/src/main.rs b/guest/pvmfw/src/main.rs
index bde03ff..4e5d82c 100644
--- a/guest/pvmfw/src/main.rs
+++ b/guest/pvmfw/src/main.rs
@@ -28,15 +28,13 @@
 mod exceptions;
 mod fdt;
 mod gpt;
-mod helpers;
 mod instance;
 mod memory;
 
 use crate::bcc::Bcc;
 use crate::dice::PartialInputs;
 use crate::entry::RebootReason;
-use crate::fdt::modify_for_next_stage;
-use crate::helpers::GUEST_PAGE_SIZE;
+use crate::fdt::{modify_for_next_stage, sanitize_device_tree};
 use crate::instance::EntryBody;
 use crate::instance::Error as InstanceError;
 use crate::instance::{get_recorded_entry, record_instance_entry};
@@ -54,21 +52,21 @@
 use pvmfw_embedded_key::PUBLIC_KEY;
 use vmbase::fdt::pci::{PciError, PciInfo};
 use vmbase::heap;
-use vmbase::memory::flush;
+use vmbase::memory::{flush, init_shared_pool, SIZE_4KB};
 use vmbase::rand;
 use vmbase::virtio::pci;
 
-const NEXT_BCC_SIZE: usize = GUEST_PAGE_SIZE;
-
 fn main(
-    fdt: &mut Fdt,
+    untrusted_fdt: &mut Fdt,
     signed_kernel: &[u8],
     ramdisk: Option<&[u8]>,
     current_bcc_handover: &[u8],
     mut debug_policy: Option<&[u8]>,
+    vm_dtbo: Option<&mut [u8]>,
+    vm_ref_dt: Option<&[u8]>,
 ) -> Result<(Range<usize>, bool), RebootReason> {
     info!("pVM firmware");
-    debug!("FDT: {:?}", fdt.as_ptr());
+    debug!("FDT: {:?}", untrusted_fdt.as_ptr());
     debug!("Signed kernel: {:?} ({:#x} bytes)", signed_kernel.as_ptr(), signed_kernel.len());
     debug!("AVB public key: addr={:?}, size={:#x} ({1})", PUBLIC_KEY.as_ptr(), PUBLIC_KEY.len());
     if let Some(rd) = ramdisk {
@@ -97,14 +95,6 @@
         debug_policy = None;
     }
 
-    // Set up PCI bus for VirtIO devices.
-    let pci_info = PciInfo::from_fdt(fdt).map_err(handle_pci_error)?;
-    debug!("PCI: {:#x?}", pci_info);
-    let mut pci_root = pci::initialize(pci_info).map_err(|e| {
-        error!("Failed to initialize PCI: {e}");
-        RebootReason::InternalError
-    })?;
-
     let verified_boot_data = verify_payload(signed_kernel, ramdisk, PUBLIC_KEY).map_err(|e| {
         error!("Failed to verify the payload: {e}");
         RebootReason::PayloadVerificationError
@@ -115,7 +105,23 @@
         info!("Please disregard any previous libavb ERROR about initrd_normal.");
     }
 
-    let next_bcc = heap::aligned_boxed_slice(NEXT_BCC_SIZE, GUEST_PAGE_SIZE).ok_or_else(|| {
+    let guest_page_size = verified_boot_data.page_size.unwrap_or(SIZE_4KB);
+    let fdt_info = sanitize_device_tree(untrusted_fdt, vm_dtbo, vm_ref_dt, guest_page_size)?;
+    let fdt = untrusted_fdt; // DT has now been sanitized.
+    let pci_info = PciInfo::from_fdt(fdt).map_err(handle_pci_error)?;
+    debug!("PCI: {:#x?}", pci_info);
+    // Set up PCI bus for VirtIO devices.
+    let mut pci_root = pci::initialize(pci_info).map_err(|e| {
+        error!("Failed to initialize PCI: {e}");
+        RebootReason::InternalError
+    })?;
+    init_shared_pool(fdt_info.swiotlb_info.fixed_range()).map_err(|e| {
+        error!("Failed to initialize shared pool: {e}");
+        RebootReason::InternalError
+    })?;
+
+    let next_bcc_size = guest_page_size;
+    let next_bcc = heap::aligned_boxed_slice(next_bcc_size, guest_page_size).ok_or_else(|| {
         error!("Failed to allocate the next-stage BCC");
         RebootReason::InternalError
     })?;
diff --git a/guest/pvmfw/src/memory.rs b/guest/pvmfw/src/memory.rs
index 35bfd3a..d2f63b5 100644
--- a/guest/pvmfw/src/memory.rs
+++ b/guest/pvmfw/src/memory.rs
@@ -15,7 +15,7 @@
 //! Low-level allocation and tracking of main memory.
 
 use crate::entry::RebootReason;
-use crate::fdt;
+use crate::fdt::{read_initrd_range_from, read_kernel_range_from};
 use core::num::NonZeroUsize;
 use core::slice;
 use log::debug;
@@ -24,7 +24,7 @@
 use log::warn;
 use vmbase::{
     layout::crosvm,
-    memory::{init_shared_pool, map_data, map_rodata, resize_available_memory},
+    memory::{map_data, map_rodata, resize_available_memory},
 };
 
 pub(crate) struct MemorySlices<'a> {
@@ -34,13 +34,7 @@
 }
 
 impl<'a> MemorySlices<'a> {
-    pub fn new(
-        fdt: usize,
-        kernel: usize,
-        kernel_size: usize,
-        vm_dtbo: Option<&mut [u8]>,
-        vm_ref_dt: Option<&[u8]>,
-    ) -> Result<Self, RebootReason> {
+    pub fn new(fdt: usize, kernel: usize, kernel_size: usize) -> Result<Self, RebootReason> {
         let fdt_size = NonZeroUsize::new(crosvm::FDT_MAX_SIZE).unwrap();
         // TODO - Only map the FDT as read-only, until we modify it right before jump_to_payload()
         // e.g. by generating a DTBO for a template DT in main() and, on return, re-map DT as RW,
@@ -51,44 +45,39 @@
         })?;
 
         // SAFETY: map_data validated the range to be in main memory, mapped, and not overlap.
-        let fdt = unsafe { slice::from_raw_parts_mut(fdt as *mut u8, fdt_size.into()) };
-
-        let info = fdt::sanitize_device_tree(fdt, vm_dtbo, vm_ref_dt)?;
-        let fdt = libfdt::Fdt::from_mut_slice(fdt).map_err(|e| {
-            error!("Failed to load sanitized FDT: {e}");
+        let untrusted_fdt = unsafe { slice::from_raw_parts_mut(fdt as *mut u8, fdt_size.into()) };
+        let untrusted_fdt = libfdt::Fdt::from_mut_slice(untrusted_fdt).map_err(|e| {
+            error!("Failed to load input FDT: {e}");
             RebootReason::InvalidFdt
         })?;
-        debug!("Fdt passed validation!");
 
-        let memory_range = info.memory_range;
+        let memory_range = untrusted_fdt.first_memory_range().map_err(|e| {
+            error!("Failed to read memory range from DT: {e}");
+            RebootReason::InvalidFdt
+        })?;
         debug!("Resizing MemoryTracker to range {memory_range:#x?}");
         resize_available_memory(&memory_range).map_err(|e| {
             error!("Failed to use memory range value from DT: {memory_range:#x?}: {e}");
             RebootReason::InvalidFdt
         })?;
 
-        init_shared_pool(info.swiotlb_info.fixed_range()).map_err(|e| {
-            error!("Failed to initialize shared pool: {e}");
-            RebootReason::InternalError
+        let kernel_range = read_kernel_range_from(untrusted_fdt).map_err(|e| {
+            error!("Failed to read kernel range: {e}");
+            RebootReason::InvalidFdt
         })?;
-
-        let (kernel_start, kernel_size) = if let Some(r) = info.kernel_range {
-            let size = r.len().try_into().map_err(|_| {
-                error!("Invalid kernel size: {:#x}", r.len());
-                RebootReason::InternalError
-            })?;
-            (r.start, size)
+        let (kernel_start, kernel_size) = if let Some(r) = kernel_range {
+            (r.start, r.len())
         } else if cfg!(feature = "legacy") {
             warn!("Failed to find the kernel range in the DT; falling back to legacy ABI");
-            let size = NonZeroUsize::new(kernel_size).ok_or_else(|| {
-                error!("Invalid kernel size: {kernel_size:#x}");
-                RebootReason::InvalidPayload
-            })?;
-            (kernel, size)
+            (kernel, kernel_size)
         } else {
             error!("Failed to locate the kernel from the DT");
             return Err(RebootReason::InvalidPayload);
         };
+        let kernel_size = kernel_size.try_into().map_err(|_| {
+            error!("Invalid kernel size: {kernel_size:#x}");
+            RebootReason::InvalidPayload
+        })?;
 
         map_rodata(kernel_start, kernel_size).map_err(|e| {
             error!("Failed to map kernel range: {e}");
@@ -99,7 +88,11 @@
         // SAFETY: map_rodata validated the range to be in main memory, mapped, and not overlap.
         let kernel = unsafe { slice::from_raw_parts(kernel, kernel_size.into()) };
 
-        let ramdisk = if let Some(r) = info.initrd_range {
+        let initrd_range = read_initrd_range_from(untrusted_fdt).map_err(|e| {
+            error!("Failed to read initrd range: {e}");
+            RebootReason::InvalidFdt
+        })?;
+        let ramdisk = if let Some(r) = initrd_range {
             debug!("Located ramdisk at {r:?}");
             let ramdisk_size = r.len().try_into().map_err(|_| {
                 error!("Invalid ramdisk size: {:#x}", r.len());
@@ -118,6 +111,6 @@
             None
         };
 
-        Ok(Self { fdt, kernel, ramdisk })
+        Ok(Self { fdt: untrusted_fdt, kernel, ramdisk })
     }
 }