further optimizations of locking

Signed-off-by: Wojciech Ozga <woz@zurich.ibm.com>

Author: wojciechozga

security-monitor/src/confidential_flow/finite_state_machine.rs

@@ -143,7 +143,7 @@ impl<'a> ConfidentialFlow<'a> {
         // Now, we are going to change the context between security domains.
         // 1) Store the hypervisor hart state that executed on this physical hart to the main memory.
         hardware_hart.hypervisor_hart_mut().save_in_main_memory();
-        match ControlDataStorage::try_confidential_vm_mut(confidential_vm_id, |mut confidential_vm| {
+        match ControlDataStorage::try_confidential_vm(confidential_vm_id, |confidential_vm| {
             confidential_vm.steal_confidential_hart(confidential_hart_id, hardware_hart)?;
             Ok(confidential_vm.allowed_external_interrupts())
         }) {
@@ -168,7 +168,7 @@ impl<'a> ConfidentialFlow<'a> {
         // Now, we are going to change the context between security domains.
         // 1) Store the confidential hart state that executed on this physical hart to the main memory.
         self.hardware_hart.confidential_hart_mut().save_in_main_memory();
-        let _ = ControlDataStorage::try_confidential_vm_mut(self.confidential_vm_id(), |mut confidential_vm| {
+        let _ = ControlDataStorage::try_confidential_vm(self.confidential_vm_id(), |confidential_vm| {
             Ok(confidential_vm.return_confidential_hart(self.hardware_hart))
         })
         // Below unwrap is safe because we are in the confidential flow that guarantees that the confidential VM with

security-monitor/src/core/control_data/confidential_vm.rs

@@ -97,23 +97,22 @@ impl ConfidentialVm {
     /// # Guarantees
     ///
     /// The physical hart is configured to enforce memory access control so that the confidential VM has access only to its own memory.
-    pub fn steal_confidential_hart(&mut self, confidential_hart_id: usize, hardware_hart: &mut HardwareHart) -> Result<(), Error> {
+    pub fn steal_confidential_hart(&self, confidential_hart_id: usize, hardware_hart: &mut HardwareHart) -> Result<(), Error> {
         ensure!(confidential_hart_id < self.confidential_harts.len(), Error::InvalidHartId())?;
-        let confidential_hart = self.confidential_harts[confidential_hart_id].get_mut();
+        let mut confidential_hart = self.confidential_harts[confidential_hart_id].write();
         // The hypervisor might try to schedule the same confidential hart on different physical harts. We detect it
         // because after a confidential_hart is scheduled for the first time, its token is stolen and the
         // ConfidentialVM is left with a dummy confidential_hart. A dummy confidential hart is a hart not associated
         // with any confidential vm.
         ensure_not!(confidential_hart.is_dummy(), Error::HartAlreadyRunning())?;
         // The hypervisor might try to schedule a confidential hart that has never been started. This is forbidden.
         ensure!(confidential_hart.is_executable(), Error::HartNotExecutable())?;
-        unsafe {
-            core::ptr::swap(hardware_hart.confidential_hart_mut() as *mut ConfidentialHart, confidential_hart as *mut ConfidentialHart);
-            // Reconfigure the hardware memory isolation mechanism to enforce that the confidential virtual machine has access only to the
-            // memory regions it owns. Below invocation is safe because we are now in the confidential flow part of the finite state
-            // machine and the virtual hart is assigned to the hardware hart.
-            self.memory_protector().enable()
-        };
+
+        core::mem::swap(hardware_hart.confidential_hart_mut(), &mut confidential_hart);
+        // Reconfigure the hardware memory isolation mechanism to enforce that the confidential virtual machine has access only to the
+        // memory regions it owns. Below invocation is safe because we are now in the confidential flow part of the finite state
+        // machine and the virtual hart is assigned to the hardware hart.
+        unsafe { self.memory_protector().enable() };
         Ok(())
     }
 
@@ -123,15 +122,12 @@ impl ConfidentialVm {
     /// # Safety
     ///
     /// A confidential hart belongs to this confidential VM and is currently assigned to the hardware hart.
-    pub fn return_confidential_hart(&mut self, hardware_hart: &mut HardwareHart) {
+    pub fn return_confidential_hart(&self, hardware_hart: &mut HardwareHart) {
         assert!(!hardware_hart.confidential_hart().is_dummy());
         assert!(Some(self.id) == hardware_hart.confidential_hart().confidential_vm_id());
         let confidential_hart_id = hardware_hart.confidential_hart().confidential_hart_id();
         assert!(self.confidential_harts.len() > confidential_hart_id);
-        let confidential_hart = self.confidential_harts[confidential_hart_id].get_mut();
-        unsafe {
-            core::ptr::swap(hardware_hart.confidential_hart_mut() as *mut ConfidentialHart, confidential_hart as *mut ConfidentialHart);
-        }
+        core::mem::swap(hardware_hart.confidential_hart_mut(), &mut self.confidential_harts[confidential_hart_id].write());
     }
 }