User iterator instead of temporary buffer for CertifyKey TCB

dpe/src/commands/certify_key.rs

@@ -4,9 +4,8 @@ use crate::{
     context::ContextHandle,
     dpe_instance::{DpeEnv, DpeInstance, DpeTypes},
     response::{CertifyKeyResp, DpeErrorCode, Response, ResponseHdr},
-    tci::TciNodeData,
     x509::{CertWriter, DirectoryString, MeasurementData, Name},
-    DPE_PROFILE, MAX_CERT_SIZE, MAX_HANDLES,
+    DPE_PROFILE, MAX_CERT_SIZE,
 };
 use bitflags::bitflags;
 #[cfg(not(feature = "no-cfi"))]
@@ -120,14 +119,6 @@ impl CommandExecution for CertifyKeyCmd {
             serial: DirectoryString::PrintableString(truncated_subj_serial),
         };
 
-        // Get TCI Nodes
-        const INITIALIZER: TciNodeData = TciNodeData::new();
-        let mut nodes = [INITIALIZER; MAX_HANDLES];
-        let tcb_count = dpe.get_tcb_nodes(idx, &mut nodes)?;
-        if tcb_count > MAX_HANDLES {
-            return Err(DpeErrorCode::InternalError);
-        }
-
         let mut subject_key_identifier = [0u8; MAX_KEY_IDENTIFIER_SIZE];
         // compute key identifier as SHA hash of the DER encoded subject public key
         let mut hasher = env.crypto.hash_initialize(DPE_PROFILE.alg_len())?;
@@ -153,7 +144,7 @@ impl CommandExecution for CertifyKeyCmd {
 
         let measurements = MeasurementData {
             label: &self.label,
-            tci_nodes: &nodes[..tcb_count],
+            tcb: dpe.get_tcb(idx),
             is_ca: self.uses_is_ca(),
             supports_recursive: dpe.support.recursive(),
             subject_key_identifier,

dpe/src/context.rs

@@ -214,6 +214,10 @@ impl ChildToRootIter<'_> {
             count: 0,
         }
     }
+
+    fn curr_idx(&self) -> usize {
+        self.idx
+    }
 }
 
 impl<'a> Iterator for ChildToRootIter<'a> {
@@ -251,6 +255,98 @@ impl<'a> Iterator for ChildToRootIter<'a> {
     }
 }
 
+pub(crate) struct RootToChildIter<'a> {
+    contexts: &'a [Context],
+    path: [u8; MAX_HANDLES],
+    done: bool,
+    curr: usize,
+}
+
+impl<'a> RootToChildIter<'a> {
+    /// Create a new iterator that will start at the leaf and go to the root node.
+    pub fn new(leaf_idx: usize, contexts: &'a [Context]) -> Result<Self, DpeErrorCode> {
+        let mut iter = ChildToRootIter::new(leaf_idx, contexts);
+        let mut path = [0u8; MAX_HANDLES];
+        let mut path_size = 0;
+
+        // Use while loop to avoid consuming iter so curr_idx can be read
+        // inside the loop
+        let mut curr = iter.curr_idx();
+        while let Some(_) = iter.next() {
+            if path_size > contexts.len() {
+                return Err(DpeErrorCode::InternalError);
+            }
+
+            // Guaranteed to fit into u8 since Contexts struct is of size
+            // MAX_HANDLES.
+            path[path_size] = curr as u8;
+            path_size += 1;
+            curr = iter.curr_idx();
+        }
+
+        let (done, curr_idx) = if path_size > 0 {
+            (false, path_size - 1)
+        } else {
+            (true, 0)
+        };
+
+        Ok(Self {
+            contexts,
+            path,
+            done,
+            curr: curr_idx,
+        })
+    }
+}
+
+impl<'a> Iterator for RootToChildIter<'a> {
+    type Item = Result<&'a Context, DpeErrorCode>;
+
+    fn next(&mut self) -> Option<Result<&'a Context, DpeErrorCode>> {
+        if self.done {
+            return None;
+        }
+
+        // PANIC FREE: curr is initialized to be less that path.len() and is
+        // only ever decremented.
+        let context = &self.contexts[self.path[self.curr] as usize];
+
+        if self.curr == 0 {
+            self.done = true
+        } else {
+            self.curr -= 1;
+        }
+        Some(Ok(context))
+    }
+}
+
+pub(crate) struct ContextTcb<'a> {
+    idx: usize,
+    contexts: &'a [Context],
+}
+
+impl<'a> ContextTcb<'a> {
+    /// Create a new iterator that will start at the leaf and go to the root node.
+    pub(crate) fn new(leaf_idx: usize, contexts: &'a [Context]) -> Self {
+        ContextTcb {
+            idx: leaf_idx,
+            contexts,
+        }
+    }
+
+    pub fn rev_iter(&self) -> Result<RootToChildIter<'a>, DpeErrorCode> {
+        RootToChildIter::new(self.idx, self.contexts)
+    }
+
+    pub fn iter(&self) -> ChildToRootIter<'a> {
+        ChildToRootIter::new(self.idx, self.contexts)
+    }
+
+    pub fn count(&self) -> usize {
+        self.iter().count()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

dpe/src/dpe_instance.rs

@@ -8,10 +8,10 @@ Abstract:
 use crate::INTERNAL_INPUT_INFO_SIZE;
 use crate::{
     commands::{Command, CommandExecution, InitCtxCmd},
-    context::{ChildToRootIter, Context, ContextHandle, ContextState},
+    context::{ChildToRootIter, Context, ContextHandle, ContextState, ContextTcb},
     response::{DpeErrorCode, GetProfileResp, Response, ResponseHdr},
     support::Support,
-    tci::{TciMeasurement, TciNodeData},
+    tci::TciMeasurement,
     U8Bool, DPE_PROFILE, MAX_HANDLES,
 };
 #[cfg(not(feature = "no-cfi"))]
@@ -312,40 +312,10 @@ impl DpeInstance {
         Ok(())
     }
 
-    /// Get the TCI nodes from the context at `start_idx` to the root node following parent
-    /// links. These are the nodes that should contribute to CDI and key
-    /// derivation for the context at `start_idx`.
-    ///
-    /// # Arguments
-    ///
-    /// * `start_idx` - Index into context array
-    /// * `nodes` - Array to write TCI nodes to
-    ///
-    /// Returns the number of TCIs written to `nodes`
-    #[cfg_attr(not(feature = "no-cfi"), cfi_impl_fn)]
-    pub(crate) fn get_tcb_nodes(
-        &self,
-        start_idx: usize,
-        nodes: &mut [TciNodeData],
-    ) -> Result<usize, DpeErrorCode> {
-        let mut out_idx = 0;
-
-        for status in ChildToRootIter::new(start_idx, &self.contexts) {
-            let curr = status?;
-            if out_idx >= nodes.len() {
-                return Err(DpeErrorCode::InternalError);
-            }
-
-            nodes[out_idx] = curr.tci;
-            out_idx += 1;
-        }
-
-        if out_idx > nodes.len() {
-            return Err(DpeErrorCode::InternalError);
-        }
-        nodes[..out_idx].reverse();
-
-        Ok(out_idx)
+    // Get an iterator that collects all the contexts in the path from
+    // `start_idx` to the root
+    pub(crate) fn get_tcb(&self, start_idx: usize) -> ContextTcb {
+        ContextTcb::new(start_idx, &self.contexts)
     }
 
     /// Adds `measurement` to `context`. The current TCI is the measurement and