Fix RISC-V vector state preservation across context switches

security-monitor/src/confidential_flow/lightweight_context_switch/enter_from_confidential_hart.S

@@ -62,7 +62,9 @@ enter_from_confidential_hart_asm:
     # Recover the stack pointer for this hart. 
     # The stack pointer is stored in the memory dump area of the hart (pointed by `mscratch`)
     csrr        a0, mscratch
-    ld	        sp, ({HART_STACK_ADDRESS_OFFSET})(a0)
+    li          t0, {HART_STACK_ADDRESS_OFFSET}
+    add         t0, a0, t0
+    ld          sp, 0(t0)
 
     # Before calling Rust code, we must guarantee that: 
     #   * `mscratch` points to the dump area of this hardware hart

security-monitor/src/core/architecture/riscv/extensions/vector_unit.rs

@@ -1,16 +1,131 @@
 // SPDX-FileCopyrightText: 2023 IBM Corporation
 // SPDX-FileContributor: Wojciech Ozga <woz@zurich.ibm.com>, IBM Research - Zurich
 // SPDX-License-Identifier: Apache-2.0
-#![allow(unused)]
+
+use crate::core::architecture::riscv::control_status_registers::ReadWriteRiscvCsr;
+use crate::core::architecture::riscv::specification::{CSR_VCSR, CSR_VL, CSR_VSTART, CSR_VTYPE};
+use crate::core::hardware_setup::HardwareSetup;
+use crate::error::Error;
 use alloc::vec::Vec;
-use core::ops::Range;
+use core::arch::asm;
 
-pub const MAX_NUMBER_OF_REGISTER_LENGTH: usize = 32;
-const NUMBER_OF_VALUES_IN_REGISTER: usize = MAX_NUMBER_OF_REGISTER_LENGTH >> 3;
 const NUMBER_OF_VECTOR_REGISTERS: usize = 32;
 
 #[repr(C)]
-pub struct VectorRegister(pub [u64; NUMBER_OF_VALUES_IN_REGISTER]);
+pub struct VectorUnit {
+    registers: Vec<u8>,
+    vlenb: usize,
+    vl: ReadWriteRiscvCsr<CSR_VL>,
+    vtype: ReadWriteRiscvCsr<CSR_VTYPE>,
+    vstart: ReadWriteRiscvCsr<CSR_VSTART>,
+    vcsr: ReadWriteRiscvCsr<CSR_VCSR>,
+}
 
-#[repr(C)]
-pub struct VectorRegisters(pub [usize; NUMBER_OF_VECTOR_REGISTERS]);
+impl VectorUnit {
+    pub fn empty() -> Result<Self, Error> {
+        let vlenb = HardwareSetup::vector_vlenb()?;
+
+        let length = NUMBER_OF_VECTOR_REGISTERS.checked_mul(vlenb).ok_or(Error::InvalidVectorRegisterSize())?;
+
+        let mut registers = Vec::new();
+        registers.try_reserve_exact(length).map_err(|_| Error::NotEnoughMemory())?;
+        registers.resize(length, 0);
+
+        Ok(Self {
+            registers,
+            vlenb,
+            vl: ReadWriteRiscvCsr::new(),
+            vtype: ReadWriteRiscvCsr::new(),
+            vstart: ReadWriteRiscvCsr::new(),
+            vcsr: ReadWriteRiscvCsr::new(),
+        })
+    }
+
+    /// Saves vector state in main memory.
+    ///
+    /// # Safety
+    ///
+    /// The vector extension must be supported and enabled.
+    pub unsafe fn save_in_main_memory(&mut self) {
+        self.vl.save_in_main_memory();
+        self.vtype.save_in_main_memory();
+        self.vstart.save_in_main_memory();
+        self.vcsr.save_in_main_memory();
+
+        self.vstart.write(0);
+
+        macro_rules! store {
+            ($index:literal) => {
+                unsafe {
+                    asm!(
+                        ".option push",
+                        ".option arch, +v",
+                        concat!(
+                            "vs8r.v v",
+                            stringify!($index),
+                            ", ({address})"
+                        ),
+                        ".option pop",
+                        address = in(reg) self.registers
+                            .as_mut_ptr()
+                            .add($index * self.vlenb),
+                    );
+                }
+            };
+        }
+
+        store!(0);
+        store!(8);
+        store!(16);
+        store!(24);
+    }
+
+    /// Restores vector state from main memory.
+    ///
+    /// # Safety
+    ///
+    /// The vector extension must be supported and enabled.
+    /// The current hart must have a compatible vector configuration.
+    pub unsafe fn restore_from_main_memory(&self) {
+        self.vstart.write(0);
+
+        macro_rules! load {
+            ($index:literal) => {
+                unsafe {
+                    asm!(
+                        ".option push",
+                        ".option arch, +v",
+                        concat!(
+                            "vl8re8.v v",
+                            stringify!($index),
+                            ", ({address})"
+                        ),
+                        ".option pop",
+                        address = in(reg) self.registers
+                            .as_ptr()
+                            .add($index * self.vlenb),
+                    );
+                }
+            };
+        }
+
+        load!(0);
+        load!(8);
+        load!(16);
+        load!(24);
+
+        unsafe {
+            asm!(
+                ".option push",
+                ".option arch, +v",
+                "vsetvl zero, {saved_vl}, {saved_vtype}",
+                ".option pop",
+                saved_vl = in(reg) self.vl.read_from_main_memory(),
+                saved_vtype = in(reg) self.vtype.read_from_main_memory(),
+            );
+        }
+
+        self.vcsr.restore_from_main_memory();
+        self.vstart.restore_from_main_memory();
+    }
+}

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

