feat: Add conversion of trapframe to kprobe's pt_regs.

Add custom callback for break exception.
Add debug handler for x86_64.
Add kprobe support to aarch64 arch and extend break handler parameter.
Add debug exception for x86_64 and fix the break exception definition error in aarch64.

Signed-off-by: Godones <chenlinfeng25@outlook.com>

Author: Godones

Cargo.lock

@@ -27,6 +27,7 @@ fp-simd = []
 tls = []
 uspace = []
 arm-el2 = []
+kprobe = ["dep:kprobe"]
 
 [dependencies]
 axbacktrace = "0.1"
@@ -36,6 +37,7 @@ cfg-if = "1.0"
 memory_addr = "0.4"
 page_table_entry = "0.5"
 static_assertions = "1.1.0"
+kprobe = { git = "https://github.com/Starry-OS/kprobe", optional = true }
 
 [target.'cfg(target_arch = "x86_64")'.dependencies]
 lazyinit = "0.2"

Cargo.toml

@@ -31,6 +31,34 @@ impl fmt::Debug for TrapFrame {
     }
 }
 
+#[cfg(feature = "kprobe")]
+impl From<&TrapFrame> for kprobe::PtRegs {
+    fn from(tf: &TrapFrame) -> Self {
+        use core::u64;
+
+        use kprobe::PtRegs;
+        PtRegs {
+            regs: tf.x,
+            sp: u64::MAX,
+            pc: tf.elr,
+            pstate: tf.spsr,
+            orig_x0: u64::MAX,
+            syscallno: -1,
+            unused2: 0,
+        }
+    }
+}
+
+#[cfg(feature = "kprobe")]
+impl TrapFrame {
+    /// Update the TrapFrame from kprobe::PtRegs
+    pub fn update_from_ptregs(&mut self, ptregs: kprobe::PtRegs) {
+        self.x = ptregs.regs;
+        self.spsr = ptregs.pstate;
+        self.elr = ptregs.pc;
+    }
+}
+
 impl TrapFrame {
     /// Gets the 0th syscall argument.
     pub const fn arg0(&self) -> usize {

src/aarch64/context.rs

@@ -41,6 +41,14 @@ pub(super) fn is_valid_page_fault(iss: u64) -> bool {
     matches!(iss & 0b111100, 0b0100 | 0b1100) // IFSC or DFSC bits
 }
 
+fn handle_breakpoint(tf: &mut TrapFrame, iss: u64) {
+    debug!("BRK #{:#x} @ {:#x} ", iss, tf.elr);
+    if core::hint::likely(handle_trap!(BREAK_HANDLER, tf, iss)) {
+        return;
+    }
+    tf.elr += 4;
+}
+
 fn handle_page_fault(tf: &mut TrapFrame, access_flags: PageFaultFlags) {
     let vaddr = va!(FAR_EL1.get() as usize);
     if handle_trap!(PAGE_FAULT, vaddr, access_flags) {
@@ -99,10 +107,7 @@ fn aarch64_trap_handler(tf: &mut TrapFrame, kind: TrapKind, source: TrapSource)
                         },
                     );
                 }
-                Some(ESR_EL1::EC::Value::Brk64) => {
-                    debug!("BRK #{:#x} @ {:#x} ", iss, tf.elr);
-                    tf.elr += 4;
-                }
+                Some(ESR_EL1::EC::Value::Brk64) => handle_breakpoint(tf, iss),
                 e => {
                     let vaddr = va!(FAR_EL1.get() as usize);
                     panic!(

src/aarch64/trap.rs

@@ -150,7 +150,7 @@ impl ExceptionInfo {
     /// Returns a generalized kind of this exception.
     pub fn kind(&self) -> ExceptionKind {
         match self.esr.read_as_enum(ESR_EL1::EC) {
-            Some(ESR_EL1::EC::Value::BreakpointLowerEL) => ExceptionKind::Breakpoint,
+            Some(ESR_EL1::EC::Value::Brk64) => ExceptionKind::Breakpoint,
             Some(ESR_EL1::EC::Value::IllegalExecutionState) => ExceptionKind::IllegalInstruction,
             Some(ESR_EL1::EC::Value::PCAlignmentFault)
             | Some(ESR_EL1::EC::Value::SPAlignmentFault) => ExceptionKind::Misaligned,

src/aarch64/uspace.rs

@@ -1,3 +1,4 @@
+#![feature(likely_unlikely)]
 #![cfg_attr(not(test), no_std)]
 #![cfg_attr(docsrs, feature(doc_cfg))]
 #![feature(cold_path)]

src/lib.rs

@@ -259,8 +259,8 @@ impl TaskContext {
 
     /// Switches to another task.
     ///
-    /// It first saves the current task's context from CPU to this place, and then
-    /// restores the next task's context from `next_ctx` to CPU.
+    /// It first saves the current task's context from CPU to this place, and
+    /// then restores the next task's context from `next_ctx` to CPU.
     pub fn switch_to(&mut self, next_ctx: &Self) {
         #[cfg(feature = "tls")]
         {
@@ -353,3 +353,44 @@ unsafe extern "C" fn context_switch(_current_task: &mut TaskContext, _next_task:
         ret",
     )
 }
+
+#[cfg(feature = "kprobe")]
+impl From<&TrapFrame> for kprobe::PtRegs {
+    fn from(tf: &TrapFrame) -> Self {
+        let regs = [0; 32];
+        unsafe {
+            core::ptr::copy_nonoverlapping(
+                &tf.regs as *const GeneralRegisters as *const usize,
+                regs.as_ptr() as *mut usize,
+                32,
+            );
+        }
+        kprobe::PtRegs {
+            regs,
+            orig_a0: 0,
+            csr_era: tf.era,
+            csr_badvaddr: 0,
+            csr_crmd: 0,
+            csr_prmd: tf.prmd,
+            csr_euen: 0,
+            csr_ecfg: 0,
+            csr_estat: 0,
+        }
+    }
+}
+
+#[cfg(feature = "kprobe")]
+impl TrapFrame {
+    /// Update the TrapFrame from kprobe::PtRegs
+    pub fn update_from_ptregs(&mut self, ptregs: kprobe::PtRegs) {
+        unsafe {
+            core::ptr::copy_nonoverlapping(
+                ptregs.regs.as_ptr() as *const usize,
+                &mut self.regs as *mut GeneralRegisters as *mut usize,
+                32,
+            );
+        }
+        self.era = ptregs.csr_era;
+        self.prmd = ptregs.csr_prmd;
+    }
+}

src/loongarch64/context.rs

@@ -12,9 +12,12 @@ core::arch::global_asm!(
     trapframe_size = const (core::mem::size_of::<TrapFrame>()),
 );
 
-fn handle_breakpoint(era: &mut usize) {
-    debug!("Exception(Breakpoint) @ {era:#x} ");
-    *era += 4;
+fn handle_breakpoint(tf: &mut TrapFrame) {
+    debug!("Exception(Breakpoint) @ {:#x} ", tf.era);
+    if core::hint::likely(handle_trap!(BREAK_HANDLER, tf, 0)) {
+        return;
+    }
+    tf.era += 4;
 }
 
 fn handle_page_fault(tf: &mut TrapFrame, access_flags: PageFaultFlags) {
@@ -54,7 +57,7 @@ fn loongarch64_trap_handler(tf: &mut TrapFrame) {
         | Trap::Exception(Exception::PageNonExecutableFault) => {
             handle_page_fault(tf, PageFaultFlags::EXECUTE);
         }
-        Trap::Exception(Exception::Breakpoint) => handle_breakpoint(&mut tf.era),
+        Trap::Exception(Exception::Breakpoint) => handle_breakpoint(tf),
         Trap::Exception(Exception::AddressNotAligned) => unsafe {
             tf.emulate_unaligned().unwrap();
         },

src/loongarch64/trap.rs

@@ -1,4 +1,5 @@
 use core::arch::naked_asm;
+
 use memory_addr::VirtAddr;
 use riscv::register::sstatus::{self, FS};
 
@@ -85,7 +86,8 @@ impl FpState {
 
     /// Handles floating-point state context switching
     ///
-    /// Saves the current task's FP state (if needed) and restores the next task's FP state
+    /// Saves the current task's FP state (if needed) and restores the next
+    /// task's FP state
     pub fn switch_to(&mut self, next_fp_state: &FpState) {
         // get the real FP state of the current task
         let current_fs = sstatus::read().fs();
@@ -98,9 +100,12 @@ impl FpState {
         }
         // restore the next task's FP state
         match next_fp_state.fs {
-            FS::Clean => next_fp_state.restore(), // the next task's FP state is clean, we should restore it
-            FS::Initial => FpState::clear(),      // restore the FP state as constant values(all 0)
-            FS::Off => {}                         // do nothing
+            // the next task's FP state is clean, we should restore it
+            FS::Clean => next_fp_state.restore(),
+            // restore the FP state as constant values(all 0)
+            FS::Initial => FpState::clear(),
+            // do nothing
+            FS::Off => {}
             FS::Dirty => unreachable!("FP state of the next task should not be dirty"),
         }
         unsafe { sstatus::set_fs(next_fp_state.fs) }; // set the FP state to the next task's FP state
@@ -326,8 +331,8 @@ impl TaskContext {
 
     /// Switches to another task.
     ///
-    /// It first saves the current task's context from CPU to this place, and then
-    /// restores the next task's context from `next_ctx` to CPU.
+    /// It first saves the current task's context from CPU to this place, and
+    /// then restores the next task's context from `next_ctx` to CPU.
     pub fn switch_to(&mut self, next_ctx: &Self) {
         #[cfg(feature = "tls")]
         {
@@ -425,3 +430,87 @@ unsafe extern "C" fn context_switch(_current_task: &mut TaskContext, _next_task:
         ret",
     )
 }
+
+#[cfg(feature = "kprobe")]
+impl From<&TrapFrame> for kprobe::PtRegs {
+    fn from(tf: &TrapFrame) -> Self {
+        kprobe::PtRegs {
+            epc: tf.sepc,
+            ra: tf.regs.ra,
+            sp: tf.regs.sp,
+            gp: tf.regs.gp,
+            tp: tf.regs.tp,
+            t0: tf.regs.t0,
+            t1: tf.regs.t1,
+            t2: tf.regs.t2,
+            s0: tf.regs.s0,
+            s1: tf.regs.s1,
+            a0: tf.regs.a0,
+            a1: tf.regs.a1,
+            a2: tf.regs.a2,
+            a3: tf.regs.a3,
+            a4: tf.regs.a4,
+            a5: tf.regs.a5,
+            a6: tf.regs.a6,
+            a7: tf.regs.a7,
+            s2: tf.regs.s2,
+            s3: tf.regs.s3,
+            s4: tf.regs.s4,
+            s5: tf.regs.s5,
+            s6: tf.regs.s6,
+            s7: tf.regs.s7,
+            s8: tf.regs.s8,
+            s9: tf.regs.s9,
+            s10: tf.regs.s10,
+            s11: tf.regs.s11,
+            t3: tf.regs.t3,
+            t4: tf.regs.t4,
+            t5: tf.regs.t5,
+            t6: tf.regs.t6,
+            status: tf.sstatus.bits(),
+            // todo : other fields
+            badaddr: 0,
+            cause: 0,
+            orig_a0: 0,
+        }
+    }
+}
+
+impl TrapFrame {
+    /// Update the TrapFrame from kprobe::PtRegs
+    pub fn update_from_ptregs(&mut self, ptregs: kprobe::PtRegs) {
+        self.sepc = ptregs.epc;
+        self.regs.ra = ptregs.ra;
+        self.regs.sp = ptregs.sp;
+        self.regs.gp = ptregs.gp;
+        self.regs.tp = ptregs.tp;
+        self.regs.t0 = ptregs.t0;
+        self.regs.t1 = ptregs.t1;
+        self.regs.t2 = ptregs.t2;
+        self.regs.s0 = ptregs.s0;
+        self.regs.s1 = ptregs.s1;
+        self.regs.a0 = ptregs.a0;
+        self.regs.a1 = ptregs.a1;
+        self.regs.a2 = ptregs.a2;
+        self.regs.a3 = ptregs.a3;
+        self.regs.a4 = ptregs.a4;
+        self.regs.a5 = ptregs.a5;
+        self.regs.a6 = ptregs.a6;
+        self.regs.a7 = ptregs.a7;
+        self.regs.s2 = ptregs.s2;
+        self.regs.s3 = ptregs.s3;
+        self.regs.s4 = ptregs.s4;
+        self.regs.s5 = ptregs.s5;
+        self.regs.s6 = ptregs.s6;
+        self.regs.s7 = ptregs.s7;
+        self.regs.s8 = ptregs.s8;
+        self.regs.s9 = ptregs.s9;
+        self.regs.s10 = ptregs.s10;
+        self.regs.s11 = ptregs.s11;
+        self.regs.t3 = ptregs.t3;
+        self.regs.t4 = ptregs.t4;
+        self.regs.t5 = ptregs.t5;
+        self.regs.t6 = ptregs.t6;
+        self.sstatus = sstatus::Sstatus::from_bits(ptregs.status);
+    }
+}