pvmfw: Support #iommu-cells = <1>

Bug: 277993056
Test: atest libpvmfw.device_assignment.test, launch protected VM
Change-Id: Ie71a8f5759976b39245b74343414752caa4aca1b
diff --git a/pvmfw/src/device_assignment.rs b/pvmfw/src/device_assignment.rs
index 3f84a8d..60bf21c 100644
--- a/pvmfw/src/device_assignment.rs
+++ b/pvmfw/src/device_assignment.rs
@@ -55,6 +55,8 @@
     InvalidInterrupts,
     /// Invalid <iommus>
     InvalidIommus,
+    /// Invalid pvIOMMU node
+    InvalidPvIommu,
     /// Too many pvIOMMU
     TooManyPvIommu,
     /// Duplicated pvIOMMU IDs exist
@@ -83,6 +85,7 @@
             ),
             Self::InvalidInterrupts => write!(f, "Invalid <interrupts>"),
             Self::InvalidIommus => write!(f, "Invalid <iommus>"),
+            Self::InvalidPvIommu => write!(f, "Invalid pvIOMMU node"),
             Self::TooManyPvIommu => write!(
                 f,
                 "Too many pvIOMMU node. Insufficient pre-populated pvIOMMUs in platform DT"
@@ -195,11 +198,22 @@
 
 impl PvIommu {
     fn parse(node: &FdtNode) -> Result<Self> {
-        let id = node.getprop_u32(cstr!("id"))?.ok_or(DeviceAssignmentError::InvalidIommus)?;
+        let iommu_cells = node
+            .getprop_u32(cstr!("#iommu-cells"))?
+            .ok_or(DeviceAssignmentError::InvalidPvIommu)?;
+        // Ensures <#iommu-cells> = 1. It means that `<iommus>` entry contains pair of
+        // (pvIOMMU ID, vSID)
+        if iommu_cells != 1 {
+            return Err(DeviceAssignmentError::InvalidPvIommu);
+        }
+        let id = node.getprop_u32(cstr!("id"))?.ok_or(DeviceAssignmentError::InvalidPvIommu)?;
         Ok(Self { id })
     }
 }
 
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
+struct Vsid(u32);
+
 /// Assigned device information parsed from crosvm DT.
 /// Keeps everything in the owned data because underlying FDT will be reused for platform DT.
 #[derive(Debug, Eq, PartialEq)]
@@ -212,8 +226,8 @@
     reg: Vec<u8>,
     // <interrupts> property from the crosvm DT
     interrupts: Vec<u8>,
-    // Parsed <iommus> property from the crosvm DT.
-    iommus: Vec<PvIommu>,
+    // Parsed <iommus> property from the crosvm DT. Tuple of PvIommu and vSID.
+    iommus: Vec<(PvIommu, Vsid)>,
 }
 
 impl AssignedDeviceInfo {
@@ -233,17 +247,26 @@
     }
 
     // TODO(b/277993056): Also validate /__local_fixups__ to ensure that <iommus> has phandle.
