From 52bc29f5573d50704d59220ab111904b399b340b 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, 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
+        );
     }
 }