Reduce stack usage for DPE command input structures

Use zerocopy to pass references to DPE command structures, using the input slice as backing storage. This saves a few hundred bytes over copying these values to structs on the stack.
chipsalliance · Oct 18, 2024 · f025f93 · f025f93
1 parent a0a334d
commit f025f93
Show file tree

Hide file tree

Showing 8 changed files with 24 additions and 22 deletions.
diff --git a/dpe/src/commands/certify_key.rs b/dpe/src/commands/certify_key.rs
@@ -355,7 +355,7 @@ mod tests {
             .to_vec();
         command.extend(TEST_CERTIFY_KEY_CMD.as_bytes());
         assert_eq!(
-            Ok(Command::CertifyKey(TEST_CERTIFY_KEY_CMD)),
+            Ok(Command::CertifyKey(&TEST_CERTIFY_KEY_CMD)),
             Command::deserialize(&command)
         );
     }

diff --git a/dpe/src/commands/derive_context.rs b/dpe/src/commands/derive_context.rs
@@ -390,7 +390,7 @@ mod tests {
             .to_vec();
         command.extend(TEST_DERIVE_CONTEXT_CMD.as_bytes());
         assert_eq!(
-            Ok(Command::DeriveContext(TEST_DERIVE_CONTEXT_CMD)),
+            Ok(Command::DeriveContext(&TEST_DERIVE_CONTEXT_CMD)),
             Command::deserialize(&command)
         );
     }

diff --git a/dpe/src/commands/destroy_context.rs b/dpe/src/commands/destroy_context.rs
@@ -119,7 +119,7 @@ mod tests {
             .to_vec();
         command.extend(TEST_DESTROY_CTX_CMD.as_bytes());
         assert_eq!(
-            Ok(Command::DestroyCtx(TEST_DESTROY_CTX_CMD)),
+            Ok(Command::DestroyCtx(&TEST_DESTROY_CTX_CMD)),
             Command::deserialize(&command)
         );
     }

diff --git a/dpe/src/commands/get_certificate_chain.rs b/dpe/src/commands/get_certificate_chain.rs
@@ -74,7 +74,9 @@ mod tests {
             .to_vec();
         command.extend(TEST_GET_CERTIFICATE_CHAIN_CMD.as_bytes());
         assert_eq!(
-            Ok(Command::GetCertificateChain(TEST_GET_CERTIFICATE_CHAIN_CMD)),
+            Ok(Command::GetCertificateChain(
+                &TEST_GET_CERTIFICATE_CHAIN_CMD
+            )),
             Command::deserialize(&command)
         );
     }

diff --git a/dpe/src/commands/initialize_context.rs b/dpe/src/commands/initialize_context.rs
@@ -132,7 +132,7 @@ mod tests {
             .to_vec();
         command.extend(TEST_INIT_CTX_CMD.as_bytes());
         assert_eq!(
-            Ok(Command::InitCtx(TEST_INIT_CTX_CMD)),
+            Ok(Command::InitCtx(&TEST_INIT_CTX_CMD)),
             Command::deserialize(&command)
         );
     }

diff --git a/dpe/src/commands/mod.rs b/dpe/src/commands/mod.rs
@@ -20,7 +20,7 @@ use crate::{
     DPE_PROFILE,
 };
 use core::mem::size_of;
-use zerocopy::FromBytes;
+use zerocopy::{FromBytes, Immutable, KnownLayout};
 
 mod certify_key;
 mod derive_context;
@@ -32,19 +32,19 @@ mod rotate_context;
 mod sign;
 
 #[derive(Debug, PartialEq, Eq)]
-pub enum Command {
+pub enum Command<'a> {
     GetProfile,
-    InitCtx(InitCtxCmd),
-    DeriveContext(DeriveContextCmd),
-    CertifyKey(CertifyKeyCmd),
-    Sign(SignCmd),
+    InitCtx(&'a InitCtxCmd),
+    DeriveContext(&'a DeriveContextCmd),
+    CertifyKey(&'a CertifyKeyCmd),
+    Sign(&'a SignCmd),
     #[cfg(not(feature = "disable_rotate_context"))]
-    RotateCtx(RotateCtxCmd),
-    DestroyCtx(DestroyCtxCmd),
-    GetCertificateChain(GetCertificateChainCmd),
+    RotateCtx(&'a RotateCtxCmd),
+    DestroyCtx(&'a DestroyCtxCmd),
+    GetCertificateChain(&'a GetCertificateChainCmd),
 }
 
-impl Command {
+impl Command<'_> {
     pub const GET_PROFILE: u32 = 0x01;
     pub const INITIALIZE_CONTEXT: u32 = 0x07;
     pub const DERIVE_CONTEXT: u32 = 0x08;
@@ -80,17 +80,17 @@ impl Command {
         }
     }
 
-    fn parse_command<T: FromBytes>(
-        build: impl FnOnce(T) -> Command,
-        bytes: &[u8],
+    fn parse_command<'a, T: FromBytes + KnownLayout + Immutable + 'a>(
+        build: impl FnOnce(&'a T) -> Command,
+        bytes: &'a [u8],
     ) -> Result<Command, DpeErrorCode> {
         let (prefix, _remaining_bytes) =
-            T::read_from_prefix(bytes).map_err(|_| DpeErrorCode::InvalidArgument)?;
+            T::ref_from_prefix(bytes).map_err(|_| DpeErrorCode::InvalidArgument)?;
         Ok(build(prefix))
     }
 }
 
-impl From<Command> for u32 {
+impl From<Command<'_>> for u32 {
     fn from(cmd: Command) -> u32 {
         match cmd {
             Command::GetProfile => Command::GET_PROFILE,

diff --git a/dpe/src/commands/rotate_context.rs b/dpe/src/commands/rotate_context.rs
@@ -151,7 +151,7 @@ mod tests {
             .to_vec();
         command.extend(TEST_ROTATE_CTX_CMD.as_bytes());
         assert_eq!(
-            Ok(Command::RotateCtx(TEST_ROTATE_CTX_CMD)),
+            Ok(Command::RotateCtx(&TEST_ROTATE_CTX_CMD)),
             Command::deserialize(&command)
         );
     }

diff --git a/dpe/src/commands/sign.rs b/dpe/src/commands/sign.rs
@@ -230,7 +230,7 @@ mod tests {
         let mut command = CommandHdr::new_for_test(Command::SIGN).as_bytes().to_vec();
         command.extend(TEST_SIGN_CMD.as_bytes());
         assert_eq!(
-            Ok(Command::Sign(TEST_SIGN_CMD)),
+            Ok(Command::Sign(&TEST_SIGN_CMD)),
             Command::deserialize(&command)
         );
     }