Trap Handler Implementation

Relevant source files

• src/regs.rs
• src/trap.S
• src/trap.rs

This document covers the low-level assembly language trap handlers that manage guest virtual machine entry and exit. These handlers perform the critical register context switching between hypervisor and guest execution contexts, enabling secure isolation and efficient virtualization on RISC-V systems with H-extension support.

For higher-level trap processing and VM exit handling, see VM Exit Processing. For register data structure definitions, see Register Management.

Overview

The trap handler implementation consists of two primary assembly language routines that manage the transition between hypervisor and guest execution contexts:

• Guest Entry (_run_guest) - Saves hypervisor state, loads guest state, and transfers control to the guest VM
• Guest Exit (_guest_exit) - Captures guest state upon VM exit and restores hypervisor execution context

These handlers work in conjunction with carefully designed data structures in src/regs.rs to ensure complete CPU state preservation during virtualization transitions.

Architecture and Data Flow

The trap handler operates through a precisely choreographed sequence of register and CSR manipulations:

flowchart TD
subgraph subGraph0["VmCpuRegisters Data Structure"]
    HypRegs["hyp_regs: HypervisorCpuState"]
    GuestRegs["guest_regs: GuestCpuState"]
    VSCSRs["vs_csrs: GuestVsCsrs"]
    TrapCSRs["trap_csrs: VmCpuTrapState"]
end
HypCall["Hypervisor calls _run_guest()"]
SaveHyp["Save Hypervisor State"]
LoadGuest["Load Guest State"]
SetTrap["Set trap vector to _guest_exit"]
SRET["Execute sret to enter guest"]
GuestExec["Guest Execution"]
TrapOccurs["Guest Trap/Exit Occurs"]
SaveGuest["Save Guest State"]
RestoreHyp["Restore Hypervisor State"]
Return["Return to Hypervisor"]

GuestExec --> TrapOccurs
HypCall --> SaveHyp
LoadGuest --> GuestRegs
LoadGuest --> SetTrap
RestoreHyp --> HypRegs
RestoreHyp --> Return
SRET --> GuestExec
SaveGuest --> GuestRegs
SaveGuest --> RestoreHyp
SaveHyp --> HypRegs
SaveHyp --> LoadGuest
SetTrap --> SRET
TrapOccurs --> SaveGuest

Sources: src/trap.S(L1 - L183)  src/regs.rs(L116 - L196) 

Register Context Data Structures

The trap handlers operate on the VmCpuRegisters structure, which contains separate state containers for hypervisor and guest contexts:

The assembly code accesses these structures through computed offsets provided by the Rust compilation system:

flowchart TD
subgraph subGraph2["Offset Calculation Functions"]
    VmCpuRegs["VmCpuRegisters"]
    HypState["hyp_regs"]
    GuestState["guest_regs"]
    HypGPRs["HypervisorCpuState.gprs"]
    GuestGPRs["GuestCpuState.gprs"]
    HypGPROffset["hyp_gpr_offset()"]
    GuestGPROffset["guest_gpr_offset()"]
    HypCSROffset["hyp_csr_offset!()"]
    GuestCSROffset["guest_csr_offset!()"]
end
subgraph subGraph1["Assembly Constants"]
    HypRA["hyp_ra offset"]
    HypSP["hyp_sp offset"]
    GuestRA["guest_ra offset"]
    GuestSP["guest_sp offset"]
    HypSSTATUS["hyp_sstatus offset"]
    GuestSSTATUS["guest_sstatus offset"]
end
subgraph subGraph0["Rust Data Structure"]
    VmCpuRegs["VmCpuRegisters"]
    HypState["hyp_regs"]
    GuestState["guest_regs"]
    HypGPRs["HypervisorCpuState.gprs"]
    GuestGPRs["GuestCpuState.gprs"]
    HypGPROffset["hyp_gpr_offset()"]
    GuestGPROffset["guest_gpr_offset()"]
end

GuestCSROffset --> GuestSSTATUS
GuestGPROffset --> GuestRA
GuestGPROffset --> GuestSP
GuestState --> GuestGPRs
HypCSROffset --> HypSSTATUS
HypGPROffset --> HypRA
HypGPROffset --> HypSP
HypState --> HypGPRs
VmCpuRegs --> GuestState
VmCpuRegs --> HypState

Sources: src/trap.rs(L7 - L33)  src/regs.rs(L1 - L114)  src/regs.rs(L116 - L138) 

Guest Entry Sequence (_run_guest)

The _run_guest function implements the critical transition from hypervisor to guest execution:

Hypervisor State Preservation

The handler first saves the hypervisor's execution context, excluding temporary registers and the a0 register which contains the VmCpuRegisters pointer:

