From 3d6d05dbde49cebf6e5b330588c323aed3c7143f Mon Sep 17 00:00:00 2001
From: Anthony Rocha <Anthony.Rocha@amd.com>
Date: Fri, 7 Jun 2024 10:48:34 -0700
Subject: [PATCH] Adjust the maximum number of avaiable DPE handles.

This commit adjusts the maximum number of DPE handles via cargo feature. Exception and NMI stacks
for the mutable code (fmc, runtime) had to be adjusted to make room for the extra DPE storage.
---
 builder/src/firmware.rs      |  4 ++--
 drivers/src/memory_layout.rs |  2 +-
 drivers/src/persistent.rs    | 15 ++++++++++-----
 runtime/build.rs             |  1 +
 runtime/build.sh             |  3 +--
 runtime/src/lib.rs           |  3 +++
 6 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/builder/src/firmware.rs b/builder/src/firmware.rs
index 8e759ceb94..c9dab807ba 100644
--- a/builder/src/firmware.rs
+++ b/builder/src/firmware.rs
@@ -48,13 +48,13 @@ pub const FMC_FAKE_WITH_UART: FwId = FwId {
 pub const APP: FwId = FwId {
     crate_name: "caliptra-runtime",
     bin_name: "caliptra-runtime",
-    features: &["fips_self_test"],
+    features: &["fips_self_test", "arbitrary_max_handles"],
 };
 
 pub const APP_WITH_UART: FwId = FwId {
     crate_name: "caliptra-runtime",
     bin_name: "caliptra-runtime",
-    features: &["emu", "fips_self_test"],
+    features: &["emu", "fips_self_test", "arbitrary_max_handles"],
 };
 
 pub const APP_WITH_UART_FPGA: FwId = FwId {
diff --git a/drivers/src/memory_layout.rs b/drivers/src/memory_layout.rs
index 1de4524454..e9760e07d6 100644
--- a/drivers/src/memory_layout.rs
+++ b/drivers/src/memory_layout.rs
@@ -68,7 +68,7 @@ pub const RTALIAS_TBS_SIZE: u32 = 1024;
 pub const PCR_LOG_SIZE: u32 = 1024;
 pub const MEASUREMENT_LOG_SIZE: u32 = 1024;
 pub const FUSE_LOG_SIZE: u32 = 1024;
-pub const DPE_SIZE: u32 = 4 * 1024;
+pub const DPE_SIZE: u32 = 5 * 1024;
 pub const PCR_RESET_COUNTER_SIZE: u32 = 1024;
 pub const DATA_SIZE: u32 = 78 * 1024;
 pub const STACK_SIZE: u32 = 22 * 1024;
diff --git a/drivers/src/persistent.rs b/drivers/src/persistent.rs
index e58ed69706..4e0bfddcdc 100644
--- a/drivers/src/persistent.rs
+++ b/drivers/src/persistent.rs
@@ -22,6 +22,15 @@ pub const PCR_LOG_MAX_COUNT: usize = 17;
 pub const FUSE_LOG_MAX_COUNT: usize = 62;
 pub const MEASUREMENT_MAX_COUNT: usize = 8;
 
+#[cfg(feature = "runtime")]
+const DPE_DCCM_STORAGE: usize = size_of::<DpeInstance>()
+    + size_of::<u32>() * MAX_HANDLES
+    + size_of::<U8Bool>() * MAX_HANDLES
+    + size_of::<U8Bool>();
+
+#[cfg(feature = "runtime")]
+const _: () = assert!(DPE_DCCM_STORAGE < memory_layout::DPE_SIZE as usize);
+
 pub type PcrLogArray = [PcrLogEntry; PCR_LOG_MAX_COUNT];
 pub type FuseLogArray = [FuseLogEntry; FUSE_LOG_MAX_COUNT];
 pub type StashMeasurementArray = [MeasurementLogEntry; MEASUREMENT_MAX_COUNT];
@@ -63,11 +72,7 @@ pub struct PersistentData {
     #[cfg(feature = "runtime")]
     pub attestation_disabled: U8Bool,
     #[cfg(feature = "runtime")]
-    reserved6: [u8; memory_layout::DPE_SIZE as usize
-        - size_of::<DpeInstance>()
-        - size_of::<u32>() * MAX_HANDLES
-        - size_of::<U8Bool>() * MAX_HANDLES
-        - size_of::<U8Bool>()],
+    reserved6: [u8; memory_layout::DPE_SIZE as usize - DPE_DCCM_STORAGE],
     #[cfg(not(feature = "runtime"))]
     dpe: [u8; memory_layout::DPE_SIZE as usize],
     #[cfg(feature = "runtime")]
diff --git a/runtime/build.rs b/runtime/build.rs
index dc1c5989f9..284605d429 100644
--- a/runtime/build.rs
+++ b/runtime/build.rs
@@ -34,6 +34,7 @@ fn main() {
 
             println!("cargo:rustc-link-arg=-Tlink.x");
             println!("cargo:rerun-if-changed=build.rs");
+            println!("cargo:rustc-env=ARBITRARY_MAX_HANDLES=32");
         }
     }
 }
diff --git a/runtime/build.sh b/runtime/build.sh
index 0699c1e463..4d38168cfd 100755
--- a/runtime/build.sh
+++ b/runtime/build.sh
@@ -4,8 +4,7 @@
 
 cd "$(dirname "${BASH_SOURCE[0]}")"
 
-ARBITRARY_MAX_HANDLES=24 cargo build \
- --features arbitrary_max_handles \
+cargo build \
   --locked \
   --target riscv32imc-unknown-none-elf \
   --profile=firmware \
diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs
index eb4d8e77b4..b2d6b879e0 100644
--- a/runtime/src/lib.rs
+++ b/runtime/src/lib.rs
@@ -105,7 +105,10 @@ pub const DPE_SUPPORT: Support = Support::all();
 pub const MAX_CERT_CHAIN_SIZE: usize = 4096;
 
 pub const PL0_PAUSER_FLAG: u32 = 1;
+#[cfg(not(feature = "arbitrary_max_handles"))]
 pub const PL0_DPE_ACTIVE_CONTEXT_THRESHOLD: usize = 8;
+#[cfg(feature = "arbitrary_max_handles")]
+pub const PL0_DPE_ACTIVE_CONTEXT_THRESHOLD: usize = 16;
 pub const PL1_DPE_ACTIVE_CONTEXT_THRESHOLD: usize = 16;
 
 pub struct CptraDpeTypes;