-    // TODO(b/277993056): Also keep vSID.
-    // TODO(b/277993056): Validate #iommu-cells values.
-    fn parse_iommus(node: &FdtNode, pviommus: &BTreeMap<Phandle, PvIommu>) -> Result<Vec<PvIommu>> {
+    fn parse_iommus(
+        node: &FdtNode,
+        pviommus: &BTreeMap<Phandle, PvIommu>,
+    ) -> Result<Vec<(PvIommu, Vsid)>> {
         let mut iommus = vec![];
-        let Some(cells) = node.getprop_cells(cstr!("iommus"))? else {
+        let Some(mut cells) = node.getprop_cells(cstr!("iommus"))? else {
             return Ok(iommus);
         };
-        for cell in cells {
+        while let Some(cell) = cells.next() {
+            // Parse pvIOMMU ID
             let phandle = Phandle::try_from(cell).or(Err(DeviceAssignmentError::InvalidIommus))?;
             let pviommu = pviommus.get(&phandle).ok_or(DeviceAssignmentError::InvalidIommus)?;
-            iommus.push(*pviommu)
+
+            // Parse vSID
+            let Some(cell) = cells.next() else {
+                return Err(DeviceAssignmentError::InvalidIommus);
+            };
+            let vsid = Vsid(cell);
+
+            iommus.push((*pviommu, vsid));
         }
         Ok(iommus)
     }
@@ -275,15 +298,12 @@
         let mut dst = fdt.node_mut(&self.node_path)?.unwrap();
         dst.setprop(cstr!("reg"), &self.reg)?;
         dst.setprop(cstr!("interrupts"), &self.interrupts)?;
-
-        let iommus: Vec<u8> = self
-            .iommus
-            .iter()
-            .flat_map(|pviommu| {
-                let phandle = pviommu_phandles.get(pviommu).unwrap();
-                u32::from(*phandle).to_be_bytes()
-            })
-            .collect();
+        let mut iommus = Vec::with_capacity(8 * self.iommus.len());
+        for (pviommu, vsid) in &self.iommus {
+            let phandle = pviommu_phandles.get(pviommu).unwrap();
+            iommus.extend_from_slice(&u32::from(*phandle).to_be_bytes());
+            iommus.extend_from_slice(&vsid.0.to_be_bytes());
+        }
         dst.setprop(cstr!("iommus"), &iommus)?;
 
         Ok(())
@@ -303,7 +323,6 @@
     /// Parses pvIOMMUs in fdt
     // Note: This will validate pvIOMMU ids' uniqueness, even when unassigned.
     fn parse_pviommus(fdt: &Fdt) -> Result<BTreeMap<Phandle, PvIommu>> {
-        // TODO(b/277993056): Validated `<#iommu-cells>`.
         let mut pviommus = BTreeMap::new();
         for compatible in fdt.compatible_nodes(Self::PVIOMMU_COMPATIBLE)? {
             let Some(phandle) = compatible.get_phandle()? else {
@@ -460,9 +479,10 @@
                 .getprop(cstr!("interrupts"))?
                 .ok_or(DeviceAssignmentError::InvalidInterrupts)?;
             let mut iommus = vec![];
-            if let Some(cells) = node.getprop_cells(cstr!("iommus"))? {
-                for cell in cells {
-                    let phandle = Phandle::try_from(cell)?;
+            if let Some(mut cells) = node.getprop_cells(cstr!("iommus"))? {
+                while let Some(pviommu_id) = cells.next() {
+                    // pvIOMMU id
+                    let phandle = Phandle::try_from(pviommu_id)?;
                     let pviommu = fdt
                         .node_with_phandle(phandle)?
                         .ok_or(DeviceAssignmentError::InvalidIommus)?;
@@ -474,6 +494,12 @@
                         .getprop_u32(cstr!("id"))?
                         .ok_or(DeviceAssignmentError::InvalidIommus)?;
                     iommus.push(id);
+
+                    // vSID
+                    let Some(vsid) = cells.next() else {
+                        return Err(DeviceAssignmentError::InvalidIommus);
+                    };
+                    iommus.push(vsid);
                 }
             }
             Ok(Self { path: path.into(), reg: reg.into(), interrupts: interrupts.into(), iommus })
@@ -630,7 +656,7 @@
             path: CString::new("/rng").unwrap(),
             reg: into_fdt_prop(vec![0x0, 0x9, 0x0, 0xFF]),
             interrupts: into_fdt_prop(vec![0x0, 0xF, 0x4]),
-            iommus: vec![0x4],
+            iommus: vec![0x4, 0xFF0],
         };
 
         let node = AssignedDeviceNode::parse(platform_dt, &expected.path);
@@ -665,13 +691,13 @@
                 path: CString::new("/rng").unwrap(),
                 reg: into_fdt_prop(vec![0x0, 0x9, 0x0, 0xFF]),
                 interrupts: into_fdt_prop(vec![0x0, 0xF, 0x4]),
-                iommus: vec![0x4, 0x9],
+                iommus: vec![0x4, 0xFF0, 0x9, 0xFF1],
             },
             AssignedDeviceNode {
                 path: CString::new("/light").unwrap(),
                 reg: into_fdt_prop(vec![0x100, 0x9]),
                 interrupts: into_fdt_prop(vec![0x0, 0xF, 0x5]),
-                iommus: vec![0x40, 0x50, 0x60],
+                iommus: vec![0x40, 0xFFA, 0x50, 0xFFB, 0x60, 0xFFC],
             },
         ];
 
@@ -708,13 +734,13 @@
                 path: CString::new("/rng").unwrap(),
                 reg: into_fdt_prop(vec![0x0, 0x9, 0x0, 0xFF]),
                 interrupts: into_fdt_prop(vec![0x0, 0xF, 0x4]),
-                iommus: vec![0x4, 0x9],
+                iommus: vec![0x4, 0xFF0, 0x9, 0xFF1],
             },
             AssignedDeviceNode {
                 path: CString::new("/light").unwrap(),
                 reg: into_fdt_prop(vec![0x100, 0x9]),
                 interrupts: into_fdt_prop(vec![0x0, 0xF, 0x5]),
-                iommus: vec![0x9, 0x40],
+                iommus: vec![0x9, 0xFF1, 0x40, 0xFFA],
             },
         ];
 
diff --git a/pvmfw/testdata/test_pvmfw_devices_with_iommu_id_conflict.dts b/pvmfw/testdata/test_pvmfw_devices_with_iommu_id_conflict.dts
index f0a7162..199a5ce 100644
--- a/pvmfw/testdata/test_pvmfw_devices_with_iommu_id_conflict.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_with_iommu_id_conflict.dts
@@ -54,13 +54,13 @@
     pviommu_0: pviommu0 {
         compatible = "pkvm,pviommu";
         id = <0x4>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 
     pviommu_1: pviommu1 {
         compatible = "pkvm,pviommu";
         id = <0x9>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 
     light@70000000 {
@@ -73,12 +73,12 @@
     pviommu_a: pviommua {
         compatible = "pkvm,pviommu";
         id = <0x40>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 
     pviommu_b: pviommub {
         compatible = "pkvm,pviommu";
         id = <0x9>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 };
diff --git a/pvmfw/testdata/test_pvmfw_devices_with_iommu_sharing.dts b/pvmfw/testdata/test_pvmfw_devices_with_iommu_sharing.dts
index d6952fa..4906064 100644
--- a/pvmfw/testdata/test_pvmfw_devices_with_iommu_sharing.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_with_iommu_sharing.dts
@@ -48,31 +48,31 @@
         interrupts = <0x0 0xF 0x4>;
         google,eh,ignore-gctrl-reset;
         status = "okay";
-        iommus = <&pviommu_0>, <&pviommu_1>;
+        iommus = <&pviommu_0 0xFF0>, <&pviommu_1 0xFF1>;
     };
 
     light@70000000 {
         compatible = "android,light";
         reg = <0x100 0x9>;
         interrupts = <0x0 0xF 0x5>;
-        iommus = <&pviommu_1>, <&pviommu_a>;
+        iommus = <&pviommu_1 0xFF1>, <&pviommu_a 0xFFA>;
     };
 
     pviommu_0: pviommu0 {
         compatible = "pkvm,pviommu";
         id = <0x4>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 
     pviommu_1: pviommu1 {
         compatible = "pkvm,pviommu";
         id = <0x9>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 
     pviommu_a: pviommua {
         compatible = "pkvm,pviommu";
         id = <0x40>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 };
diff --git a/pvmfw/testdata/test_pvmfw_devices_with_multiple_devices_iommus.dts b/pvmfw/testdata/test_pvmfw_devices_with_multiple_devices_iommus.dts
index 2609c45..959cd23 100644
--- a/pvmfw/testdata/test_pvmfw_devices_with_multiple_devices_iommus.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_with_multiple_devices_iommus.dts
@@ -48,43 +48,43 @@
         interrupts = <0x0 0xF 0x4>;
         google,eh,ignore-gctrl-reset;
         status = "okay";
-        iommus = <&pviommu_0>, <&pviommu_1>;
+        iommus = <&pviommu_0 0xFF0>, <&pviommu_1 0xFF1>;
     };
 
     pviommu_0: pviommu0 {
         compatible = "pkvm,pviommu";
         id = <0x4>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 
     pviommu_1: pviommu1 {
         compatible = "pkvm,pviommu";
         id = <0x9>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 
     light@70000000 {
         compatible = "android,light";
         reg = <0x100 0x9>;
         interrupts = <0x0 0xF 0x5>;
-        iommus = <&pviommu_a>, <&pviommu_b>, <&pviommu_c>;
+        iommus = <&pviommu_a 0xFFA>, <&pviommu_b 0xFFB>, <&pviommu_c 0xFFC>;
     };
 
     pviommu_a: pviommua {
         compatible = "pkvm,pviommu";
         id = <0x40>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 
     pviommu_b: pviommub {
         compatible = "pkvm,pviommu";
         id = <0x50>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 
     pviommu_c: pviommuc {
         compatible = "pkvm,pviommu";
         id = <0x60>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 };
diff --git a/pvmfw/testdata/test_pvmfw_devices_with_rng_iommu.dts b/pvmfw/testdata/test_pvmfw_devices_with_rng_iommu.dts
index 6a5068c..8c04b39 100644
--- a/pvmfw/testdata/test_pvmfw_devices_with_rng_iommu.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_with_rng_iommu.dts
@@ -48,12 +48,12 @@
         interrupts = <0x0 0xF 0x4>;
         google,eh,ignore-gctrl-reset;
         status = "okay";
-        iommus = <&pviommu_0>;
+        iommus = <&pviommu_0 0xFF0>;
     };
 
     pviommu_0: pviommu0 {
         compatible = "pkvm,pviommu";
         id = <0x4>;
-        #iommu-cells = <0>;
+        #iommu-cells = <1>;
     };
 };