User iterator instead of temporary buffer for CertifyKey TCB

When encoding TcbInfos, do not make a temporary buffer to copy all the
TCI measurements. This adds multiple KiB of stack usage. Instead, add a
function to allow the X.509 library to create an iterator for the TCB of
the node.

This allows the main DPE context array to be used as the backing memory
when encoding TcbInfos instead of using extra stack space.

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,95 @@ impl<'a> Iterator for ChildToRootIter<'a> {
     }
 }
 
+pub(crate) struct RootToChildIter<'a> {
+    contexts: &'a [Context],
+    path: [u8; MAX_HANDLES],
+    path_size: usize,
+    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();
+        }
+
+        Ok(Self {
+            contexts,
+            path,
+            path_size,
+            done: false,
+            curr: path_size - 1,
+        })
+    }
+}
+
+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;
+        }
+        if self.curr >= self.contexts.len() || self.curr >= self.path_size {
+            return Some(Err(DpeErrorCode::InternalError));
+        }
+
+        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