• General-purpose registers (ra, gp, tp, s0-s11, a1-a7, sp)
• Critical CSRs (sstatus, hstatus, scounteren, stvec, sscratch)

Guest State Loading

The handler then configures the CPU for guest execution:

• Swaps hypervisor CSRs with guest CSRs using csrrw instructions
• Loads guest program counter into sepc
• Sets trap vector (stvec) to _guest_exit for handling guest traps
• Stores VmCpuRegisters pointer in sscratch for quick access during exit
• Restores all guest general-purpose registers

Context Switch Completion

The sequence concludes with an sret instruction that:

• Switches to guest privilege level
• Jumps to the guest program counter (sepc)
• Enables guest execution with appropriate privilege and virtualization settings

Sources: src/trap.S(L3 - L91)  src/trap.rs(L35 - L102) 

Guest Exit Sequence (_guest_exit)

When the guest encounters a trap, interrupt, or exception, control transfers to _guest_exit:

Guest State Capture

The exit handler immediately preserves the guest's execution state:

• Retrieves VmCpuRegisters pointer from sscratch
• Saves all guest general-purpose registers to the guest state structure
• Captures guest's a0 register value from sscratch
• Records guest program counter from sepc

Hypervisor State Restoration

The handler restores the hypervisor execution environment:

• Swaps guest CSRs back to hypervisor CSRs using csrrw
• Restores hypervisor trap vector (stvec)
• Restores hypervisor scratch register (sscratch)
• Reloads all hypervisor general-purpose registers

Return to Hypervisor

The handler concludes with a ret instruction, returning control to the hypervisor code that called _run_guest.

sequenceDiagram
    participant Hypervisor as Hypervisor
    participant Guest as Guest
    participant _run_guest as _run_guest
    participant _guest_exit as _guest_exit

    Hypervisor ->> _run_guest: Call with VmCpuRegisters*
    _run_guest ->> _run_guest: Save hypervisor GPRs
    _run_guest ->> _run_guest: Save hypervisor CSRs
    _run_guest ->> _run_guest: Load guest CSRs
    _run_guest ->> _run_guest: Set stvec to _guest_exit
    _run_guest ->> _run_guest: Load guest GPRs
    _run_guest ->> Guest: sret (enter guest)
    Guest ->> Guest: Execute guest code
    Guest ->> _guest_exit: Trap/Exception occurs
    _guest_exit ->> _guest_exit: Save guest GPRs
    _guest_exit ->> _guest_exit: Save guest CSRs
    _guest_exit ->> _guest_exit: Restore hypervisor CSRs
    _guest_exit ->> _guest_exit: Restore hypervisor GPRs
    _guest_exit ->> Hypervisor: ret (return to hypervisor)

Sources: src/trap.S(L92 - L183) 

Assembly-Rust Interface

The integration between assembly and Rust code relies on computed memory offsets and the global_asm! macro:

Offset Calculation

The Rust code provides compile-time offset calculations for accessing structure fields:

• hyp_gpr_offset() computes offsets into hypervisor register arrays
• guest_gpr_offset() computes offsets into guest register arrays
• hyp_csr_offset!() and guest_csr_offset!() macros compute CSR field offsets

Assembly Integration

The core::arch::global_asm! macro embeds the assembly code with substituted constants:

core::arch::global_asm!(
    include_str!("trap.S"),
    hyp_ra = const hyp_gpr_offset(GprIndex::RA),
    guest_ra = const guest_gpr_offset(GprIndex::RA),
    hyp_sstatus = const hyp_csr_offset!(sstatus),
    // ... additional offset constants
);

This approach ensures type safety and automatic offset calculation while maintaining the performance benefits of hand-optimized assembly code.

Sources: src/trap.rs(L1 - L102)  src/regs.rs(L5 - L86) 

Special Register Handling

sscratch Register Usage

The sscratch CSR serves a dual purpose in the trap handling mechanism:

• During guest execution: Contains the VmCpuRegisters pointer for fast access
• During hypervisor execution: Contains the hypervisor's original sscratch value

Register A0 Special Case

The a0 register receives special handling because it carries the VmCpuRegisters pointer:

• Not saved during initial hypervisor state preservation
• Swapped with guest a0 via sscratch during guest entry
• Guest's a0 value recovered from sscratch during exit

Zero Register Enforcement

The GeneralPurposeRegisters::set_reg() method enforces RISC-V architectural requirements by preventing writes to the zero register (x0), ensuring compliance with the RISC-V specification.

Sources: src/trap.S(L53 - L55)  src/trap.S(L94 - L95)  src/trap.S(L129 - L131)  src/regs.rs(L96 - L102)