Update kprobe function handling to enforce modern signature requirements.

Author: congwang-mk

examples/ringbuf_on_event_demo.ks

@@ -43,7 +43,7 @@ var security_events : ringbuf<SecurityEvent>(8192)
   return XDP_PASS
 }
 
-@kprobe fn security_monitor(ctx: *pt_regs) -> i32 {
+@kprobe("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/ebpf_c_codegen.ml

@@ -1470,7 +1470,7 @@ let rec generate_c_value ?(auto_deref_map_access=false) ctx ir_val =
             sprintf "({ struct %s_config *cfg = get_%s_config(); cfg ? cfg->%s : 0; })" 
               config_name config_name field_name
         | _ -> name
-      (* Check if this is a kprobe function parameter in new format *)
+      (* Check if this is a kprobe function parameter *)
       else if ctx.current_function_context_type = Some "kprobe" then
         (try
           (* Try to use kprobe parameter mapping to generate PT_REGS_PARM* access *)
@@ -3340,17 +3340,13 @@ let generate_c_function ctx ir_func =
   in
 
   let params_str = 
-    (* Special handling for kprobe functions with new signature format *)
+    (* Special handling for kprobe functions *)
     match ir_func.func_program_type with
-    | Some Ast.Kprobe when not (List.exists (fun (_, param_type) -> 
-        match param_type with
-        | IRPointer (IRStruct ("pt_regs", _, _), _) -> true
-        | _ -> false
-      ) ir_func.parameters) ->
-        (* New kprobe format: always use struct pt_regs *ctx parameter *)
+    | Some Ast.Kprobe ->
+        (* Kprobe functions always use struct pt_regs *ctx parameter *)
         "struct pt_regs *ctx"
     | _ ->
-        (* Traditional format: use parameters as-is *)
+        (* Other program types: use parameters as-is *)
         String.concat ", " 
           (List.map (fun (name, param_type) ->
              sprintf "%s %s" (ebpf_type_from_ir_type param_type) name

src/ir_function_system.ml

@@ -42,26 +42,19 @@ let validate_function_signature (ir_func : ir_function) : signature_info =
     | _ -> false
   in
 
-  (* Check if this is a kprobe function with new signature format *)
-  let is_kprobe_new_format = match ir_func.func_program_type with
-    | Some Ast.Kprobe -> 
-        (* New format: parameters don't include pt_regs context *)
-        not (List.exists (fun (_, param_type) -> 
-          match param_type with
-          | IRPointer (IRStruct ("pt_regs", _, _), _) -> true
-          | _ -> false
-        ) ir_func.parameters)
+  (* Check if this is a kprobe function *)
+  let is_kprobe_function = match ir_func.func_program_type with
+    | Some Ast.Kprobe -> true
     | _ -> false
   in
 
-  if ir_func.is_main && not is_struct_ops_function && not is_kprobe_new_format then (
+  if ir_func.is_main && not is_struct_ops_function && not is_kprobe_function then (
     if param_count <> 1 then
       errors := "Main function must have exactly one parameter (context)" :: !errors;
     match ir_func.parameters with
     | [(_, IRContext _)] -> ()
     | [(_, IRPointer (IRContext _, _))] -> ()
     | [(_, IRPointer (IRStruct ("__sk_buff", _, _), _))] -> ()  (* Also recognize __sk_buff as TC context *)
-    | [(_, IRPointer (IRStruct ("pt_regs", _, _), _))] -> ()  (* Also recognize pt_regs as kprobe context *)
     | _ -> errors := "Main function parameter must be a context type" :: !errors;
 
     (* Check return type based on context type *)
@@ -72,28 +65,22 @@ let validate_function_signature (ir_func : ir_function) : signature_info =
       | _ -> false
     in
 
-    let is_kprobe_program = match ir_func.parameters with
-      | [(_, IRPointer (IRStruct ("pt_regs", _, _), _))] -> true  (* Recognize pt_regs as kprobe *)
-      | _ -> false
-    in
-    
     match ir_func.return_type with
-    | Some (IRAction _) when not is_tc_program && not is_kprobe_program -> ()  (* Action types for programs that use actions *)
-    | Some (IRI32) when is_tc_program || is_kprobe_program -> ()  (* int return type for TC and kprobe programs *)
-    | Some (IRU32) when is_tc_program || is_kprobe_program -> ()  (* Allow u32/int for TC and kprobe programs *)
+    | Some (IRAction _) when not is_tc_program -> ()  (* Action types for programs that use actions *)
+    | Some (IRI32) when is_tc_program -> ()  (* int return type for TC programs *)
+    | Some (IRU32) when is_tc_program -> ()  (* Allow u32/int for TC programs *)
     | Some _ when is_tc_program -> errors := "TC programs must return int (i32)" :: !errors;
-    | Some _ when is_kprobe_program -> errors := "Kprobe programs must return int (i32)" :: !errors;
-    | Some _ -> errors := "Main function must return an action type (or int for TC/kprobe programs)" :: !errors;
+    | Some _ -> errors := "Main function must return an action type (or int for TC programs)" :: !errors;
     | None -> errors := "Main function must have a return type" :: !errors
   );
 
-  (* Special validation for kprobe functions with new signature format *)
-  if ir_func.is_main && is_kprobe_new_format then (
-    (* Allow flexible parameter count for new kprobe format *)
+  (* Validation for kprobe functions *)
+  if ir_func.is_main && is_kprobe_function then (
+    (* Kprobe functions support up to 6 parameters (kernel function signature) *)
     if param_count > 6 then
       errors := "Kprobe functions support maximum 6 parameters" :: !errors;
 
-    (* Validate return type for new kprobe format *)
+    (* Validate return type for kprobe functions *)
     match ir_func.return_type with
     | Some (IRI32) -> ()  (* Standard kprobe return type *)
     | Some (IRU32) -> ()  (* Allow u32 as well *)

src/ir_generator.ml

@@ -2790,6 +2790,19 @@ let lower_multi_program ast symbol_table source_name =
                     prog_structs = [];
                     prog_pos = attr_func.attr_pos;
                   })
+         | AttributeWithArg (attr_name, _target_func) :: _ ->
+             (* Handle attributes with arguments like @kprobe("sys_read") *)
+             (match attr_name with
+              | "kprobe" -> 
+                  Some {
+                    Ast.prog_name = attr_func.attr_function.func_name;
+                    prog_type = Ast.Kprobe;
+                    prog_functions = [attr_func.attr_function];
+                    prog_maps = [];
+                    prog_structs = [];
+                    prog_pos = attr_func.attr_pos;
+                  }
+              | _ -> None)
          | _ -> None)
     | _ -> None
   ) ast in

src/type_checker.ml

@@ -2936,22 +2936,28 @@ let rec type_check_and_annotate_ast ?symbol_table:(provided_symbol_table=None) ?
           Hashtbl.add ctx.test_functions attr_func.attr_function.func_name ();
 
         (* Extract program type from attribute for context *)
-        let prog_type = match attr_func.attr_list with
+        let (prog_type, kprobe_target) = match attr_func.attr_list with
           | SimpleAttribute prog_type_str :: _ ->
               (match prog_type_str with
-               | "xdp" -> Some Xdp
-               | "tc" -> Some Tc  
-               | "kprobe" -> Some Kprobe
-               | "uprobe" -> Some Uprobe
-               | "tracepoint" -> Some Tracepoint
-               | "lsm" -> Some Lsm
-               | "cgroup_skb" -> Some CgroupSkb
-               | "kfunc" -> None  (* kfuncs don't have program types *)
-               | "private" -> None  (* private functions don't have program types *)
-               | "helper" -> None  (* helper functions don't have program types *)
-               | "test" -> None  (* test functions don't have program types *)
-               | _ -> None)
-          | _ -> None
+               | "xdp" -> (Some Xdp, None)
+               | "tc" -> (Some Tc, None)  
+               | "kprobe" -> 
+                   (* Reject old format: @kprobe without target function *)
+                   type_error ("@kprobe requires target function specification. Use @kprobe(\"function_name\") instead.") attr_func.attr_pos
+               | "uprobe" -> (Some Uprobe, None)
+               | "tracepoint" -> (Some Tracepoint, None)
+               | "lsm" -> (Some Lsm, None)
+               | "cgroup_skb" -> (Some CgroupSkb, None)
+               | "kfunc" -> (None, None)  (* kfuncs don't have program types *)
+               | "private" -> (None, None)  (* private functions don't have program types *)
+               | "helper" -> (None, None)  (* helper functions don't have program types *)
+               | "test" -> (None, None)  (* test functions don't have program types *)
+               | _ -> (None, None))
+          | AttributeWithArg (attr_name, target_func) :: _ ->
+              (match attr_name with
+               | "kprobe" -> (Some Kprobe, Some target_func)
+               | _ -> (None, None))
+          | _ -> (None, None)
         in
 
         (* Validate attributed function signatures based on program type *)
@@ -3003,11 +3009,24 @@ let rec type_check_and_annotate_ast ?symbol_table:(provided_symbol_table=None) ?
                | Some ret_type -> Some (resolve_user_type ctx ret_type)
                | None -> None in
 
-             (* For kprobe functions, accept only the new function signature format with actual kernel function parameters *)
-             let valid_signature = 
-               (* Accept any number of parameters for new format (kernel function signature) *)
-               List.length params >= 0 && List.length params <= 6  (* x86_64 supports max 6 function parameters *)
-             in
+             (* Validate kprobe function - only modern format supported *)
+             (match kprobe_target with
+             | Some _target_func ->
+                 (* Modern format with target function specified *)
+                 (* Check for invalid pt_regs parameter usage *)
+                 List.iter (fun (_, param_type) ->
+                   match param_type with
+                   | Pointer (UserType "pt_regs") ->
+                       type_error ("@kprobe functions should not use pt_regs parameter. Use kernel function parameters directly.") attr_func.attr_pos
+                   | _ -> ()
+                 ) params;
+                 (* Validate signature against BTF if available *)
+                 if List.length params > 6 then
+                   type_error ("Kprobe functions support maximum 6 parameters") attr_func.attr_pos
+             | None ->
+                 (* This case should never be reached due to earlier validation *)
+                 failwith "Internal error: kprobe without target function should have been rejected earlier"
+             );
 
              (* Allow both i32 (standard eBPF kprobe return) and void (matching kernel function signature) *)
              let valid_return_type = match resolved_return_type with
@@ -3017,8 +3036,8 @@ let rec type_check_and_annotate_ast ?symbol_table:(provided_symbol_table=None) ?
                | _ -> false
              in
 
-             if not valid_signature || not valid_return_type then
-               type_error ("@kprobe attributed function must have valid signature (max 6 params) with return type -> i32, u32, or void") attr_func.attr_pos
+             if not valid_return_type then
+               type_error ("@kprobe attributed function must return i32, u32, or void") attr_func.attr_pos
            | Some _ -> () (* Other program types - validation can be added later *)
            | None -> type_error ("Invalid or unsupported attribute") attr_func.attr_pos);
 

tests/test_program_ref.ml

@@ -47,7 +47,7 @@ fn main() -> i32 {
 (** Test program reference with different program types *)
 let test_different_program_types () =
   let program_text = {|
-@kprobe fn kprobe_tracer(ctx: *pt_regs) -> i32 {
+@kprobe("sys_read") fn kprobe_tracer(fd: u32, buf: *u8, count: usize) -> i32 {
   return 0
 }
 

tests/test_ringbuf.ml

@@ -459,7 +459,7 @@ var security_events : ringbuf<SecurityEvent>(8192)
   return XDP_PASS
 }
 
-@kprobe fn security_prog(ctx: *pt_regs) -> i32 {
+@kprobe("sys_read") fn security_prog(fd: u32, buf: *u8, count: usize) -> i32 {
   var sec_event = security_events.reserve()
   if (sec_event != null) {
     security_events.submit(sec_event)
@@ -1043,7 +1043,7 @@ fn handle_security(event: *SecurityEvent) -> i32 {
   return XDP_PASS
 }
 
-@kprobe fn security_monitor(ctx: *pt_regs) -> i32 {
+@kprobe("sys_openat") fn security_monitor(dfd: i32, filename: *u8, flags: i32, mode: u16) -> i32 {
   var sec_event = security_events.reserve()
   if (sec_event != null) {
     sec_event->severity = 1