[hyperlight_common] Add resource table abstraction

Signed-off-by: Lucy Menon <[email protected]>

Author: syntactically

src/hyperlight_common/Cargo.toml

@@ -21,6 +21,7 @@ log = "0.4.26"
 tracing = { version = "0.1.41", optional = true }
 strum = {version = "0.27",  default-features = false, features = ["derive"]}
 arbitrary = {version = "1.4.1", optional = true, features = ["derive"]}
+spin = "0.9.8"
 
 [features]
 default = ["tracing"]

src/hyperlight_common/src/lib.rs

@@ -36,3 +36,5 @@ pub mod flatbuffer_wrappers;
 mod flatbuffers;
 /// cbindgen:ignore
 pub mod mem;
+
+pub mod resource;

src/hyperlight_common/src/resource.rs

@@ -0,0 +1,141 @@
+///! Shared operations around resources
+use alloc::sync::Arc;
+
+#[cfg(target_os = "linux")]
+extern crate std;
+use core::marker::{PhantomData, Send};
+use core::ops::Deref;
+#[cfg(target_os = "linux")]
+use std::sync::{RwLock, RwLockReadGuard};
+
+#[cfg(not(target_os = "linux"))]
+use spin::{RwLock, RwLockReadGuard};
+
+/// The semantics of component model resources are, pleasingly,
+/// roughly compatible with those of Rust. Less pleasingly, it's not
+/// terribly easy to show that statically.
+///
+/// In particular, if the host calls into the guest and gives it a
+/// borrow of a resource, reentrant host function calls that use that
+/// borrow need to be able to resolve the original reference and use
+/// it in an appropriately scoped manner, but it is not simple to do
+/// this, because the core Hyperlight machinery doesn't offer an easy
+/// way to augment the host's context for the span of time of a guest
+/// function call.  This may be worth revisiting at some time, but in
+/// the meantime, it's easier to just do it dynamically.
+///
+/// # Safety
+/// Informally: this only creates SharedRead references, so having a
+/// bunch of them going at once is fine.  Safe Rust in the host can't
+/// use any earlier borrows (potentially invalidating these) until
+/// borrow passed into [`ResourceEntry::lend`] has expired.  Because
+/// that borrow outlives the [`LentResourceGuard`], it will not expire
+/// until that destructor is called. That destructor ensures that (a)
+/// there are no outstanding [`BorrowedResourceGuard`]s alive (since
+/// they would be holding the read side of the [`RwLock`] if they
+/// were), and that (b) the shared flag has been set to false, so
+/// [`ResourceEntry::borrow`] will never create another borrow
+pub enum ResourceEntry<T> {
+    Empty,
+    Owned(T),
+    Borrowed(Arc<RwLock<bool>>, *const T),
+}
+unsafe impl<T: Send> Send for ResourceEntry<T> {}
+
+pub struct LentResourceGuard<'a> {
+    flag: Arc<RwLock<bool>>,
+    already_revoked: bool,
+    _phantom: core::marker::PhantomData<&'a mut ()>,
+}
+impl<'a> LentResourceGuard<'a> {
+    pub fn revoke_nonblocking(&mut self) -> bool {
+        #[cfg(target_os = "linux")]
+        let Ok(mut flag) = self.flag.try_write() else {
+            return false;
+        };
+        #[cfg(not(target_os = "linux"))]
+        let Some(mut flag) = self.flag.try_write() else {
+            return false;
+        };
+        *flag = false;
+        self.already_revoked = true;
+        true
+    }
+}
+impl<'a> Drop for LentResourceGuard<'a> {
+    fn drop(&mut self) {
+        if !self.already_revoked {
+            let mut guard = self.flag.write();
+            #[cfg(target_os = "linux")]
+            {
+                *guard.unwrap() = false;
+            }
+            #[cfg(not(target_os = "linux"))]
+            {
+                *guard = false;
+            }
+        }
+    }
+}
+pub struct BorrowedResourceGuard<'a, T> {
+    _flag: Option<RwLockReadGuard<'a, bool>>,
+    reference: &'a T,
+}
+impl<'a, T> Deref for BorrowedResourceGuard<'a, T> {
+    type Target = T;
+    fn deref(&self) -> &T {
+        self.reference
+    }
+}
+impl<T> ResourceEntry<T> {
+    pub fn give(x: T) -> ResourceEntry<T> {
+        ResourceEntry::Owned(x)
+    }
+    pub fn lend<'a>(x: &'a T) -> (LentResourceGuard<'a>, ResourceEntry<T>) {
+        let flag = Arc::new(RwLock::new(true));
+        (
+            LentResourceGuard {
+                flag: flag.clone(),
+                already_revoked: false,
+                _phantom: PhantomData {},
+            },
+            ResourceEntry::Borrowed(flag, x as *const T),
+        )
+    }
+    pub fn borrow<'a>(&'a self) -> Option<BorrowedResourceGuard<'a, T>> {
+        match self {
+            ResourceEntry::Empty => None,
+            ResourceEntry::Owned(t) => Some(BorrowedResourceGuard {
+                _flag: None,
+                reference: &t,
+            }),
+            ResourceEntry::Borrowed(flag, t) => {
+                let guard = flag.read();
+                let flag = {
+                    #[cfg(target_os = "linux")]
+                    {
+                        guard.unwrap()
+                    }
+                    #[cfg(not(target_os = "linux"))]
+                    {
+                        guard
+                    }
+                };
+                if *flag {
+                    Some(BorrowedResourceGuard {
+                        _flag: Some(flag),
+                        reference: unsafe { &**t },
+                    })
+                } else {
+                    None
+                }
+            }
+        }
+    }
+    pub fn take(&mut self) -> Option<T> {
+        match core::mem::replace(self, ResourceEntry::Empty) {
+            ResourceEntry::Owned(t) => Some(t),
+            _ => None,
+        }
+    }
+}