rust oak_runtime: Fix any possible lock inversion issues (#784)

Fixes #780

oak/server/rust/oak_runtime/src/runtime/mod.rs

@@ -502,7 +502,7 @@ impl Runtime {
         handles_capacity: usize,
     ) -> Result<Option<ReadStatus>, OakStatus> {
         self.validate_handle_access(node_id, reference)?;
-        self.channels
+        let result = self.channels
             .with_channel(self.channels.get_reader_channel(reference)?, |channel| {
                 let mut messages = channel.messages.write().unwrap();
                 match messages.front() {
@@ -517,7 +517,6 @@ impl Runtime {
                                 ReadStatus::NeedsCapacity(req_bytes_capacity, req_handles_capacity)
                             } else {
                                 let msg = messages.pop_front().expect( "Front element disappeared while we were holding the write lock!");
-                                self.track_handles_in_node(node_id, msg.channels.clone());
                                 ReadStatus::Success(msg)
                             },
                         ))
@@ -530,7 +529,15 @@ impl Runtime {
                         }
                     }
                 }
-            })
+            });
+
+        // Add handles outside the channels lock so we don't hold the node lock inside the channel
+        // lock.
+        if let Ok(Some(ReadStatus::Success(ref msg))) = result {
+            self.track_handles_in_node(node_id, msg.channels.clone());
+        }
+
+        result
     }
 
     /// Return the direction of a [`Handle`]. This is useful when reading
@@ -605,6 +612,13 @@ impl Runtime {
             .remove(&node_id)
             .expect("remove_node_id: Node didn't exist!");
     }
+
+    /// Add an [`NodeId`] [`Node`] pair to the [`Runtime`]. This method temporarily holds the node
+    /// write lock.
+    fn add_running_node(&self, reference: NodeId, node: Node) {
+        let mut nodes = self.nodes.write().unwrap();
+        nodes.insert(reference, node);
+    }
 }
 
 /// A reference to a [`Runtime`].
@@ -639,7 +653,6 @@ impl RuntimeRef {
         // to do that we first need to provide a reference to the caller node as a parameter to this
         // function.
 
-        let mut nodes = self.nodes.write().unwrap();
         let reference = self.new_node_reference();
 
         let reader = self.channels.duplicate_reference(reader)?;
@@ -666,7 +679,7 @@ impl RuntimeRef {
 
         // If the node was successfully started, insert it in the list of currently running
         // nodes.
-        nodes.insert(
+        self.add_running_node(
             reference,
             Node {
                 reference,