Abstract the virtual_exception_handler function.

framework/jinux-frame/src/arch/x86/cpu.rs

@@ -5,9 +5,11 @@ use core::fmt::Debug;
 
 use trapframe::{GeneralRegs, UserContext as RawUserContext};
 
+#[cfg(feature = "intel_tdx")]
+use crate::arch::tdx_guest::{virtual_exception_handler, TdxTrapFrame};
 use log::debug;
 #[cfg(feature = "intel_tdx")]
-use tdx_guest::{serial_println, tdcall, tdvmcall, TdxVirtualExceptionType};
+use tdx_guest::tdcall;
 use x86_64::registers::rflags::RFlags;
 
 use crate::trap::call_irq_callback_functions;
@@ -42,72 +44,52 @@ pub struct TrapInformation {
 }
 
 #[cfg(feature = "intel_tdx")]
-pub fn virtual_exception_handler(trapframe: &mut GeneralRegs, ve_info: &tdcall::TdgVeInfo) {
-    match ve_info.exit_reason.into() {
-        TdxVirtualExceptionType::Hlt => {
-            serial_println!("Ready to halt");
-            tdvmcall::hlt();
-        }
-        TdxVirtualExceptionType::Io => {
-            if !handle_ve_io(trapframe, ve_info) {
-                serial_println!("Handle tdx ioexit errors, ready to halt");
-                tdvmcall::hlt();
-            }
-        }
-        TdxVirtualExceptionType::MsrRead => {
-            let msr = tdvmcall::rdmsr(trapframe.rcx as u32).unwrap();
-            trapframe.rax = (msr as u32 & u32::MAX) as usize;
-            trapframe.rdx = ((msr >> 32) as u32 & u32::MAX) as usize;
-        }
-        TdxVirtualExceptionType::MsrWrite => {
-            let data = trapframe.rax as u64 | ((trapframe.rdx as u64) << 32);
-            tdvmcall::wrmsr(trapframe.rcx as u32, data).unwrap();
-        }
-        TdxVirtualExceptionType::CpuId => {
-            let cpuid_info = tdvmcall::cpuid(trapframe.rax as u32, trapframe.rcx as u32).unwrap();
-            let mask = 0xFFFF_FFFF_0000_0000_usize;
-            trapframe.rax = (trapframe.rax & mask) | cpuid_info.eax;
-            trapframe.rbx = (trapframe.rbx & mask) | cpuid_info.ebx;
-            trapframe.rcx = (trapframe.rcx & mask) | cpuid_info.ecx;
-            trapframe.rdx = (trapframe.rdx & mask) | cpuid_info.edx;
-        }
-        TdxVirtualExceptionType::Other => panic!("Unknown TDX vitrual exception type"),
-        _ => return,
-    }
-    trapframe.rip = trapframe.rip + ve_info.exit_instruction_length as usize;
-}
+struct VeGeneralRegs<'a>(&'a mut GeneralRegs);
 
 #[cfg(feature = "intel_tdx")]
