Symmetrical multiprocessing (#78)

* fixing multi-vcpu setups * fix broadcast of IPIs * adjust build process to the new OpenSBI patches * Minor improvements to the build process --------- Signed-off-by: Wojciech Ozga <[email protected]>
IBM · Dec 11, 2024 · 56cf466 · 56cf466
1 parent ea764b5
commit 56cf466
Show file tree

Hide file tree

Showing 15 changed files with 148 additions and 92 deletions.
diff --git a/confidential-vms/linux_vm/Makefile b/confidential-vms/linux_vm/Makefile
@@ -51,6 +51,10 @@ buildroot: setup
 		sed "s@^BR2_TARGET_ROOTFS_EXT2_SIZE=.*@BR2_TARGET_ROOTFS_EXT2_SIZE=\"$(LINUX_VM_BUILDROOT_ROOTFS_SIZE)\"@g" -i $(LINUX_VM_BUILDROOT_WORK_DIR)/.config; \
 		sed "s@^BR2_LINUX_KERNEL_CUSTOM_CONFIG_FILE=.*@BR2_LINUX_KERNEL_CUSTOM_CONFIG_FILE=\"$(LINUX_VM_BUILDROOT_WORK_DIR)/linux64-config\"@g" -i $(LINUX_VM_BUILDROOT_WORK_DIR)/.config; \
 		sed "s@^BR2_LINUX_KERNEL_PATCH=.*@BR2_LINUX_KERNEL_PATCH=\"$(LINUX_VM_PATCHES_DIR)\"@g" -i $(LINUX_VM_BUILDROOT_WORK_DIR)/.config; \
+		if [ ! -f "$(LINUX_VM_BUILDROOT_SOURCE_DIR)/Makefile" ]; then \
+			echo "Buildroot repository not initialized. Did you forget to run: git submodule update --init --recursive?" ;\
+			exit 1 ;\
+		fi ;\
 		$(MAKE) -s -C $(LINUX_VM_BUILDROOT_SOURCE_DIR) RISCV=$(RISCV_GNU_TOOLCHAIN_WORK_DIR) PATH=$(PATH) O=$(LINUX_VM_BUILDROOT_WORK_DIR) CROSS_COMPILE=$(CROSS_COMPILE) BR2_JLEVEL=0 olddefconfig; \
 		$(MAKE) -s -C $(LINUX_VM_BUILDROOT_SOURCE_DIR) RISCV=$(RISCV_GNU_TOOLCHAIN_WORK_DIR) PATH=$(PATH) O=$(LINUX_VM_BUILDROOT_WORK_DIR) BR2_JLEVEL=0; \
 	fi

diff --git a/hypervisor/Makefile b/hypervisor/Makefile
@@ -39,10 +39,15 @@ buildroot: setup
 		mkdir -p $(HYPERVISOR_OVERLAY_ROOT_DIR); \
 		mkdir -p $(HYPERVISOR_OVERLAY_DIR); \
 		cp $(HYPERVISOR_BUILDROOT_CONFIG_DIR) $(BUILDROOT_WORK_DIR)/.config; \
+		echo $(BUILDROOT_WORK_DIR)/.config; \
 		sed "s@^BR2_ROOTFS_OVERLAY=.*@BR2_ROOTFS_OVERLAY=\"$(HYPERVISOR_OVERLAY_DIR)\"@g" -i $(BUILDROOT_WORK_DIR)/.config; \
 		sed "s@^BR2_TARGET_ROOTFS_EXT2_SIZE=.*@BR2_TARGET_ROOTFS_EXT2_SIZE=\"$(HYPERVISOR_ROOTFS_SIZE)\"@g" -i $(BUILDROOT_WORK_DIR)/.config; \
 		sed "s@^BR2_LINUX_KERNEL_CUSTOM_CONFIG_FILE=.*@BR2_LINUX_KERNEL_CUSTOM_CONFIG_FILE=\"$(HYPERVISOR_LINUX_CONFIG)\"@g" -i $(BUILDROOT_WORK_DIR)/.config; \
 		sed "s@^BR2_LINUX_KERNEL_PATCH=.*@BR2_LINUX_KERNEL_PATCH=\"$(HYPERVISOR_LINUX_PATCH)\"@g" -i $(BUILDROOT_WORK_DIR)/.config; \
+		if [ ! -f "$(BUILDROOT_SOURCE_DIR)/Makefile" ]; then \
+			echo "Buildroot repository not initialized. Did you forget to run: git submodule update --init --recursive?" ;\
+			exit 1 ;\
+		fi ;\
 		$(MAKE) -s -C $(BUILDROOT_SOURCE_DIR) RISCV=$(RISCV_GNU_TOOLCHAIN_WORK_DIR) PATH=$(PATH) O=$(BUILDROOT_WORK_DIR) CROSS_COMPILE=$(CROSS_COMPILE) BR2_JLEVEL=0 olddefconfig ;\
 		$(MAKE) -s -C $(BUILDROOT_SOURCE_DIR) RISCV=$(RISCV_GNU_TOOLCHAIN_WORK_DIR) PATH=$(PATH) O=$(BUILDROOT_WORK_DIR) BR2_JLEVEL=0 ;\
 	fi

diff --git a/security-monitor/Makefile b/security-monitor/Makefile
@@ -51,7 +51,7 @@ debug: opensbi_bindings
 
 opensbi: build
 	echo "Compiling OpenSBI" ;\
-	if ! grep -q ace_setup_this_hart "$(OPENSBI_SOURCE_DIR)/lib/sbi/sbi_hart.c"; then \
+	if ! grep -q ace_setup_this_hart "$(OPENSBI_SOURCE_DIR)/lib/sbi/sbi_hsm.c"; then \
 		echo "Applying patches to OpenSBI" ;\
 		cd $(OPENSBI_SOURCE_DIR); git apply --whitespace=fix $(OPENSBI_PATCH); cd $(MAKEFILE_SOURCE_DIR);\
 	fi ;\
@@ -62,7 +62,7 @@ opensbi: build
 
 opensbi_bindings:
 	echo "Compiling OpenSBI bindings" ;\
-	if grep -q ace_setup_this_hart "$(OPENSBI_SOURCE_DIR)/lib/sbi/sbi_hart.c"; then \
+	if grep -q ace_setup_this_hart "$(OPENSBI_SOURCE_DIR)/lib/sbi/sbi_hsm.c"; then \
 		echo "Applying patches to OpenSBI bindings" ;\
 		cd $(OPENSBI_SOURCE_DIR); git apply -R --whitespace=fix $(OPENSBI_PATCH); cd $(MAKEFILE_SOURCE_DIR); \
 	fi; \

diff --git a/security-monitor/platform/generic/configs/defconfig b/security-monitor/platform/generic/configs/defconfig
@@ -1,17 +1,44 @@
-# CONFIG_FDT_GPIO=y
-# CONFIG_FDT_I2C=y
+CONFIG_PLATFORM_ALLWINNER_D1=y
+CONFIG_PLATFORM_ANDES_AE350=y
+CONFIG_PLATFORM_RENESAS_RZFIVE=y
+CONFIG_PLATFORM_SIFIVE_FU540=y
+CONFIG_PLATFORM_SIFIVE_FU740=y
+CONFIG_PLATFORM_SOPHGO_SG2042=y
+CONFIG_PLATFORM_STARFIVE_JH7110=y
+CONFIG_PLATFORM_THEAD=y
+CONFIG_FDT_GPIO=y
+CONFIG_FDT_GPIO_DESIGNWARE=y
+CONFIG_FDT_GPIO_SIFIVE=y
+CONFIG_FDT_GPIO_STARFIVE=y
+CONFIG_FDT_I2C=y
+CONFIG_FDT_I2C_SIFIVE=y
+CONFIG_FDT_I2C_DW=y
 CONFIG_FDT_IPI=y
+CONFIG_FDT_IPI_MSWI=y
+CONFIG_FDT_IPI_PLICSW=y
 CONFIG_FDT_IRQCHIP=y
 CONFIG_FDT_IRQCHIP_APLIC=y
 CONFIG_FDT_IRQCHIP_IMSIC=y
 CONFIG_FDT_IRQCHIP_PLIC=y
-# CONFIG_FDT_REGMAP=y
-# CONFIG_FDT_REGMAP_SYSCON=y
+CONFIG_FDT_REGMAP=y
+CONFIG_FDT_REGMAP_SYSCON=y
 CONFIG_FDT_RESET=y
-# CONFIG_FDT_RESET_GPIO=y
-# CONFIG_FDT_RESET_SYSCON=y
+CONFIG_FDT_RESET_ATCWDT200=y
+CONFIG_FDT_RESET_GPIO=y
+CONFIG_FDT_RESET_HTIF=y
+CONFIG_FDT_RESET_SUNXI_WDT=y
+CONFIG_FDT_RESET_SYSCON=y
 CONFIG_FDT_SERIAL=y
+CONFIG_FDT_SERIAL_CADENCE=y
+CONFIG_FDT_SERIAL_GAISLER=y
+CONFIG_FDT_SERIAL_HTIF=y
+CONFIG_FDT_SERIAL_RENESAS_SCIF=y
+CONFIG_FDT_SERIAL_SHAKTI=y
+CONFIG_FDT_SERIAL_SIFIVE=y
+CONFIG_FDT_SERIAL_LITEX=y
 CONFIG_FDT_SERIAL_UART8250=y
+CONFIG_FDT_SERIAL_XILINX_UARTLITE=y
+CONFIG_SERIAL_SEMIHOSTING=y
 CONFIG_FDT_TIMER=y
 CONFIG_FDT_TIMER_MTIMER=y
-# CONFIG_FDT_TIMER_PLMT=y
+CONFIG_FDT_TIMER_PLMT=y
diff --git a/security-monitor/platform/generic/platform.c b/security-monitor/platform/generic/platform.c
@@ -30,7 +30,7 @@
 #include <sbi_utils/serial/semihosting.h>
 
 /*
-	ACE: init_security_monitor_asm hook implemented in Rust. 	
+	ACE: init_security_monitor_asm hook implemented in Rust.
 */
 extern void init_security_monitor_asm(bool cold_boot, void *fdt);
 
@@ -101,7 +101,7 @@ unsigned long fw_platform_init(unsigned long arg0, unsigned long arg1,
 	int rc, root_offset, cpus_offset, cpu_offset, len;
 	// START ACE
 	const uint8_t *mem_reg_prop;
-	int mem_offset;	
+	int mem_offset;
 	fdt64_t new_size;
 	// END ACE
 
@@ -220,14 +220,12 @@ static int generic_final_init(bool cold_boot)
 			if (rc)
 				return rc;
 		}
+		// ACE start
+		// This is the entry point of the security monitor's initialization procedure.
+		init_security_monitor_asm(cold_boot, fdt);
+		// ACE end
 	}
 
