Add functions to validate a DPE Instance

This commit adds a function that validates a DPEInstance's
context array, to ensure it represents a properly-formed
tree. We also add functions to retrieve the current tci
from a context, and get the index of a context whose
locality and tci_type are equal to some parameters. These
functions will be called from caliptra rt-fw when it validates
the DPE instance loaded from SRAM upon a warm reset.

dpe/src/dpe_instance.rs

@@ -389,6 +389,83 @@ impl DpeInstance {
 
         hasher.finish().map_err(|_| DpeErrorCode::HashError)
     }
+
+    /// Retrieve index of the root context based on the tci_type and locality of
+    /// the context matching the inputs
+    ///
+    /// # Arguments
+    ///
+    /// * `idx` - desired index of the root context
+    /// * `locality` - desired locality of the root context
+    pub fn get_root_idx(&self, tci_type: u32, locality: u32) -> Result<usize, DpeErrorCode> {
+        // find indices of all non-inactive contexts matching tci_type and locality
+        let mut tci_nodes = self.contexts.iter().enumerate().filter(|(_, context)| {
+            context.state != ContextState::Inactive
+                && context.locality == locality
+                && context.tci.tci_type == tci_type
+        });
+        let tci = tci_nodes.next();
+        // ensure there is exactly 1 match
+        if tci.is_none() || tci_nodes.next().is_some() {
+            return Err(DpeErrorCode::InternalError);
+        }
+        Ok(tci.unwrap().0)
+    }
+
+    /// Retrieve tci_current of a context
+    ///
+    /// # Arguments
+    ///
+    /// * `idx` - index of the context
+    pub fn get_current_tci(&self, idx: usize) -> Result<TciMeasurement, DpeErrorCode> {
+        if idx >= MAX_HANDLES {
+            return Err(DpeErrorCode::InternalError);
+        }
+        Ok(self.contexts[idx].tci.tci_current)
+    }
+
+    /// Determines if the context array represents a valid tree by checking that
+    /// there is only 1 connected component and that all nodes lead up to
+    /// the root node.
+    ///
+    /// # Arguments
+    ///
+    /// * `root_idx` - The index of the root context
+    pub fn validate_context_tree(&self, root_idx: usize) -> bool {
+        let mut seen = [false; MAX_HANDLES];
+
+        // dfs from the root node and try to discover cycles
+        if self.detect_cycle(root_idx, &mut seen) {
+            return false;
+        }
+
+        for i in 0..MAX_HANDLES {
+            // If a node was not seen when doing a dfs from the root, there must be multiple
+            // connected components or the root is not actually the root
+            if i != root_idx && self.contexts[i].state != ContextState::Inactive && !seen[i] {
+                return false;
+            }
+        }
+        true
+    }
+
+    fn detect_cycle(&self, curr_idx: usize, seen: &mut [bool; MAX_HANDLES]) -> bool {
+        // if the current node was already visited we have a cycle
+        if curr_idx >= MAX_HANDLES
+            || self.contexts[curr_idx].state == ContextState::Inactive
+            || seen[curr_idx]
+        {
+            return true;
+        }
+        seen[curr_idx] = true;
+        // dfs on all child nodes
+        for child_idx in flags_iter(self.contexts[curr_idx].children, MAX_HANDLES) {
+            if child_idx >= MAX_HANDLES || self.detect_cycle(child_idx, seen) {
+                return true;
+            }
+        }
+        false
+    }
 }
 
 /// Iterate over all of the bits set to 1 in a u32. Each iteration returns the bit index 0 being the
@@ -789,4 +866,79 @@ pub mod tests {
             .unwrap();
         assert_eq!(answer, cdi_with_internal_input_dice)
     }
+
+    #[test]
+    fn test_get_root_index() {
+        let mut env = DpeEnv::<TestTypes> {
+            crypto: OpensslCrypto::new(),
+            platform: DefaultPlatform,
+        };
+        let mut dpe = DpeInstance::new(&mut env, SUPPORT).unwrap();
+
+        DeriveChildCmd {
+            handle: ContextHandle::default(),
+            data: [0; DPE_PROFILE.get_hash_size()],
+            flags: DeriveChildFlags::MAKE_DEFAULT,
+            tci_type: 1u32,
+            target_locality: 0,
+        }
+        .execute(&mut dpe, &mut env, TEST_LOCALITIES[0])
+        .unwrap();
+
+        // Finds index of context containing matching tci_type and locality
+        assert_eq!(dpe.get_root_idx(1, 0), Ok(1));
+        // No match should return internal error
+        assert_eq!(dpe.get_root_idx(1, 1), Err(DpeErrorCode::InternalError));
+
+        DeriveChildCmd {
+            handle: ContextHandle::default(),
+            data: [0; DPE_PROFILE.get_hash_size()],
+            flags: DeriveChildFlags::RETAIN_PARENT,
+            tci_type: 1u32,
+            target_locality: 0,
+        }
+        .execute(&mut dpe, &mut env, TEST_LOCALITIES[0])
+        .unwrap();
+
+        // multiple tci_type + locality matches should return InternalError
+        assert_eq!(dpe.get_root_idx(1, 0), Err(DpeErrorCode::InternalError));
+    }
+
+    #[test]
+    fn test_validate_context_tree() {
+        let mut env = DpeEnv::<TestTypes> {
+            crypto: OpensslCrypto::new(),
+            platform: DefaultPlatform,
+        };
+        let mut dpe = DpeInstance::new(&mut env, SUPPORT).unwrap();
+
+        dpe.contexts[0].state = ContextState::Active;
+        dpe.contexts[0].children = 0b100;
+        dpe.contexts[1].state = ContextState::Active;
+        dpe.contexts[1].children = 0b100;
+        dpe.contexts[2].state = ContextState::Active;
+        // validation fails on graph where child has multiple parents
+        assert_eq!(dpe.validate_context_tree(0), false);
+
+        dpe.contexts[0].children = 0b10;
+        // validation passes on a tree in the shape of a linked-list
+        assert_eq!(dpe.validate_context_tree(0), true);
+
+        dpe.contexts[2].children = 0b1;
+        // validation fails on circle graph
+        assert_eq!(dpe.validate_context_tree(0), false);
+
+        dpe.contexts[0].children |= 0b100;
+        dpe.contexts[1].children = 0;
+        dpe.contexts[2].children = 0;
+        // validation passes on a complete binary tree of size 2
+        assert_eq!(dpe.validate_context_tree(0), true);
+
+        dpe.contexts[10].state = ContextState::Active;
+        dpe.contexts[10].children = 1 << 11 | 1 << 12;
+        dpe.contexts[11].state = ContextState::Active;
+        dpe.contexts[12].state = ContextState::Active;
+        // validation fails on a graph with multiple connected components
+        assert_eq!(dpe.validate_context_tree(0), false);
+    }
 }