-pub fn handle_ve_io(trapframe: &mut GeneralRegs, ve_info: &tdcall::TdgVeInfo) -> bool {
-    let size = match ve_info.exit_qualification & 0x3 {
-        0 => 1,
-        1 => 2,
-        3 => 4,
-        _ => panic!("Invalid size value"),
-    };
-    let direction = if (ve_info.exit_qualification >> 3) & 0x1 == 0 {
-        tdvmcall::Direction::Out
-    } else {
-        tdvmcall::Direction::In
-    };
-    let string = (ve_info.exit_qualification >> 4) & 0x1 == 1;
-    let repeat = (ve_info.exit_qualification >> 5) & 0x1 == 1;
-    let operand = if (ve_info.exit_qualification >> 6) & 0x1 == 0 {
-        tdvmcall::Operand::Dx
-    } else {
-        tdvmcall::Operand::Immediate
-    };
-    let port = (ve_info.exit_qualification >> 16) as u16;
-
-    match direction {
-        tdvmcall::Direction::In => {
-            trapframe.rax = tdvmcall::io_read(size, port).unwrap() as usize;
-        }
-        tdvmcall::Direction::Out => {
-            tdvmcall::io_write(size, port, trapframe.rax as u32).unwrap();
-        }
-    };
-    true
+impl TdxTrapFrame for VeGeneralRegs<'_> {
+    fn rax(&self) -> usize {
+        self.0.rax
+    }
+    fn set_rax(&mut self, rax: usize) {
+        self.0.rax = rax;
+    }
+    fn rbx(&self) -> usize {
+        self.0.rbx
+    }
+    fn set_rbx(&mut self, rbx: usize) {
+        self.0.rbx = rbx;
+    }
+    fn rcx(&self) -> usize {
+        self.0.rcx
+    }
+    fn set_rcx(&mut self, rcx: usize) {
+        self.0.rcx = rcx;
+    }
+    fn rdx(&self) -> usize {
+        self.0.rdx
+    }
+    fn set_rdx(&mut self, rdx: usize) {
+        self.0.rdx = rdx;
+    }
+    fn rsi(&self) -> usize {
+        self.0.rsi
+    }
+    fn set_rsi(&mut self, rsi: usize) {
+        self.0.rsi = rsi;
+    }
+    fn rdi(&self) -> usize {
+        self.0.rdi
+    }
+    fn set_rdi(&mut self, rdi: usize) {
+        self.0.rdi = rdi;
+    }
+    fn rip(&self) -> usize {
+        self.0.rip
+    }
+    fn set_rip(&mut self, rip: usize) {
+        self.0.rip = rip;
+    }
 }
 
 impl UserContext {
@@ -148,7 +130,8 @@ impl UserContextApiInternal for UserContext {
                     if *exception == VIRTUALIZATION_EXCEPTION {
                         let ve_info =
                             tdcall::get_veinfo().expect("#VE handler: fail to get VE info\n");
-                        virtual_exception_handler(self.general_regs_mut(), &ve_info);
+                        let mut ve_f = VeGeneralRegs(self.general_regs_mut());
+                        virtual_exception_handler(&mut ve_f, &ve_info);
                         continue;
                     }
                     if exception.typ == CpuExceptionType::FaultOrTrap

framework/jinux-frame/src/arch/x86/mod.rs

@@ -6,6 +6,8 @@ pub(crate) mod irq;
 mod kernel;
 pub(crate) mod mm;
 pub(crate) mod pci;
+#[cfg(feature = "intel_tdx")]
+pub(crate) mod tdx_guest;
 pub(crate) mod timer;
 
 use alloc::fmt;

framework/jinux-frame/src/arch/x86/tdx_guest.rs

@@ -0,0 +1,89 @@
+use tdx_guest::{
+    tdcall::TdgVeInfo,
+    tdvmcall::{cpuid, hlt, rdmsr, wrmsr},
+    {serial_println, tdcall, tdvmcall, TdxVirtualExceptionType},
+};
+
+pub trait TdxTrapFrame {
+    fn rax(&self) -> usize;
+    fn set_rax(&mut self, rax: usize);
+    fn rbx(&self) -> usize;
+    fn set_rbx(&mut self, rbx: usize);
+    fn rcx(&self) -> usize;
+    fn set_rcx(&mut self, rcx: usize);
+    fn rdx(&self) -> usize;
+    fn set_rdx(&mut self, rdx: usize);
+    fn rsi(&self) -> usize;
+    fn set_rsi(&mut self, rsi: usize);
+    fn rdi(&self) -> usize;
+    fn set_rdi(&mut self, rdi: usize);
+    fn rip(&self) -> usize;
+    fn set_rip(&mut self, rip: usize);
+}
+
+fn io_handler(trapframe: &mut dyn TdxTrapFrame, ve_info: &tdcall::TdgVeInfo) -> bool {
+    let size = match ve_info.exit_qualification & 0x3 {
+        0 => 1,
+        1 => 2,
+        3 => 4,
+        _ => panic!("Invalid size value"),
+    };
+    let direction = if (ve_info.exit_qualification >> 3) & 0x1 == 0 {
+        tdvmcall::Direction::Out
+    } else {
+        tdvmcall::Direction::In
+    };
+    let string = (ve_info.exit_qualification >> 4) & 0x1 == 1;
+    let repeat = (ve_info.exit_qualification >> 5) & 0x1 == 1;
+    let operand = if (ve_info.exit_qualification >> 6) & 0x1 == 0 {
+        tdvmcall::Operand::Dx
+    } else {
+        tdvmcall::Operand::Immediate
+    };
+    let port = (ve_info.exit_qualification >> 16) as u16;
+
+    match direction {
+        tdvmcall::Direction::In => {
+            trapframe.set_rax(tdvmcall::io_read(size, port).unwrap() as usize);
+        }
+        tdvmcall::Direction::Out => {
+            tdvmcall::io_write(size, port, trapframe.rax() as u32).unwrap();
+        }
+    };
+    true
+}
+
+pub fn virtual_exception_handler(trapframe: &mut impl TdxTrapFrame, ve_info: &TdgVeInfo) {
+    match ve_info.exit_reason.into() {
+        TdxVirtualExceptionType::Hlt => {
+            serial_println!("Ready to halt");
+            hlt();
+        }
+        TdxVirtualExceptionType::Io => {
+            if !io_handler(trapframe, ve_info) {
+                serial_println!("Handle tdx ioexit errors, ready to halt");
+                hlt();
+            }
+        }
+        TdxVirtualExceptionType::MsrRead => {
+            let msr = rdmsr(trapframe.rcx() as u32).unwrap();
+            trapframe.set_rax((msr as u32 & u32::MAX) as usize);
+            trapframe.set_rdx(((msr >> 32) as u32 & u32::MAX) as usize);
+        }
+        TdxVirtualExceptionType::MsrWrite => {
+            let data = trapframe.rax() as u64 | ((trapframe.rdx() as u64) << 32);
+            wrmsr(trapframe.rcx() as u32, data).unwrap();
+        }
+        TdxVirtualExceptionType::CpuId => {
+            let cpuid_info = cpuid(trapframe.rax() as u32, trapframe.rcx() as u32).unwrap();
+            let mask = 0xFFFF_FFFF_0000_0000_usize;
+            trapframe.set_rax((trapframe.rax() & mask) | cpuid_info.eax);
+            trapframe.set_rbx((trapframe.rbx() & mask) | cpuid_info.ebx);
+            trapframe.set_rcx((trapframe.rcx() & mask) | cpuid_info.ecx);
+            trapframe.set_rdx((trapframe.rdx() & mask) | cpuid_info.edx);
+        }
+        TdxVirtualExceptionType::Other => panic!("Unknown TDX vitrual exception type"),
+        _ => return,
+    }
+    trapframe.set_rip(trapframe.rip() + ve_info.exit_instruction_length as usize);
+}

framework/jinux-frame/src/trap/handler.rs

@@ -1,76 +1,58 @@
 use crate::{arch::irq::IRQ_LIST, cpu::CpuException};
 
 #[cfg(feature = "intel_tdx")]
-use tdx_guest::{serial_println, tdcall, tdvmcall, TdxVirtualExceptionType};
+use crate::arch::tdx_guest::{virtual_exception_handler, TdxTrapFrame};
+#[cfg(feature = "intel_tdx")]
+use tdx_guest::tdcall;
 use trapframe::TrapFrame;
 
 #[cfg(feature = "intel_tdx")]
-pub fn virtual_exception_handler(trapframe: &mut TrapFrame, ve_info: &tdcall::TdgVeInfo) {
-    match ve_info.exit_reason.into() {
-        TdxVirtualExceptionType::Hlt => {
-            serial_println!("Ready to halt");
-            tdvmcall::hlt();
-        }
-        TdxVirtualExceptionType::Io => {
-            if !handle_ve_io(trapframe, ve_info) {
-                serial_println!("Handle tdx ioexit errors, ready to halt");
-                tdvmcall::hlt();
-            }
-        }
-        TdxVirtualExceptionType::MsrRead => {
-            let msr = tdvmcall::rdmsr(trapframe.rcx as u32).unwrap();
-            trapframe.rax = (msr as u32 & u32::MAX) as usize;
-            trapframe.rdx = ((msr >> 32) as u32 & u32::MAX) as usize;
-        }
-        TdxVirtualExceptionType::MsrWrite => {
-            let data = trapframe.rax as u64 | ((trapframe.rdx as u64) << 32);
-            tdvmcall::wrmsr(trapframe.rcx as u32, data).unwrap();
-        }
-        TdxVirtualExceptionType::CpuId => {
-            let cpuid_info = tdvmcall::cpuid(trapframe.rax as u32, trapframe.rcx as u32).unwrap();
-            let mask = 0xFFFF_FFFF_0000_0000_usize;
-            trapframe.rax = (trapframe.rax & mask) | cpuid_info.eax;
-            trapframe.rbx = (trapframe.rbx & mask) | cpuid_info.ebx;
-            trapframe.rcx = (trapframe.rcx & mask) | cpuid_info.ecx;
-            trapframe.rdx = (trapframe.rdx & mask) | cpuid_info.edx;
-        }
-        TdxVirtualExceptionType::Other => panic!("Unknown TDX vitrual exception type"),
-        _ => return,
-    }
-    trapframe.rip = trapframe.rip + ve_info.exit_instruction_length as usize;
-}
+struct VeTrapFrame<'a>(&'a mut TrapFrame);
 
 #[cfg(feature = "intel_tdx")]
-pub fn handle_ve_io(trapframe: &mut TrapFrame, ve_info: &tdcall::TdgVeInfo) -> bool {
-    let size = match ve_info.exit_qualification & 0x3 {
-        0 => 1,
-        1 => 2,
-        3 => 4,
-        _ => panic!("Invalid size value"),
-    };
-    let direction = if (ve_info.exit_qualification >> 3) & 0x1 == 0 {
-        tdvmcall::Direction::Out
-    } else {
-        tdvmcall::Direction::In
-    };
-    let string = (ve_info.exit_qualification >> 4) & 0x1 == 1;
-    let repeat = (ve_info.exit_qualification >> 5) & 0x1 == 1;
-    let operand = if (ve_info.exit_qualification >> 6) & 0x1 == 0 {
-        tdvmcall::Operand::Dx
-    } else {
-        tdvmcall::Operand::Immediate
-    };
-    let port = (ve_info.exit_qualification >> 16) as u16;
-
-    match direction {
-        tdvmcall::Direction::In => {
-            trapframe.rax = tdvmcall::io_read(size, port).unwrap() as usize;
-        }
-        tdvmcall::Direction::Out => {
-            tdvmcall::io_write(size, port, trapframe.rax as u32).unwrap();
-        }
-    };
-    true
+impl TdxTrapFrame for VeTrapFrame<'_> {
+    fn rax(&self) -> usize {
+        self.0.rax
+    }
+    fn set_rax(&mut self, rax: usize) {
+        self.0.rax = rax;
+    }
+    fn rbx(&self) -> usize {
+        self.0.rbx
+    }
+    fn set_rbx(&mut self, rbx: usize) {
+        self.0.rbx = rbx;
+    }
+    fn rcx(&self) -> usize {
+        self.0.rcx
+    }
+    fn set_rcx(&mut self, rcx: usize) {
+        self.0.rcx = rcx;
+    }
+    fn rdx(&self) -> usize {
+        self.0.rdx
+    }
+    fn set_rdx(&mut self, rdx: usize) {
+        self.0.rdx = rdx;
+    }
+    fn rsi(&self) -> usize {
+        self.0.rsi
+    }
+    fn set_rsi(&mut self, rsi: usize) {
+        self.0.rsi = rsi;
+    }
+    fn rdi(&self) -> usize {
+        self.0.rdi
+    }
+    fn set_rdi(&mut self, rdi: usize) {
+        self.0.rdi = rdi;
+    }
+    fn rip(&self) -> usize {
+        self.0.rip
+    }
+    fn set_rip(&mut self, rip: usize) {
+        self.0.rip = rip;
+    }
 }
 
 /// Only from kernel
@@ -80,7 +62,8 @@ extern "sysv64" fn trap_handler(f: &mut TrapFrame) {
         #[cfg(feature = "intel_tdx")]
         if f.trap_num as u16 == 20 {
             let ve_info = tdcall::get_veinfo().expect("#VE handler: fail to get VE info\n");
-            virtual_exception_handler(f, &ve_info);
+            let mut ve_f = VeTrapFrame(f);
+            virtual_exception_handler(&mut ve_f, &ve_info);
         }
         #[cfg(not(feature = "intel_tdx"))]
         panic!("cannot handle kernel cpu fault now, information:{:#x?}", f);