-	fdt = fdt_get_address();	
-	// ACE start
-	// This is the entry point of the security monitor's initialization procedure.
-	init_security_monitor_asm(cold_boot, fdt);
-	// ACE end
-
 	return 0;
 }
 

diff --git a/security-monitor/platform/opensbi_v1.4.patch b/security-monitor/platform/opensbi_v1.4.patch
@@ -1,50 +1,56 @@
 diff --git a/lib/sbi/sbi_hart.c b/lib/sbi/sbi_hart.c
-index 770fee0..663eec7 100644
+index 770fee0..9e08308 100644
 --- a/lib/sbi/sbi_hart.c
 +++ b/lib/sbi/sbi_hart.c
-@@ -24,6 +24,7 @@
- #include <sbi/sbi_trap.h>
- #include <sbi/sbi_hfence.h>
-
-+extern void ace_setup_this_hart();
- extern void __sbi_expected_trap(void);
- extern void __sbi_expected_trap_hext(void);
-
-@@ -369,7 +370,7 @@ static int sbi_hart_smepmp_configure(struct sbi_scratch *scratch,
+@@ -369,7 +369,7 @@ static int sbi_hart_smepmp_configure(struct sbi_scratch *scratch,
  	pmp_disable(SBI_SMEPMP_RESV_ENTRY);
- 
+
  	/* Program M-only regions when MML is not set. */
 -	pmp_idx = 0;
 +	pmp_idx = 2;
  	sbi_domain_for_each_memregion(dom, reg) {
  		/* Skip reserved entry */
  		if (pmp_idx == SBI_SMEPMP_RESV_ENTRY)
-@@ -395,7 +396,7 @@ static int sbi_hart_smepmp_configure(struct sbi_scratch *scratch,
+@@ -395,7 +395,7 @@ static int sbi_hart_smepmp_configure(struct sbi_scratch *scratch,
  	csr_set(CSR_MSECCFG, MSECCFG_MML);
- 
+
  	/* Program shared and SU-only regions */
 -	pmp_idx = 0;
 +	pmp_idx = 2;
  	sbi_domain_for_each_memregion(dom, reg) {
  		/* Skip reserved entry */
  		if (pmp_idx == SBI_SMEPMP_RESV_ENTRY)
-@@ -432,7 +433,7 @@ static int sbi_hart_oldpmp_configure(struct sbi_scratch *scratch,
+@@ -432,7 +432,7 @@ static int sbi_hart_oldpmp_configure(struct sbi_scratch *scratch,
  {
  	struct sbi_domain_memregion *reg;
  	struct sbi_domain *dom = sbi_domain_thishart_ptr();
 -	unsigned int pmp_idx = 0;
 +	unsigned int pmp_idx = 2;
  	unsigned int pmp_flags;
  	unsigned long pmp_addr;
-
-@@ -534,6 +535,10 @@ int sbi_hart_pmp_configure(struct sbi_scratch *scratch)
- 		rc = sbi_hart_oldpmp_configure(scratch, pmp_count,
- 						pmp_log2gran, pmp_addr_max);
-
+
+diff --git a/lib/sbi/sbi_hsm.c b/lib/sbi/sbi_hsm.c
+index 3d60ceb..45c602d 100644
+--- a/lib/sbi/sbi_hsm.c
++++ b/lib/sbi/sbi_hsm.c
+@@ -26,6 +26,8 @@
+ #include <sbi/sbi_timer.h>
+ #include <sbi/sbi_console.h>
+
++extern void ace_setup_this_hart();
++
+ #define __sbi_hsm_hart_change_state(hdata, oldstate, newstate)		\
+ ({									\
+ 	long state = atomic_cmpxchg(&(hdata)->state, oldstate, newstate); \
+@@ -154,6 +156,11 @@ void __noreturn sbi_hsm_hart_start_finish(struct sbi_scratch *scratch,
+ 	next_mode = scratch->next_mode;
+ 	hsm_start_ticket_release(hdata);
+
 +	// ACE START
 +	// temporal hack to reconfigure PMPs after they have been reconfigured by OpenSBI.
 +	ace_setup_this_hart();
 +	// ACE END
- 	/*
- 	 * As per section 3.7.2 of privileged specification v1.12,
- 	 * virtual address translations can be speculatively performed
++
+ 	sbi_hart_switch_mode(hartid, next_arg1, next_addr, next_mode, false);
+ }
+
diff --git a/security-monitor/src/confidential_flow/finite_state_machine.rs b/security-monitor/src/confidential_flow/finite_state_machine.rs
@@ -28,7 +28,8 @@ use crate::core::architecture::riscv::sbi::SrstExtension::*;
 use crate::core::architecture::TrapCause::*;
 use crate::core::architecture::{HartLifecycleState, TrapCause};
 use crate::core::control_data::{
-    ConfidentialHart, ConfidentialHartRemoteCommand, ConfidentialVmId, ControlDataStorage, HardwareHart, HypervisorHart, ResumableOperation,
+    ConfidentialHart, ConfidentialHartRemoteCommand, ConfidentialVm, ConfidentialVmId, ControlDataStorage, HardwareHart, HypervisorHart,
+    ResumableOperation,
 };
 use crate::error::Error;
 use crate::non_confidential_flow::{DeclassifyToHypervisor, NonConfidentialFlow};
@@ -99,7 +100,13 @@ impl<'a> ConfidentialFlow<'a> {
             VirtualInstruction => VirtualInstruction::from_confidential_hart(flow.confidential_hart()).handle(flow),
             GuestStorePageFault => MmioStoreRequest::from_confidential_hart(flow.confidential_hart()).handle(flow),
             trap_reason => {
-                debug!("Bug: Not supported trap cause {:?}, maybe due to incorrect exception delegation?", trap_reason);
+                debug!(
+                    "Bug when executing confidential hart {}. Not supported trap cause {:?}. mepc={:x} mtval={:x}",
+                    flow.confidential_hart().confidential_hart_id(),
+                    trap_reason,
+                    flow.confidential_hart().csrs().mepc.read_from_main_memory(),
+                    flow.confidential_hart().csrs().mtval.read()
+                );
                 ShutdownRequest::from_confidential_hart(flow.confidential_hart()).handle(flow)
             }
         }
@@ -206,17 +213,17 @@ impl<'a> ConfidentialFlow<'a> {
 impl<'a> ConfidentialFlow<'a> {
     /// Broadcasts the inter hart request to confidential harts of the currently executing confidential VM. Returns error if sending an IPI
     /// to other confidential hart failed or if there is too many pending IPI queued.
-    pub fn broadcast_remote_command(&mut self, confidential_hart_remote_command: ConfidentialHartRemoteCommand) -> Result<(), Error> {
-        ControlDataStorage::try_confidential_vm_mut(self.confidential_vm_id(), |mut confidential_vm| {
-            // Hack: For the time-being, we rely on the OpenSBI's implementation of physical IPIs. To use OpenSBI functions we
-            // must set the mscratch register to the value expected by OpenSBI. We do it here, because we have access to the `HardwareHart`
-            // that knows the original value of the mscratch expected by OpenSBI.
-            self.hardware_hart.swap_mscratch();
-            let result = confidential_vm.broadcast_remote_command(confidential_hart_remote_command);
-            // We must revert the content of mscratch back to the value expected by our context switched.
-            self.hardware_hart.swap_mscratch();
-            result
-        })
+    pub fn broadcast_remote_command(
+        &mut self, confidential_vm: &mut ConfidentialVm, confidential_hart_remote_command: ConfidentialHartRemoteCommand,
+    ) -> Result<(), Error> {
+        // Hack: For the time-being, we rely on the OpenSBI's implementation of physical IPIs. To use OpenSBI functions we
+        // must set the mscratch register to the value expected by OpenSBI. We do it here, because we have access to the `HardwareHart`
+        // that knows the original value of the mscratch expected by OpenSBI.
+        self.hardware_hart.swap_mscratch();
+        let result = confidential_vm.broadcast_remote_command(confidential_hart_remote_command);
+        // We must revert the content of mscratch back to the value expected by our context switched.
+        self.hardware_hart.swap_mscratch();
+        result
     }
 
     /// Processes pending requests from other confidential harts by applying the corresponding state transformation to

diff --git a/security-monitor/src/confidential_flow/handlers/shared_page/share_page_complete.rs b/security-monitor/src/confidential_flow/handlers/shared_page/share_page_complete.rs
@@ -30,23 +30,24 @@ impl SharePageComplete {
         }
     }
 
-    pub fn handle(self, confidential_flow: ConfidentialFlow) -> ! {
+    pub fn handle(self, mut confidential_flow: ConfidentialFlow) -> ! {
         let transformation = self
-            .map_shared_page(&confidential_flow)
+            .map_shared_page(&mut confidential_flow)
             .and_then(|_| Ok(SbiResponse::success()))
             .unwrap_or_else(|error| SbiResponse::error(error));
         confidential_flow.apply_and_exit_to_confidential_hart(ApplyToConfidentialHart::SbiResponse(transformation))
     }
 
-    fn map_shared_page(&self, confidential_flow: &ConfidentialFlow) -> Result<(), Error> {
+    fn map_shared_page(&self, confidential_flow: &mut ConfidentialFlow) -> Result<(), Error> {
         ensure!(self.response_code == 0, Error::Failed())?;
         // Security: check that the start address is located in the non-confidential memory
         let hypervisor_address = NonConfidentialMemoryAddress::new(self.hypervisor_page_address as *mut usize)?;
 
         ControlDataStorage::try_confidential_vm_mut(confidential_flow.confidential_vm_id(), |mut confidential_vm| {
             let page_size = confidential_vm.memory_protector_mut().map_shared_page(hypervisor_address, self.request.address)?;
             let request = RemoteHfenceGvmaVmid::all_harts(&self.request.address, page_size, confidential_flow.confidential_vm_id());
-            confidential_vm.broadcast_remote_command(ConfidentialHartRemoteCommand::RemoteHfenceGvmaVmid(request))?;
+            confidential_flow
+                .broadcast_remote_command(&mut confidential_vm, ConfidentialHartRemoteCommand::RemoteHfenceGvmaVmid(request))?;
             Ok(())
         })
     }

diff --git a/security-monitor/src/confidential_flow/handlers/shared_page/unshare_page_request.rs b/security-monitor/src/confidential_flow/handlers/shared_page/unshare_page_request.rs
@@ -6,9 +6,7 @@ use crate::confidential_flow::handlers::symmetrical_multiprocessing::RemoteHfenc
 use crate::confidential_flow::{ApplyToConfidentialHart, ConfidentialFlow};
 use crate::core::architecture::riscv::sbi::CovgExtension;
 use crate::core::architecture::{GeneralPurposeRegister, SharedPage};
-use crate::core::control_data::{
-    ConfidentialHart, ConfidentialHartRemoteCommand, ConfidentialVmId, ControlDataStorage, ResumableOperation,
-};
+use crate::core::control_data::{ConfidentialHart, ConfidentialHartRemoteCommand, ControlDataStorage, ResumableOperation};
 use crate::core::memory_layout::ConfidentialVmPhysicalAddress;
 use crate::error::Error;
 use crate::non_confidential_flow::DeclassifyToHypervisor;
@@ -27,8 +25,8 @@ impl UnsharePageRequest {
         }
     }
 
-    pub fn handle(self, confidential_flow: ConfidentialFlow) -> ! {
-        match self.unmap_shared_page(confidential_flow.confidential_vm_id()) {
+    pub fn handle(self, mut confidential_flow: ConfidentialFlow) -> ! {
+        match self.unmap_shared_page(&mut confidential_flow) {
             Ok(_) => confidential_flow
                 .set_resumable_operation(ResumableOperation::SbiRequest())
                 .into_non_confidential_flow()
@@ -43,14 +41,16 @@ impl UnsharePageRequest {
         SbiRequest::new(CovgExtension::EXTID, CovgExtension::SBI_EXT_COVG_UNSHARE_MEMORY, self.address.usize(), self.size)
     }
 
-    fn unmap_shared_page(&self, confidential_vm_id: ConfidentialVmId) -> Result<(), Error> {
+    fn unmap_shared_page(&self, confidential_flow: &mut ConfidentialFlow) -> Result<(), Error> {
         ensure!(self.address.usize() % SharedPage::SIZE.in_bytes() == 0, Error::AddressNotAligned())?;
         ensure!(self.size == SharedPage::SIZE.in_bytes(), Error::InvalidParameter())?;
 
+        let confidential_vm_id = confidential_flow.confidential_vm_id();
         ControlDataStorage::try_confidential_vm_mut(confidential_vm_id, |mut confidential_vm| {
             let unmapped_page_size = confidential_vm.memory_protector_mut().unmap_shared_page(&self.address)?;
             let request = RemoteHfenceGvmaVmid::all_harts(&self.address, unmapped_page_size, confidential_vm_id);
-            confidential_vm.broadcast_remote_command(ConfidentialHartRemoteCommand::RemoteHfenceGvmaVmid(request))?;
+            confidential_flow
+                .broadcast_remote_command(&mut confidential_vm, ConfidentialHartRemoteCommand::RemoteHfenceGvmaVmid(request))?;
             Ok(())
         })
     }

diff --git a/security-monitor/src/confidential_flow/handlers/shutdown/shutdown_vm.rs b/security-monitor/src/confidential_flow/handlers/shutdown/shutdown_vm.rs
@@ -4,7 +4,7 @@
 use crate::confidential_flow::handlers::sbi::SbiResponse;
 use crate::confidential_flow::handlers::shutdown::shutdown_confidential_hart;
 use crate::confidential_flow::{ApplyToConfidentialHart, ConfidentialFlow};
-use crate::core::control_data::{ConfidentialHart, ConfidentialHartRemoteCommand};
+use crate::core::control_data::{ConfidentialHart, ConfidentialHartRemoteCommand, ControlDataStorage};
 
 /// Handles the system reset call of the SBI's SRST extension. This call is a request to shutdown or reboot the
 /// confidential virtual machine. The security monitor allows only for the full shutdown of the confidential virtual
@@ -24,11 +24,12 @@ impl ShutdownRequest {
     }
 
     pub fn handle(self, mut confidential_flow: ConfidentialFlow) -> ! {
-        match confidential_flow.broadcast_remote_command(ConfidentialHartRemoteCommand::ShutdownRequest(self)) {
+        match ControlDataStorage::try_confidential_vm_mut(confidential_flow.confidential_vm_id(), |mut confidential_vm| {
+            confidential_flow.broadcast_remote_command(&mut confidential_vm, ConfidentialHartRemoteCommand::ShutdownRequest(self))
+        }) {
             Ok(_) => shutdown_confidential_hart(confidential_flow),
             Err(error) => {
-                let transformation = ApplyToConfidentialHart::SbiResponse(SbiResponse::error(error));
-                confidential_flow.apply_and_exit_to_confidential_hart(transformation)
+                confidential_flow.apply_and_exit_to_confidential_hart(ApplyToConfidentialHart::SbiResponse(SbiResponse::error(error)))
             }
         }
     }

diff --git a/security-monitor/src/confidential_flow/handlers/symmetrical_multiprocessing/fence_i.rs b/security-monitor/src/confidential_flow/handlers/symmetrical_multiprocessing/fence_i.rs
@@ -4,7 +4,9 @@
 use crate::confidential_flow::handlers::sbi::SbiResponse;
 use crate::confidential_flow::handlers::symmetrical_multiprocessing::Ipi;
 use crate::confidential_flow::{ApplyToConfidentialHart, ConfidentialFlow};
-use crate::core::control_data::{ConfidentialHart, ConfidentialHartRemoteCommand, ConfidentialHartRemoteCommandExecutable};
+use crate::core::control_data::{
+    ConfidentialHart, ConfidentialHartRemoteCommand, ConfidentialHartRemoteCommandExecutable, ControlDataStorage,
+};
 
 /// Handles a request from one confidential hart to execute fence.i instruction on remote confidential harts.
 #[derive(Clone)]
@@ -18,11 +20,11 @@ impl RemoteFenceI {
     }
 
     pub fn handle(self, mut confidential_flow: ConfidentialFlow) -> ! {
-        let transformation = confidential_flow
-            .broadcast_remote_command(ConfidentialHartRemoteCommand::RemoteFenceI(self))
-            .and_then(|_| Ok(SbiResponse::success()))
-            .unwrap_or_else(|error| SbiResponse::error(error));
-        confidential_flow.apply_and_exit_to_confidential_hart(ApplyToConfidentialHart::SbiResponse(transformation))
+        let result = ControlDataStorage::try_confidential_vm_mut(confidential_flow.confidential_vm_id(), |mut confidential_vm| {
+            confidential_flow.broadcast_remote_command(&mut confidential_vm, ConfidentialHartRemoteCommand::RemoteFenceI(self))
+        })
+        .map_or_else(|error| SbiResponse::error(error), |_| SbiResponse::success());
+        confidential_flow.apply_and_exit_to_confidential_hart(ApplyToConfidentialHart::SbiResponse(result))
     }
 }