From ef977e9265ba6a36d55984cd9fa03df0eb330792 Mon Sep 17 00:00:00 2001
From: Marius Knaust <mknaust@amazon.com>
Date: Wed, 10 Jul 2024 15:17:57 +0200
Subject: [PATCH] nitro-cli: Update enclave boot timeout based on allocated
 memory

Before pull request #310 the enlave boot timeout was calculated based on
the allocated memory size. We reintroduce this heuristic in addition to
the heuristic based on the efi size, since we observed timeout
violations for large allocated memory sizes.

Signed-off-by: Marius Knaust <mknaust@amazon.com>
---
 src/enclave_proc/resource_manager.rs | 45 ++++++++++++++++++++++------
 1 file changed, 36 insertions(+), 9 deletions(-)

diff --git a/src/enclave_proc/resource_manager.rs b/src/enclave_proc/resource_manager.rs
index d67d6eea7..4153d0a42 100644
--- a/src/enclave_proc/resource_manager.rs
+++ b/src/enclave_proc/resource_manager.rs
@@ -589,7 +589,7 @@ impl EnclaveHandle {
             .len();
 
         // Update the poll timeout based on the eif size or allocated memory
-        let poll_timeout = calculate_necessary_timeout(eif_size);
+        let poll_timeout = calculate_necessary_timeout(eif_size, self.allocated_memory_mib * MiB);
 
         enclave_ready(listener, poll_timeout).map_err(|err| {
             let err_msg = format!("Waiting on enclave to boot failed with error {:?}", err);
@@ -1188,21 +1188,26 @@ pub fn between_packets_delay() -> Option<Duration> {
 /// # Arguments
 ///
 /// * `eif_size` - The EIF size in bytes
+/// * `allocated_memory` - The memory size in bytes
 ///
 /// # Examples
 ///
 /// ```
 /// use nitro_cli::enclave_proc::resource_manager::calculate_necessary_timeout;
 /// use nitro_cli::enclave_proc::utils::GiB;
-/// // Returns the timeout based on the 8GiB EIF size
-/// let timeout = calculate_necessary_timeout(8 * GiB);
+/// // Returns the timeout based on the 8GiB EIF size and 32GiB of allocated memory
+/// let timeout = calculate_necessary_timeout(8 * GiB, 32 * GiB);
 /// ```
-pub fn calculate_necessary_timeout(eif_size: u64) -> c_int {
+pub fn calculate_necessary_timeout(eif_size: u64, allocated_memory: u64) -> c_int {
     // in case we have a valid eif_size give TIMEOUT_MINUTE_MS ms for each 6GiB
-    let poll_timeout: c_int =
+    let efi_size_timeout: c_int =
         ((1 + (eif_size - 1) / (6 * GiB)) as i32).saturating_mul(TIMEOUT_MINUTE_MS);
 
-    poll_timeout
+    // Update the poll timeout to be TIMEOUT_MINUTE_MS per 100 GiB of enclave memory.
+    let allocated_memory_timeout: c_int =
+        ((1 + (allocated_memory - 1) / (100 * GiB)) as i32).saturating_mul(TIMEOUT_MINUTE_MS);
+
+    std::cmp::max(efi_size_timeout, allocated_memory_timeout)
 }
 
 #[cfg(test)]
@@ -1213,8 +1218,30 @@ mod tests {
 
     #[test]
     fn test_timeout_calculation() {
-        assert_eq!(calculate_necessary_timeout(2 * GiB), TIMEOUT_MINUTE_MS);
-        assert_eq!(calculate_necessary_timeout(6 * GiB), TIMEOUT_MINUTE_MS);
-        assert_eq!(calculate_necessary_timeout(10 * GiB), 2 * TIMEOUT_MINUTE_MS);
+        assert_eq!(
+            calculate_necessary_timeout(2 * GiB, 32 * GiB),
+            TIMEOUT_MINUTE_MS
+        );
+        assert_eq!(
+            calculate_necessary_timeout(6 * GiB, 32 * GiB),
+            TIMEOUT_MINUTE_MS
+        );
+        assert_eq!(
+            calculate_necessary_timeout(10 * GiB, 32 * GiB),
+            2 * TIMEOUT_MINUTE_MS
+        );
+
+        assert_eq!(
+            calculate_necessary_timeout(2 * GiB, 128 * GiB),
+            2 * TIMEOUT_MINUTE_MS
+        );
+        assert_eq!(
+            calculate_necessary_timeout(6 * GiB, 128 * GiB),
+            2 * TIMEOUT_MINUTE_MS
+        );
+        assert_eq!(
+            calculate_necessary_timeout(10 * GiB, 128 * GiB),
+            2 * TIMEOUT_MINUTE_MS
+        );
     }
 }