Refactor kprobe handling to a unified probe system with fprobe.

Author: congwang-mk

README.md

@@ -93,8 +93,8 @@ KernelScript supports all major eBPF program types with typed contexts:
     return TC_ACT_OK
 }
 
-// Kprobe for kernel function tracing
-@kprobe fn trace_syscall(ctx: *pt_regs) -> i32 {
+// Probe for kernel function tracing
+@probe fn trace_syscall(ctx: *pt_regs) -> i32 {
     // Trace system call entry
     return 0
 }
@@ -247,8 +247,8 @@ kernelscript init xdp my_packet_filter
 # Create TC project  
 kernelscript init tc my_traffic_shaper
 
-# Create kprobe project
-kernelscript init kprobe/sys_read my_tracer
+# Create probe project
+kernelscript init probe/sys_read my_tracer
 
 # Create project with custom BTF path
 kernelscript init --btf-vmlinux-path /custom/path/vmlinux xdp my_project
@@ -268,7 +268,7 @@ my_project/
 **Available program types:**
 - `xdp` - XDP programs for packet processing
 - `tc` - Traffic control programs  
-- `kprobe` - Kernel function tracing
+- `probe` - Kernel function probing
 - `tracepoint` - Kernel tracepoint programs
 
 **Available struct_ops:**

SPEC.md

@@ -174,7 +174,7 @@ ebpf_program = attribute_list "fn" identifier "(" parameter_list ")" "->" return
 
 attribute_list = attribute { attribute }
 attribute = "@" attribute_name [ "(" attribute_args ")" ]
-attribute_name = "xdp" | "tc" | "kprobe" | "tracepoint" |
+attribute_name = "xdp" | "tc" | "probe" | "tracepoint" |
                  "struct_ops" | "kfunc" | "private" | "helper" | "test"
 attribute_args = string_literal | identifier
 
@@ -185,44 +185,64 @@ return_type = type_annotation
 
 **Note:** eBPF programs are now simple attributed functions. All configuration is done through global named config blocks.
 
-#### 3.1.1 Advanced Kprobe Functions with BTF Signature Extraction
+#### 3.1.1 Advanced Probe Functions with BTF Signature Extraction and Intelligent Probe Type Selection
 
-KernelScript automatically extracts kernel function signatures from BTF (BPF Type Format) for kprobe functions, eliminating the need for `KprobeContext` and providing type-safe access to function parameters.
+KernelScript automatically extracts kernel function signatures from BTF (BPF Type Format) for probe functions and intelligently chooses between fprobe (function entrance) and kprobe (arbitrary address) based on the target specification.
 
 ```kernelscript
-@kprobe("sys_read")
-fn new_style(fd: u32, buf: *u8, count: usize) -> i32 {
+// Function entrance probe (uses fprobe)
+@probe("sys_read")
+fn function_entrance(fd: u32, buf: *u8, count: usize) -> i32 {
     // Direct access to function parameters with correct types
     // Compiler automatically extracts signature from BTF:
     // long sys_read(unsigned int fd, char __user *buf, size_t count)
+    // Uses fprobe for better performance at function entrance
     
     print("Reading %d bytes from fd %d", count, fd)
     return 0
 }
+
+// Arbitrary address probe (uses kprobe)  
+@probe("vfs_read+109")
+fn arbitrary_address() -> i32 {
+    // Probes specific instruction offset within vfs_read
+    // Uses kprobe for arbitrary address probing
+    // No direct parameters available at arbitrary addresses
+    
+    print("Probing vfs_read at offset +109")
+    return 0
+}
 ```
 
 **Key Benefits:**
-- **Type Safety**: Parameters have correct types extracted from kernel BTF information
-- **No Magic Numbers**: Direct parameter access instead of `ctx.arg_*(index)`
-- **Self-Documenting**: Function signature matches the actual kernel function
+- **Intelligent Probe Selection**: Automatically chooses fprobe for function entrance (better performance) or kprobe for arbitrary addresses
+- **Type Safety**: Function entrance probes have correct types extracted from kernel BTF information
+- **No Magic Numbers**: Direct parameter access for function entrance probes
+- **Self-Documenting**: Function signature matches the actual kernel function for entrance probes
 - **Compile-Time Validation**: Invalid parameter access caught at compile time
 
-**BTF Signature Mapping:**
+**Probe Type Selection:**
+- `@probe("function_name")` → Uses **fprobe** for function entrance with direct parameter access
+- `@probe("function_name+offset")` → Uses **kprobe** for arbitrary address probing
+
+**BTF Signature Mapping for Function Entrance:**
 ```kernelscript
 // Kernel function: long sys_openat(int dfd, const char __user *filename, int flags, umode_t mode)
-@kprobe("sys_openat")
+@probe("sys_openat")
 fn trace_openat(dfd: i32, filename: *u8, flags: i32, mode: u16) -> i32 {
-    // Parameters automatically mapped to PT_REGS_PARM1, PT_REGS_PARM2, etc.
+    // Direct parameter access with fprobe (no PT_REGS needed)
     print("Opening file with flags %d", flags)
     return 0
 }
 
-// Kernel function: long sys_write(unsigned int fd, const char __user *buf, size_t count)
-@kprobe("sys_write")
-fn trace_write(fd: u32, buf: *u8, count: usize) -> i32 {
-    // Type-safe parameter access
-    if (count > 1024) {
-        print("Large write detected: %d bytes to fd %d", count, fd)
+// For arbitrary address probing:
+@probe("sys_write+50")  
+fn trace_write_offset() -> i32 {
+    // Uses kprobe for arbitrary offset - no direct parameters available
+    print("Probing sys_write at offset +50")
+    return 0
+}
+
     }
     return 0
 }
@@ -1940,7 +1960,7 @@ struct PersonInfo {
 }
 
 // Kernel space usage - kprobe with BTF-extracted function signature
-@kprobe("sys_open")
+@probe("sys_open")
 fn user_monitor(dfd: i32, filename: *u8, flags: i32, mode: u16) -> i32 {
     var process_name: ProcessName = get_current_process_name()
     var file_path: FilePath = get_file_path_from_filename(filename)
@@ -2381,7 +2401,7 @@ fn security_analyzer(ctx: LsmContext) -> i32 {
 pin var global_counters : array<u32, GlobalCounter>(256)
 pin var event_stream : hash<u32, Event>(1024)
 
-@kprobe("sys_read")
+@probe("sys_read")
 fn producer(fd: u32, buf: *u8, count: usize) -> i32 {
     var pid = bpf_get_current_pid_tgid() as u32
     
@@ -2401,7 +2421,7 @@ fn producer(fd: u32, buf: *u8, count: usize) -> i32 {
     return 0
 }
 
-@kprobe("sys_write")
+@probe("sys_write")
 fn consumer(fd: u32, buf: *u8, count: usize) -> i32 {
     var pid = bpf_get_current_pid_tgid() as u32
     
@@ -4335,7 +4355,7 @@ fn measure_write_time() -> u64 {
     return bpf_ktime_get_ns()
 }
 
-@kprobe("sys_read")
+@probe("sys_read")
 fn perf_monitor(fd: u32, buf: *u8, count: usize) -> i32 {
     var pid = bpf_get_current_pid_tgid() as u32
     var call_info = CallInfo {
@@ -4371,7 +4391,7 @@ fn perf_monitor_return(ret_value: isize) -> i32 {
     return 0
 }
 
-@kprobe("sys_write")
+@probe("sys_write")
 fn write_monitor(fd: u32, buf: *u8, count: usize) -> i32 {
     var pid = bpf_get_current_pid_tgid() as u32
     var duration = measure_write_time()  // No context needed

examples/map_operations_demo.ks

@@ -205,7 +205,7 @@ fn event_logger(ctx: *trace_event_raw_sys_enter) -> i32 {
 }
 
 // Program 4: Sequential access pattern demonstration
-@kprobe("vfs_read")
+@probe("vfs_read")
 fn data_processor(file: *file, buf: *u8, count: usize, pos: *i64) -> i32 {
     // Sequential access pattern - will be detected and optimized
     for (i in 0..32) {

examples/kprobe_do_exit.ks renamed to examples/probe_do_exit.ks

@@ -1,21 +1,21 @@
 // Kprobe Example: Monitor process exit events
 // 
-// This example demonstrates how to use kprobe to intercept and monitor
+// This example demonstrates how to use probe to intercept and monitor
 // the do_exit() kernel function, which is called when a process exits.
 // We print the exit code parameter to see why processes are exiting.
 
 // Target kernel function signature:
-// do_exit(code: u64) -> void
+// do_exit(code: i64) -> void
 // 
 // The 'code' parameter contains the exit status/signal that caused
-// the process to exit.
+// the process to exit. In the kernel, it's declared as 'long' (signed 64-bit).
 
 
-@kprobe("do_exit")
-fn do_exit(code: u64) -> void {
+@probe("do_exit")
+fn do_exit(code: i64) -> void {
     // Print the exit code parameter
     // This will show us the exit status/signal for the exiting process
-    print("Process exiting with code: %u", code)    
+    print("Process exiting with code: %ld", code)
     return 0
 }
 
@@ -24,17 +24,17 @@ fn main() -> i32 {
     var result = attach(prog, "do_exit", 0)
 
     if (result == 0) {
-        print("kprobe program attached to do_exit successfully")
+        print("probe program attached to do_exit successfully")
         print("Monitoring process exits...")
 
         // In a real scenario, you would wait for events or run for a specific time
         // For this example, we'll just clean up immediately
 
         // Detach the program
         detach(prog)
-        print("kprobe program detached")
+        print("probe program detached")
     } else {
-        print("Failed to attach kprobe program")
+        print("Failed to attach probe program")
         return 1
     }
 

examples/ringbuf_demo.ks

@@ -92,7 +92,8 @@ fn get_timestamp() -> u64 {
 }
 
 // Security monitoring program
-@kprobe("sys_openat") fn security_monitor(dfd: i32, filename: *u8, flags: i32, mode: u16) -> i32 {
+@probe("sys_openat")
+fn security_monitor(dfd: i32, filename: *u8, flags: i32, mode: u16) -> i32 {
   var reserved = security_events.reserve()
   if (reserved != null) {
     // Successfully reserved space - populate security event inline

examples/ringbuf_on_event_demo.ks

@@ -43,7 +43,8 @@ var security_events : ringbuf<SecurityEvent>(8192)
   return XDP_PASS
 }
 
-@kprobe("sys_openat") fn security_monitor(dfd: i32, filename: *u8, flags: i32, mode: u16) -> i32 {
+@probe("sys_openat")
+fn security_monitor(dfd: i32, filename: *u8, flags: i32, mode: u16) -> i32 {
   var reserved = security_events.reserve()
   security_events.submit(reserved)
   return 0

src/ast.ml

@@ -33,9 +33,14 @@ type attribute =
   | SimpleAttribute of string  (* @xdp *)
   | AttributeWithArg of string * string  (* @kprobe("sys_read") *)
 
+(** Probe types for distinguishing between fprobe and kprobe *)
+type probe_type = 
+  | Fprobe    (* Function entrance/exit probe - no offset *)
+  | Kprobe    (* Kernel probe with offset support *)
+
 (** Program types supported by KernelScript *)
 type program_type = 
-  | Xdp | Tc | Kprobe | Tracepoint | StructOps
+  | Xdp | Tc | Probe of probe_type | Tracepoint | StructOps
 
 (** Map types for eBPF maps *)
 type map_type =
@@ -553,7 +558,8 @@ let string_of_position pos =
 let string_of_program_type = function
   | Xdp -> "xdp"
   | Tc -> "tc"
-  | Kprobe -> "kprobe"
+  | Probe Fprobe -> "fprobe"
+  | Probe Kprobe -> "kprobe"
   | Tracepoint -> "tracepoint"
   | StructOps -> "struct_ops"
 

src/btf_parser.ml

@@ -348,7 +348,7 @@ let generate_kernelscript_source template project_name =
         | [] -> ("target_function", "fn() -> i32")
       in
       let func_def = generate_kprobe_function_from_signature first_func first_sig in
-      (comment, first_func, func_def, Some (sprintf "@kprobe(\"%s\")" first_func))
+      (comment, first_func, func_def, Some (sprintf "@probe(\"%s\")" first_func))
     else if template.program_type = "tracepoint" && template.function_signatures <> [] then
       let signature_lines = List.map (fun (event_name, signature) ->
         sprintf "// Tracepoint event: %s -> %s" event_name signature

src/context/context_codegen.ml

@@ -42,6 +42,7 @@ type context_codegen = {
   generate_includes: unit -> string list;
   generate_field_access: string -> string -> string; (* ctx_var -> field_name -> C expression *)
   map_action_constant: int -> string option; (* Map integer to action constant *)
+  generate_function_signature: (string -> (string * string) list -> string -> string) option; (* func_name -> parameters -> return_type -> signature *)
 }
 
 (** Registry for context code generators *)
@@ -95,6 +96,15 @@ let get_context_action_constants ctx_type =
       List.rev (collect_constants [] 0)
   | None -> []
 
+(** Generate custom function signature for a context type *)
+let generate_context_function_signature ctx_type func_name parameters return_type =
+  match get_context_codegen ctx_type with
+  | Some codegen ->
+      (match codegen.generate_function_signature with
+      | Some gen_func -> Some (gen_func func_name parameters return_type)
+      | None -> None)
+  | None -> None
+
 (** Get struct field definitions for a context type as (name, c_type) pairs *)
 let get_context_struct_fields ctx_type =
   match get_context_codegen ctx_type with
@@ -221,6 +231,7 @@ let create_context_codegen_from_btf ctx_type_name btf_type_info =
     generate_includes;
     generate_field_access;
     map_action_constant;
+    generate_function_signature = None;
   }
 
 (** Register context codegen from BTF type information *)

src/context/context_codegen.mli

@@ -40,6 +40,7 @@ type context_codegen = {
   generate_includes: unit -> string list;
   generate_field_access: string -> string -> string;
   map_action_constant: int -> string option;
+  generate_function_signature: (string -> (string * string) list -> string -> string) option;
 }
 
 (** Register a context code generator *)
@@ -63,6 +64,9 @@ val map_context_action_constant : string -> int -> string option
 (** Get all action constants for a context type *)
 val get_context_action_constants : string -> (string * int) list
 
+(** Generate custom function signature for a context type *)
+val generate_context_function_signature : string -> string -> (string * string) list -> string -> string option
+
 (** Get struct field definitions for a context type *)
 val get_context_struct_fields : string -> (string * string) list