From 901771674d4ff4d43fa07fbed65d1284474445c6 Mon Sep 17 00:00:00 2001
From: Raoul Strackx <raoul.strackx@fortanix.com>
Date: Thu, 4 Jan 2024 09:57:03 +0100
Subject: [PATCH] Addressed reviewer comments

---
 intel-sgx/async-usercalls/src/lib.rs    | 15 +++++----------
 intel-sgx/async-usercalls/src/queues.rs |  5 ++++-
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/intel-sgx/async-usercalls/src/lib.rs b/intel-sgx/async-usercalls/src/lib.rs
index 7485a31d..71acb4d8 100644
--- a/intel-sgx/async-usercalls/src/lib.rs
+++ b/intel-sgx/async-usercalls/src/lib.rs
@@ -205,16 +205,11 @@ impl CallbackHandler {
             0 => return 0,
             n => &returns[..n],
         };
-        // 2. try to lock the mutex, if successful, receive all pending callbacks and put them in the hash map
-        let mut guard = match self.callbacks.try_lock() {
-            Ok(mut callbacks) => {
-                for (id, cb) in self.callback_rx.try_iter() {
-                    callbacks.insert(id, cb);
-                }
-                callbacks
-            }
-            _ => self.callbacks.lock().unwrap(),
-        };
+        // 2. Receive all pending callbacks and put them in the hash map
+        let mut guard = self.callbacks.lock().unwrap();
+        for (id, cb) in self.callback_rx.try_iter() {
+            guard.insert(id, cb);
+        }
         // 3. remove callbacks for returns received in step 1 from the hash map
         let mut ret_callbacks = Vec::with_capacity(returns.len());
         for ret in returns {
diff --git a/intel-sgx/async-usercalls/src/queues.rs b/intel-sgx/async-usercalls/src/queues.rs
index a0ded0c4..08ba74f0 100644
--- a/intel-sgx/async-usercalls/src/queues.rs
+++ b/intel-sgx/async-usercalls/src/queues.rs
@@ -186,6 +186,9 @@ mod map {
                 let id = self.next_id;
                 // We intentionally ignore the overflow here, thus allowing `next_id` to jump back to 0
                 // after `u32::MAX` number of insertions.
+                // TODO: We should have a way of limiting the size of this queue to avoid
+                //   potentially checking 2^32 items and huge memory consumption
+                //   https://github.com/fortanix/rust-sgx/issues/550
                 self.next_id = self.next_id.overflowing_add(1).0;
                 if !self.map.contains_key(&id) {
                     self.map.insert(id, value);
@@ -206,4 +209,4 @@ mod map {
             self.map.remove(&id)
         }
     }
-}
\ No newline at end of file
+}