Make protocol db race free

2 files changed, 28 insertions, 18 deletions

diff --git a/src/efi/bootservices.rs b/src/efi/bootservices.rs
index 2b072ba..50870ff 100644
--- a/src/efi/bootservices.rs
+++ b/src/efi/bootservices.rs
@@ -10,6 +10,7 @@ use crate::efi::RTSDATA_ALLOCATOR;
 use crate::efi::install_configtable;
 use crate::efi::memmap;
 use crate::efi::memmap::Placement;
+use crate::efi::ProtocolPtr;
 
 use crate::efi::devicepath;
 use crate::efi::bootservices::AllocateType::*;
@@ -435,8 +436,8 @@ extern "C" fn handle_protocol(
 ) -> Status {
     let protocol = unsafe { &*_protocol };
     let key = (_handle, *protocol);
-    if let Some(ptr) = unsafe { PROTOCOL_DB.get(&key) } {
-        unsafe { *_interface = *ptr };
+    if let Some(ptr) = PROTOCOL_DB.lock().get(&key) {
+        unsafe { *_interface = ptr.0 };
         Status::EFI_SUCCESS
     } else {
         Status::EFI_UNSUPPORTED
@@ -461,7 +462,7 @@ extern "C" fn locate_handle(
     Status::EFI_NOT_FOUND
 }
 
-fn compare_device_path(entry: (&(usize, Guid), &*const ()),
+fn compare_device_path(entry: (&(usize, Guid), &ProtocolPtr),
                        protocol: &Guid,
                        device_path: &DevicePath,
                        db: &ProtocolDb,
@@ -477,7 +478,7 @@ fn compare_device_path(entry: (&(usize, Guid), &*const ()),
     // Check whether the device path in the protocol database
     // is a prefix of the provided device path
     let bytes_equal = devicepath::is_prefix(
-        unsafe { &*(*entry.1 as *const DevicePath) },
+        unsafe { &*(entry.1.0 as *const DevicePath) },
         device_path,
         );
     if bytes_equal == 0 {
@@ -495,7 +496,7 @@ extern "C" fn locate_device_path(
     device_path: *mut *const DevicePath,
     device: *mut Handle
 ) -> Status {
-    let db = unsafe { &PROTOCOL_DB };
+    let db = PROTOCOL_DB.lock();
     let protocol = unsafe { &*protocol };
     let devpath = unsafe { &**device_path };
 
@@ -505,8 +506,8 @@ extern "C" fn locate_device_path(
     // provided one, if any
     if let Some(entry) = db.iter()
         .filter_map(
-            |entry: (&(usize, Guid), &*const ())|
-                compare_device_path(entry, protocol, devpath, db))
+            |entry: (&(usize, Guid), &ProtocolPtr)|
+                compare_device_path(entry, protocol, devpath, &db))
         .max_by(|a, b| a.0.cmp(&b.0)) {
 
         unsafe {
@@ -551,11 +552,11 @@ extern "C" fn exit(
     _exit_data_size: usize,
     _exit_data: *const Char16
 ) -> Status {
-    let db = unsafe { &PROTOCOL_DB };
+    let db = PROTOCOL_DB.lock();
     let key = (image_handle, EFI_LOADED_IMAGE_PROTOCOL_GUID);
     if let Some(ptr) = db.get(&key) {
         unsafe {
-            let li = &mut *(*ptr as *mut LoadedImage);
+            let li = &mut *(ptr.0 as *mut LoadedImage);
             if li.reserved != 0 {
                 let sp = li.reserved;
                 li.reserved = 0;
@@ -672,11 +673,11 @@ extern "C" fn locate_protocol(
     _registration: *const (),
     _interface: *mut *const ()
 ) -> Status {
-    let db = unsafe { &PROTOCOL_DB };
+    let db = PROTOCOL_DB.lock();
     let protocol = unsafe { *_protocol };
     if let Some(entry) = db.iter()
-        .find(|e: &(&(usize, Guid), &*const ())| e.0.1 == protocol) {
-        unsafe { *_interface = *entry.1; }
+        .find(|e: &(&(usize, Guid), &ProtocolPtr)| e.0.1 == protocol) {
+        unsafe { *_interface = entry.1.0; }
         Status::EFI_SUCCESS
     } else {
         Status::EFI_NOT_FOUND
diff --git a/src/efi/mod.rs b/src/efi/mod.rs
index 0058c86..2c54663 100644
--- a/src/efi/mod.rs
+++ b/src/efi/mod.rs
@@ -19,6 +19,7 @@ use alloc::collections::BTreeMap;
 use alloc::vec::Vec;
 
 use linked_list_allocator::LockedHeap;
+use spinning_top::Spinlock;
 
 const UEFI_REVISION: u32 = (2 << 16) | 100; // 2.10
 
@@ -154,9 +155,19 @@ impl MemoryAttributesTable {
 
 pub(crate) use guid;
 
-type ProtocolDb = BTreeMap::<(Handle, Guid), *const ()>;
+struct ProtocolPtr(*const ());
+unsafe impl Send for ProtocolPtr {}
+unsafe impl Sync for ProtocolPtr {}
 
-static mut PROTOCOL_DB: ProtocolDb = BTreeMap::new();
+impl ProtocolPtr {
+    pub fn from<T>(p: &T) -> ProtocolPtr {
+        ProtocolPtr(p as *const _ as *const())
+    }
+}
+
+type ProtocolDb = BTreeMap::<(Handle, Guid), ProtocolPtr>;
+
+static PROTOCOL_DB: Spinlock<ProtocolDb> = Spinlock::new(BTreeMap::new());
 static mut CONFIGTABLE_DB: BTreeMap::<Guid, ConfigurationTable> = BTreeMap::new();
 
 static RTSDATA_ALLOCATOR: LockedHeap = LockedHeap::empty();
@@ -181,8 +192,7 @@ pub fn install_protocol<T>(
     guid: &'static Guid,
     protocol: &T
 ) {
-    let db = unsafe { &mut PROTOCOL_DB };
-    db.insert((handle, *guid), protocol as *const _ as *const());
+    PROTOCOL_DB.lock().insert((handle, *guid), ProtocolPtr::from(protocol));
 }
 
 pub fn uninstall_protocol<T>(
@@ -190,8 +200,7 @@ pub fn uninstall_protocol<T>(
     guid: &'static Guid,
     _protocol: &T
 ) {
-    let db = unsafe { &mut PROTOCOL_DB };
-    db.remove(&(handle, *guid));
+    PROTOCOL_DB.lock().remove(&(handle, *guid));
 }
 
 pub fn install_configtable(