Improve interaction with OpenSBI

security-monitor/src/confidential_flow/finite_state_machine.rs

@@ -25,8 +25,9 @@ use crate::core::architecture::riscv::sbi::IpiExtension::*;
 use crate::core::architecture::riscv::sbi::RfenceExtension::*;
 use crate::core::architecture::riscv::sbi::SbiExtension::*;
 use crate::core::architecture::riscv::sbi::SrstExtension::*;
+use crate::core::architecture::specification::CSR_MSTATUS_MPV;
 use crate::core::architecture::TrapCause::*;
-use crate::core::architecture::{HartLifecycleState, TrapCause};
+use crate::core::architecture::{is_bit_enabled, HartLifecycleState, TrapCause, CSR};
 use crate::core::control_data::{
     ConfidentialHart, ConfidentialHartRemoteCommand, ConfidentialVm, ConfidentialVmId, ControlDataStorage, HardwareHart, HypervisorHart,
     ResumableOperation,
@@ -54,7 +55,6 @@ pub struct ConfidentialFlow<'a> {
 }
 
 impl<'a> ConfidentialFlow<'a> {
-    const CTX_SWITCH_ERROR_MSG: &'static str = "Bug: Invalid argument provided by the assembly context switch";
     const DUMMY_HART_ERROR_MSG: &'static str = "Bug: Found dummy hart instead of a confidential hart";
 
     /// Routes the control flow to a handler that will process the confidential hart interrupt or exception. This is an entry point to
@@ -65,8 +65,18 @@ impl<'a> ConfidentialFlow<'a> {
     /// ConfidentialFlow but accessible to the assembly code performing the context switch.
     #[no_mangle]
     unsafe extern "C" fn route_trap_from_confidential_hart(hardware_hart_pointer: *mut HardwareHart) -> ! {
-        let flow = Self { hardware_hart: unsafe { hardware_hart_pointer.as_mut().expect(Self::CTX_SWITCH_ERROR_MSG) } };
+        let flow = Self { hardware_hart: unsafe { &mut *hardware_hart_pointer } };
         assert!(!flow.hardware_hart.confidential_hart().is_dummy());
+
+        if !is_bit_enabled(CSR.mstatus.read(), CSR_MSTATUS_MPV) {
+            debug!(
+                "TSM exception {} when executing confidential hart {}",
+                flow.confidential_hart().confidential_hart_state().csrs().mcause.read(),
+                flow.confidential_hart().confidential_hart_id(),
+            );
+            ShutdownRequest::from_confidential_hart(flow.confidential_hart()).handle(flow)
+        }
+
         match TrapCause::from_hart_architectural_state(flow.confidential_hart().confidential_hart_state()) {
             Interrupt => HandleInterrupt::from_confidential_hart(flow.confidential_hart()).handle(flow),
             VsEcall(Base(GetSpecVersion)) => SbiGetSpecVersion::from_confidential_hart(flow.confidential_hart()).handle(flow),
@@ -100,13 +110,7 @@ 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 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()
-                );
+                debug!("Not supported trap cause {:?}", trap_reason);
                 ShutdownRequest::from_confidential_hart(flow.confidential_hart()).handle(flow)
             }
         }

security-monitor/src/core/architecture/riscv/control_status_registers.rs

@@ -301,6 +301,8 @@ pub struct ControlStatusRegister {
     pub mvendorid: ReadWriteRiscvCsr<CSR_MVENDORID>,
     pub marchid: ReadWriteRiscvCsr<CSR_MARCHID>,
     pub mimpid: ReadWriteRiscvCsr<CSR_MIMPID>,
+    pub mcause: ReadWriteRiscvCsr<CSR_MCAUSE>,
+    pub mstatus: ReadWriteRiscvCsr<CSR_MSTATUS>,
     pub mscratch: ReadWriteRiscvCsr<CSR_MSCRATCH>,
     pub hgatp: ReadWriteRiscvCsr<CSR_HGATP>,
     pub pmpcfg0: ReadWriteRiscvCsr<CSR_PMPCFG0>,
@@ -313,6 +315,8 @@ pub const CSR: &ControlStatusRegister = &ControlStatusRegister {
     mvendorid: ReadWriteRiscvCsr::new(),
     marchid: ReadWriteRiscvCsr::new(),
     mimpid: ReadWriteRiscvCsr::new(),
+    mcause: ReadWriteRiscvCsr::new(),
+    mstatus: ReadWriteRiscvCsr::new(),
     mscratch: ReadWriteRiscvCsr::new(),
     hgatp: ReadWriteRiscvCsr::new(),
     pmpcfg0: ReadWriteRiscvCsr::new(),
@@ -415,6 +419,17 @@ impl<const V: u16> ReadWriteRiscvCsr<V> {
         }
         r
     }
+
+    pub fn swap(&self, value: usize) -> usize {
+        let r: usize;
+        unsafe {
+            asm!("csrrw {rd}, {csr}, {rs1}",
+                     rd = out(reg) r,
+                     csr = const V,
+                     rs1 = in(reg) value);
+        }
+        r
+    }
 }
 
 #[derive(Copy, Clone)]

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

@@ -5,6 +5,13 @@ use crate::core::architecture::CSR;
 use crate::core::control_data::{ConfidentialHart, HypervisorHart};
 use crate::core::memory_protector::HypervisorMemoryProtector;
 use crate::core::page_allocator::{Allocated, Page, UnAllocated};
+use opensbi_sys::sbi_trap_regs;
+
+extern "C" {
+    /// Currently, we rely on OpenSBI to handle some of the interrupts or exceptions. Below function is the entry point
+    /// to OpenSBI trap handler.
+    fn sbi_trap_handler(regs: *mut sbi_trap_regs) -> *mut sbi_trap_regs;
+}
 
 pub const HART_STACK_ADDRESS_OFFSET: usize = memoffset::offset_of!(HardwareHart, stack_address);
 
@@ -46,24 +53,33 @@ impl HardwareHart {
         }
     }
 
