Merge "Check for null before dereference" into main
diff --git a/pvmfw/src/device_assignment.rs b/pvmfw/src/device_assignment.rs
index 9c3e566..86ad0f0 100644
--- a/pvmfw/src/device_assignment.rs
+++ b/pvmfw/src/device_assignment.rs
@@ -73,6 +73,8 @@
     DuplicatedIommuIds,
     /// Duplicated pvIOMMU IDs exist
     DuplicatedPvIommuIds,
+    /// Unsupported path format. Only supports full path.
+    UnsupportedPathFormat,
     /// Unsupported overlay target syntax. Only supports <target-path> with full path.
     UnsupportedOverlayTarget,
     /// Unsupported PhysIommu,
@@ -120,6 +122,9 @@
             Self::DuplicatedPvIommuIds => {
                 write!(f, "Duplicated pvIOMMU IDs exist. IDs must unique among iommu node")
             }
+            Self::UnsupportedPathFormat => {
+                write!(f, "Unsupported UnsupportedPathFormat. Only supports full path")
+            }
             Self::UnsupportedOverlayTarget => {
                 write!(f, "Unsupported overlay target. Only supports 'target-path = \"/\"'")
             }
@@ -140,6 +145,67 @@
 
 pub type Result<T> = core::result::Result<T, DeviceAssignmentError>;
 
+#[derive(Clone, Default, Ord, PartialOrd, Eq, PartialEq)]
+pub struct DtPathTokens<'a> {
+    tokens: Vec<&'a [u8]>,
+}
+
+impl<'a> fmt::Debug for DtPathTokens<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut list = f.debug_list();
+        for token in &self.tokens {
+            let mut bytes = token.to_vec();
+            bytes.push(b'\0');
+            match CString::from_vec_with_nul(bytes) {
+                Ok(string) => list.entry(&string),
+                Err(_) => list.entry(token),
+            };
+        }
+        list.finish()
+    }
+}
+
+impl<'a> DtPathTokens<'a> {
+    fn new(path: &'a CStr) -> Result<Self> {
+        if path.to_bytes().first() != Some(&b'/') {
+            return Err(DeviceAssignmentError::UnsupportedPathFormat);
+        }
+        let tokens: Vec<_> = path
+            .to_bytes()
+            .split(|char| *char == b'/')
+            .filter(|&component| !component.is_empty())
+            .collect();
+        Ok(Self { tokens })
+    }
+
+    fn to_overlay_target_path(&self) -> Result<Self> {
+        if !self.is_overlayable_node() {
+            return Err(DeviceAssignmentError::InvalidDtbo);
+        }
+        Ok(Self { tokens: self.tokens.as_slice()[2..].to_vec() })
+    }
+
+    fn to_cstring(&self) -> CString {
+        if self.tokens.is_empty() {
+            return CString::new(*b"/\0").unwrap();
+        }
+
+        let size = self.tokens.iter().fold(0, |sum, token| sum + token.len() + 1);
+        let mut path = Vec::with_capacity(size + 1);
+        for token in &self.tokens {
+            path.push(b'/');
+            path.extend_from_slice(token);
+        }
+        path.push(b'\0');
+
+        CString::from_vec_with_nul(path).unwrap()
+    }
+
+    fn is_overlayable_node(&self) -> bool {
+        self.tokens.get(1) == Some(&&b"__overlay__"[..])
+    }
+}
+
 /// Represents VM DTBO
 #[repr(transparent)]
 pub struct VmDtbo(Fdt);
@@ -179,14 +245,9 @@
     // node and/or properties. The enforcement is for compatibility reason.
     fn locate_overlay_target_path(
         &self,
-        dtbo_node_path: &CStr,
+        dtbo_node_path: &DtPathTokens,
         dtbo_node: &FdtNode,
     ) -> Result<CString> {
-        let dtbo_node_path_bytes = dtbo_node_path.to_bytes();
-        if dtbo_node_path_bytes.first() != Some(&b'/') {
-            return Err(DeviceAssignmentError::UnsupportedOverlayTarget);
-        }
-
         let fragment_node = dtbo_node.supernode_at_depth(1)?;
         let target_path = fragment_node
             .getprop_str(cstr!("target-path"))?
@@ -195,22 +256,8 @@
             return Err(DeviceAssignmentError::UnsupportedOverlayTarget);
         }
 
-        let mut components = dtbo_node_path_bytes
-            .split(|char| *char == b'/')
-            .filter(|&component| !component.is_empty())
-            .skip(1);
-        let overlay_node_name = components.next();
-        if overlay_node_name != Some(b"__overlay__") {
-            return Err(DeviceAssignmentError::InvalidDtbo);
-        }
-        let mut overlaid_path = Vec::with_capacity(dtbo_node_path_bytes.len());
-        for component in components {
-            overlaid_path.push(b'/');
-            overlaid_path.extend_from_slice(component);
-        }
-        overlaid_path.push(b'\0');
-
-        Ok(CString::from_vec_with_nul(overlaid_path).unwrap())
+        let overlaid_path = dtbo_node_path.to_overlay_target_path()?;
+        Ok(overlaid_path.to_cstring())
     }
 
     fn parse_physical_iommus(physical_node: &FdtNode) -> Result<BTreeMap<Phandle, PhysIommu>> {
@@ -281,15 +328,17 @@
         let phys_iommus = Self::parse_physical_iommus(&physical_node)?;
         Self::parse_physical_devices_with_iommus(&physical_node, &phys_iommus)
     }