@@ -3,7 +3,10 @@
 // SPDX-License-Identifier: Apache-2.0
 use crate::core::architecture::riscv::extensions::supervisor_timer_extension::SupervisorTimerExtension;
 use crate::core::architecture::riscv::specification::*;
-use crate::core::architecture::riscv::{ControlStatusRegisters, FloatingPointUnit, GeneralPurposeRegister, GeneralPurposeRegisters};
+use crate::core::architecture::riscv::{
+    ControlStatusRegisters, FloatingPointUnit, GeneralPurposeRegister, GeneralPurposeRegisters, VectorUnit,
+};
+use crate::error::Error;
 
 /// Defines the state of a processor's core (hart) when stored in main memory.
 #[repr(C)]
@@ -12,16 +15,26 @@ pub struct HartArchitecturalState {
     csrs: ControlStatusRegisters,
     fprs: FloatingPointUnit,
     sstc: SupervisorTimerExtension,
+    vrs: VectorUnit,
 }
 
 impl HartArchitecturalState {
-    pub fn empty() -> Self {
-        Self {
+    pub fn empty() -> Result<Self, Error> {
+        Ok(Self {
             gprs: GeneralPurposeRegisters::empty(),
             csrs: ControlStatusRegisters::empty(),
             fprs: FloatingPointUnit::empty(),
             sstc: SupervisorTimerExtension::empty(),
-        }
+            vrs: VectorUnit::empty()?,
+        })
+    }
+
+    pub fn vrs(&self) -> &VectorUnit {
+        &self.vrs
+    }
+
+    pub fn vrs_mut(&mut self) -> &mut VectorUnit {
+        &mut self.vrs
     }
 
     pub fn set_gprs(&mut self, gprs: GeneralPurposeRegisters) {

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

@@ -2,11 +2,13 @@
 // SPDX-FileContributor: Wojciech Ozga <woz@zurich.ibm.com>, IBM Research - Zurich
 // SPDX-License-Identifier: Apache-2.0
 #![allow(unused)]
+pub(crate) use control_status_registers::ReadWriteRiscvCsr;
 pub use control_status_registers::{CSR, ControlStatusRegister, ControlStatusRegisters};
 pub use extensions::HardwareExtension;
 pub use extensions::compressed_instructions::decode_result_register;
 pub use extensions::floating_point_unit::FloatingPointUnit;
 pub use extensions::supervisor_timer_extension::SupervisorTimerExtension;
+pub use extensions::vector_unit::VectorUnit;
 pub use general_purpose_registers::{GeneralPurposeRegister, GeneralPurposeRegisters};
 pub use hart_architectural_state::HartArchitecturalState;
 pub use hart_lifecycle_state::HartLifecycleState;

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

@@ -422,6 +422,9 @@ pub const MTVEC_BASE_SHIFT: usize = 2;
 
 pub const CSR_STATUS_SIE: usize = 1;
 
+pub const SR_VS: usize = 0x00000600;
+pub const SR_VS_INITIAL: usize = 0x00000200;
+pub const SR_VS_DIRTY: usize = 0x00000600;
 pub const SR_FS: usize = 0x00006000;
 pub const SR_FS_OFF: usize = 0x00000000;
 pub const SR_FS_INITIAL: usize = 0x00002000;

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

@@ -67,23 +67,23 @@ impl ConfidentialHart {
     /// confidential hart has been assigned to a hardware hart for execution. A dummy confidential hart has the id of the hardware hart it
     /// represents. This is required to track information on which hardware hart a stolen confidential hart executes (for example to
     /// interrupt it with IPI).
-    pub fn dummy(hardware_hart_id: usize) -> Self {
+    pub fn dummy(hardware_hart_id: usize) -> Result<Self, Error> {
         // We set the lifecycle state of the dummy hart to `Started` because this state will be used as the state of the confidential hart
         // that has been assigned to physical hart. Specifically, after a confidential hart gets assigned to hardware hart, the confidential
         // VM will be left with a dummy hart. Any other code or request that will inspect the state of the (stolen) confidential hart, will
         // in reality look into the state of the dummy hart which will correctly say that the confidential hart is in the `Started` state.
-        Self {
+        Ok(Self {
             confidential_vm_id: None,
-            confidential_hart_state: HartArchitecturalState::empty(),
+            confidential_hart_state: HartArchitecturalState::empty()?,
             lifecycle_state: HartLifecycleState::Started,
             resumable_operation: None,
             id: hardware_hart_id,
-        }
+        })
     }
 
     /// Constructs a confidential hart with the state after a reset.
-    pub fn from_vm_hart_reset(id: usize, htimedelta: usize, shared_memory: &NaclSharedMemory) -> Self {
-        let mut confidential_hart_state = HartArchitecturalState::empty();
+    pub fn from_vm_hart_reset(id: usize, htimedelta: usize, shared_memory: &NaclSharedMemory) -> Result<Self, Error> {
+        let mut confidential_hart_state = HartArchitecturalState::empty()?;
         confidential_hart_state.csrs_mut().mstatus.save_value_in_main_memory(Self::INITIAL_MSTATUS);
         confidential_hart_state.csrs_mut().sstatus.save_value_in_main_memory(Self::INITIAL_SSTATUS);
         confidential_hart_state.csrs_mut().hstatus.save_value_in_main_memory(Self::INITIAL_HSTATUS);
@@ -116,22 +116,27 @@ impl ConfidentialHart {
             confidential_hart_state.csrs_mut().mstatus.enable_bits_on_saved_value(SR_FS_INITIAL);
         }
 
-        Self {
+        if HardwareSetup::is_extension_supported(HardwareExtension::VectorExtension) {
+            confidential_hart_state.csrs_mut().mstatus.enable_bits_on_saved_value(SR_VS_INITIAL);
+            confidential_hart_state.csrs_mut().vsstatus.enable_bits_on_saved_value(SR_VS_INITIAL);
+        }
+
+        Ok(Self {
             confidential_vm_id: None,
             confidential_hart_state,
             lifecycle_state: HartLifecycleState::Stopped,
             resumable_operation: None,
             id,
-        }
+        })
     }
 
     /// Constructs a confidential hart with the state of the non-confidential hart that made a call to promote the VM to confidential VM
     pub fn from_vm_hart(
         id: usize, program_counter: usize, fdt_address: &ConfidentialVmPhysicalAddress, htimedelta: usize, shared_memory: &NaclSharedMemory,
-    ) -> Self {
+    ) -> Result<Self, Error> {
         // We first create a confidential hart in the reset state and then fill this state with the runtime state of the hart that made a
         // call to promote to confidential VM. This state consists of GPRs and VS-level CSRs.
-        let mut confidential_hart = Self::from_vm_hart_reset(id, htimedelta, shared_memory);
+        let mut confidential_hart = Self::from_vm_hart_reset(id, htimedelta, shared_memory)?;
         let confidential_hart_state = &mut confidential_hart.confidential_hart_state;
         confidential_hart_state.set_gprs(shared_memory.gprs());
         confidential_hart_state.gprs_mut().write(GeneralPurposeRegister::a0, id);
@@ -147,7 +152,7 @@ impl ConfidentialHart {
         // Store the program counter of the VM, so that we can resume confidential VM at the point it became promoted.
         confidential_hart_state.csrs_mut().mepc.save_value_in_main_memory(program_counter);
         confidential_hart.lifecycle_state = HartLifecycleState::Started;
-        confidential_hart
+        Ok(confidential_hart)
     }
 
     /// Saves the state of the confidential hart in the main memory. This is part of the heavy context switch.
@@ -171,6 +176,14 @@ impl ConfidentialHart {
             self.csrs_mut().vsstatus.disable_bits_on_saved_value(SR_FS_DIRTY);
             self.csrs_mut().vsstatus.enable_bits_on_saved_value(SR_FS_CLEAN);
         }
+
+        if HardwareSetup::is_extension_supported(HardwareExtension::VectorExtension) {
+            // Enable the vector unit in mstatus to access vector
+            // registers and execute vector instructions.
+            self.csrs().mstatus.read_and_set_bits(SR_VS_DIRTY);
+
+            unsafe { self.confidential_hart_state.vrs_mut().save_in_main_memory() };
+        }
     }
 
     /// Restores the state of the confidential hart from the main memory. This is part of the heavy context switch.
@@ -186,6 +199,14 @@ impl ConfidentialHart {
             // below unsafe is ok because we checked that FPU hardware is available and we enabled it in mstatus.
             unsafe { self.confidential_hart_state.fprs_mut().restore_from_main_memory() };
         }
+
+        if HardwareSetup::is_extension_supported(HardwareExtension::VectorExtension) {
+            // Enable the vector unit in mstatus to access vector
+            // registers and execute vector instructions.
+            self.csrs().mstatus.read_and_set_bits(SR_VS_DIRTY);
+
+            unsafe { self.confidential_hart_state.vrs().restore_from_main_memory() };
+        }
     }
 
     pub fn measure(&self, digest: &mut MeasurementDigest) {

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

@@ -5,6 +5,7 @@ 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 crate::error::Error;
 use opensbi_sys::sbi_trap_regs;
 
 unsafe extern "C" {
@@ -43,14 +44,14 @@ pub struct HardwareHart {
 impl HardwareHart {
     /// Creates the instance of a state associated with the physical hart. Id is the unique physical hart identifier used to by hardware
     /// during inter process interrupts (IPIs), e.g., on RISC-V this will be a number reported by mhartid.
-    pub fn init(id: usize, stack: Page<UnAllocated>, hypervisor_memory_protector: HypervisorMemoryProtector) -> Self {
-        Self {
-            hypervisor_hart: HypervisorHart::new(hypervisor_memory_protector),
-            confidential_hart: ConfidentialHart::dummy(id),
+    pub fn init(id: usize, stack: Page<UnAllocated>, hypervisor_memory_protector: HypervisorMemoryProtector) -> Result<Self, Error> {
+        Ok(Self {
+            hypervisor_hart: HypervisorHart::new(hypervisor_memory_protector)?,
+            confidential_hart: ConfidentialHart::dummy(id)?,
             stack_address: stack.end_address(),
             stack: stack.zeroize(),
             previous_mscratch: 0,
-        }
+        })
     }
 
     // Safety: this function can be called only once during initialization. We cannot do it in constructor, because physical harts art not

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

@@ -2,7 +2,7 @@
 // SPDX-FileContributor: Wojciech Ozga <woz@zurich.ibm.com>, IBM Research - Zurich
 // SPDX-License-Identifier: Apache-2.0
 use crate::core::architecture::riscv::sbi::NaclSharedMemory;
-use crate::core::architecture::riscv::specification::SR_FS;
+use crate::core::architecture::riscv::specification::{SR_FS, SR_VS_DIRTY};
 use crate::core::architecture::specification::SR_FS_DIRTY;
 use crate::core::architecture::{
     ControlStatusRegisters, GeneralPurposeRegisters, HardwareExtension, HartArchitecturalState, SupervisorTimerExtension,
@@ -28,12 +28,12 @@ pub struct HypervisorHart {
 }
 
 impl HypervisorHart {
-    pub fn new(hypervisor_memory_protector: HypervisorMemoryProtector) -> Self {
-        Self {
-            hypervisor_hart_state: HartArchitecturalState::empty(),
+    pub fn new(hypervisor_memory_protector: HypervisorMemoryProtector) -> Result<Self, Error> {
+        Ok(Self {
+            hypervisor_hart_state: HartArchitecturalState::empty()?,
             shared_memory: NaclSharedMemory::uninitialized(),
             hypervisor_memory_protector,
-        }
+        })
     }
 
     /// Saves the state of the hypervisor hart in the main memory. This is part of the heavy context switch.
@@ -52,6 +52,13 @@ impl HypervisorHart {
                 }
             }
         }
+
+        if HardwareSetup::is_extension_supported(HardwareExtension::VectorExtension) {
+            // Enable the vector unit in mstatus to access vector
+            // registers and execute vector instructions.
+            self.csrs().mstatus.read_and_set_bits(SR_VS_DIRTY);
+            unsafe { self.hypervisor_hart_state.vrs_mut().save_in_main_memory() };
+        }
     }
 
     /// Restores the state of the hypervisor hart from the main memory. This is part of the heavy context switch.
@@ -65,6 +72,13 @@ impl HypervisorHart {
             // below unsafe is ok because we checked that FPU hardware is available and we enabled it in mstatus.
             unsafe { self.hypervisor_hart_state.fprs_mut().restore_from_main_memory() };
         }
+
+        if HardwareSetup::is_extension_supported(HardwareExtension::VectorExtension) {
+            // Enable the vector unit in mstatus to access vector
+            // registers and execute vector instructions.
+            self.csrs().mstatus.read_and_set_bits(SR_VS_DIRTY);
+            unsafe { self.hypervisor_hart_state.vrs().restore_from_main_memory() };
+        }
     }
 
     pub fn set_shared_memory(&mut self, base_address: NonConfidentialMemoryAddress) -> Result<(), Error> {