-    /// Calling OpenSBI handler to process the SBI call requires setting the mscratch register to a specific value which
-    /// we replaced during the system initialization. We store the original mscratch value expected by the OpenSBI in
-    /// the previous_mscratch field.
-    pub fn swap_mscratch(&mut self) {
-        let previous_mscratch = self.previous_mscratch;
-        let current_mscratch = CSR.mscratch.read();
-        CSR.mscratch.write(previous_mscratch);
-        self.previous_mscratch = current_mscratch;
+    // Safety: this function can be called only once during initialization. We cannot do it in constructor, because physical harts art not
+    // initialized yet then.
+    pub fn set_ace_mscratch(&mut self, value: usize) {
+        self.previous_mscratch = CSR.mscratch.swap(value);
     }
 
+    #[inline(always)]
     pub fn opensbi_context<F>(&mut self, function: F) -> Result<(), crate::error::Error>
     where F: FnOnce() -> Result<(), crate::error::Error> {
-        self.swap_mscratch();
+        let ace_mscratch = CSR.mscratch.swap(self.previous_mscratch);
         let result = function();
-        self.swap_mscratch();
+        CSR.mscratch.write(ace_mscratch);
         result
     }
 
+    #[inline(always)]
+    pub fn call_opensbi_trap_handler(&mut self) {
+        // Safety: We play with fire here. We must make sure statically that the OpenSBI's input structure is bitwise same as the ACE's hart
+        // state.
+        let trap_regs = self.hypervisor_hart_mut().hypervisor_hart_state_mut() as *mut _ as *mut sbi_trap_regs;
+        let _ = self.opensbi_context(|| {
+            Ok(unsafe {
+                sbi_trap_handler(trap_regs);
+            })
+        });
+    }
+
     pub fn confidential_hart(&self) -> &ConfidentialHart {
         &self.confidential_hart
     }

security-monitor/src/core/initialization/mod.rs

@@ -229,10 +229,8 @@ extern "C" fn ace_setup_this_hart() {
     // opensbi_mscratch in the internal hart state. OpenSBI stored in mscratch a pointer to the
     // `opensbi_mscratch` region of this hart before calling the security monitor's initialization
     // procedure. Thus, the swap will move the mscratch register value into the dump state of the hart
-    hart.swap_mscratch();
-    let hart_address = hart.hypervisor_hart().address();
-    hart.hypervisor_hart_mut().csrs_mut().mscratch.write(hart_address);
-    debug!("Hardware hart id={} has state area region at {:x}", hart_id, hart_address);
+    hart.set_ace_mscratch(hart.hypervisor_hart().address());
+    debug!("Hardware hart id={} has state area region at {:x}", hart_id, hart.hypervisor_hart().address());
 
     // Configure the memory isolation mechanism that can limit memory view of the hypervisor to the memory region
     // owned by the hypervisor. The setup method enables the memory isolation. It is safe to call it because

security-monitor/src/non_confidential_flow/finite_state_machine.rs

@@ -7,8 +7,9 @@ use crate::core::architecture::riscv::sbi::CovhExtension::*;
 use crate::core::architecture::riscv::sbi::NaclExtension::*;
 use crate::core::architecture::riscv::sbi::NaclSharedMemory;
 use crate::core::architecture::riscv::sbi::SbiExtension::*;
-use crate::core::architecture::TrapCause;
+use crate::core::architecture::specification::CAUSE_SUPERVISOR_ECALL;
 use crate::core::architecture::TrapCause::*;
+use crate::core::architecture::{TrapCause, CSR};
 use crate::core::control_data::{ConfidentialVmId, HardwareHart, HypervisorHart};
 use crate::error::Error;
 use crate::non_confidential_flow::handlers::cove_host_extension::{
@@ -24,10 +25,6 @@ extern "C" {
     /// never returns to KVM with other state. For example, only a subset of exceptions/interrupts can be handled by KVM.
     /// KVM kill the vcpu if it receives unexpected exception because it does not know what to do with it.
     fn exit_to_hypervisor_asm() -> !;
-
-    /// Currently, we rely on OpenSBI to handle some of the interrupts or exceptions. Below function is the entry point
-    /// to OpenSBI trap handler.
-    fn sbi_trap_handler(regs: *mut sbi_trap_regs) -> *mut sbi_trap_regs;
 }
 
 /// Represents the non-confidential part of the finite state machine (FSM), implementing router and exit nodes. It encapsulates the
@@ -37,8 +34,6 @@ pub struct NonConfidentialFlow<'a> {
 }
 
 impl<'a> NonConfidentialFlow<'a> {
-    const CTX_SWITCH_ERROR_MSG: &'static str = "Bug: invalid argument provided by the assembly context switch";
-
     /// Creates an instance of the `NonConfidentialFlow`. A confidential hart must not be assigned to the hardware hart.
     pub fn create(hardware_hart: &'a mut HardwareHart) -> Self {
         assert!(hardware_hart.confidential_hart().is_dummy());
@@ -59,7 +54,15 @@ impl<'a> NonConfidentialFlow<'a> {
         // hardware hart's dump area in main memory. This area in main memory is exclusively owned by the physical hart executing this code.
         // Specifically, every physical hart has its own are in the main memory and its `mscratch` register stores the address. See the
         // `initialization` procedure for more details.
-        let mut flow = unsafe { Self::create(hart_ptr.as_mut().expect(Self::CTX_SWITCH_ERROR_MSG)) };
+        let hardware_hart = unsafe { &mut *hart_ptr };
+        // Performance optimziation: we do not want to add overhead when delegating calls to OpenSBI, thus make sure we do not
+        // create extra objects on heap or make unnecessary load/stores.
+        if CSR.mcause.read() != CAUSE_SUPERVISOR_ECALL.into() {
+            hardware_hart.call_opensbi_trap_handler();
+            unsafe { exit_to_hypervisor_asm() };
+        }
+
+        let mut flow = Self::create(hardware_hart);
         match TrapCause::from_hart_architectural_state(flow.hypervisor_hart().hypervisor_hart_state()) {
             HsEcall(Base(ProbeExtension)) => ProbeSbiExtension::from_hypervisor_hart(flow.hypervisor_hart_mut()).handle(flow),
             HsEcall(Covh(TsmGetInfo)) => GetSecurityMonitorInfo::from_hypervisor_hart(flow.hypervisor_hart()).handle(flow),
@@ -75,13 +78,7 @@ impl<'a> NonConfidentialFlow<'a> {
     }
 
     pub fn delegate_to_opensbi(self) -> ! {
-        // Safety: We play with fire here. We must statically make sure that OpenSBI's input structure is bitwise same as ACE's hart state.
-        let trap_regs = self.hardware_hart.hypervisor_hart_mut().hypervisor_hart_state_mut() as *mut _ as *mut sbi_trap_regs;
-        let _ = self.hardware_hart.opensbi_context(|| {
-            Ok(unsafe {
-                sbi_trap_handler(trap_regs);
-            })
-        });
+        self.hardware_hart.call_opensbi_trap_handler();
         unsafe { exit_to_hypervisor_asm() }
     }