-}
 
-fn is_overlayable_node(dtbo_path: &CStr) -> bool {
-    dtbo_path
-        .to_bytes()
-        .split(|char| *char == b'/')
-        .filter(|&component| !component.is_empty())
-        .nth(1)
-        .map_or(false, |name| name == b"__overlay__")
+    fn node(&self, path: &DtPathTokens) -> Result<Option<FdtNode>> {
+        let mut node = self.as_ref().root();
+        for token in &path.tokens {
+            let Some(subnode) = node.subnode_with_name_bytes(token)? else {
+                return Ok(None);
+            };
+            node = subnode;
+        }
+        Ok(Some(node))
+    }
 }
 
 fn filter_dangling_symbols(fdt: &mut Fdt) -> Result<()> {
@@ -458,8 +507,6 @@
 struct AssignedDeviceInfo {
     // Node path of assigned device (e.g. "/rng")
     node_path: CString,
-    // DTBO node path of the assigned device (e.g. "/fragment@rng/__overlay__/rng")
-    dtbo_node_path: CString,
     // <reg> property from the crosvm DT
     reg: Vec<DeviceReg>,
     // <interrupts> property from the crosvm DT
@@ -568,13 +615,13 @@
     fn parse(
         fdt: &Fdt,
         vm_dtbo: &VmDtbo,
-        dtbo_node_path: &CStr,
+        dtbo_node_path: &DtPathTokens,
         physical_devices: &BTreeMap<Phandle, PhysicalDeviceInfo>,
         pviommus: &BTreeMap<Phandle, PvIommu>,
         hypervisor: &dyn DeviceAssigningHypervisor,
     ) -> Result<Option<Self>> {
         let dtbo_node =
-            vm_dtbo.as_ref().node(dtbo_node_path)?.ok_or(DeviceAssignmentError::InvalidSymbols)?;
+            vm_dtbo.node(dtbo_node_path)?.ok_or(DeviceAssignmentError::InvalidSymbols)?;
         let node_path = vm_dtbo.locate_overlay_target_path(dtbo_node_path, &dtbo_node)?;
 
         let Some(node) = fdt.node(&node_path)? else { return Ok(None) };
@@ -595,7 +642,7 @@
         let iommus = Self::parse_iommus(&node, pviommus)?;
         Self::validate_iommus(&iommus, &physical_device.iommus, hypervisor)?;
 
-        Ok(Some(Self { node_path, dtbo_node_path: dtbo_node_path.into(), reg, interrupts, iommus }))
+        Ok(Some(Self { node_path, reg, interrupts, iommus }))
     }
 
     fn patch(&self, fdt: &mut Fdt, pviommu_phandles: &BTreeMap<PvIommu, Phandle>) -> Result<()> {
@@ -681,13 +728,14 @@
             let symbol_prop_value = symbol_prop.value()?;
             let dtbo_node_path = CStr::from_bytes_with_nul(symbol_prop_value)
                 .or(Err(DeviceAssignmentError::InvalidSymbols))?;
-            if !is_overlayable_node(dtbo_node_path) {
+            let dtbo_node_path = DtPathTokens::new(dtbo_node_path)?;
+            if !dtbo_node_path.is_overlayable_node() {
                 continue;
             }
             let assigned_device = AssignedDeviceInfo::parse(
                 fdt,
                 vm_dtbo,
-                dtbo_node_path,
+                &dtbo_node_path,
                 &physical_devices,
                 &pviommus,
                 hypervisor,
@@ -695,7 +743,7 @@
             if let Some(assigned_device) = assigned_device {
                 assigned_devices.push(assigned_device);
             } else {
-                filtered_dtbo_paths.push(dtbo_node_path.into());
+                filtered_dtbo_paths.push(dtbo_node_path.to_cstring());
             }
         }
         if assigned_devices.is_empty() {
@@ -922,7 +970,6 @@
 
         let expected = [AssignedDeviceInfo {
             node_path: CString::new("/bus0/backlight").unwrap(),
-            dtbo_node_path: cstr!("/fragment@0/__overlay__/bus0/backlight").into(),
             reg: vec![[0x9, 0xFF].into()],
             interrupts: into_fdt_prop(vec![0x0, 0xF, 0x4]),
             iommus: vec![],
@@ -946,7 +993,6 @@
 
         let expected = [AssignedDeviceInfo {
             node_path: CString::new("/rng").unwrap(),
-            dtbo_node_path: cstr!("/fragment@0/__overlay__/rng").into(),
             reg: vec![[0x9, 0xFF].into()],
             interrupts: into_fdt_prop(vec![0x0, 0xF, 0x4]),
             iommus: vec![(PvIommu { id: 0x4 }, Vsid(0xFF0))],