Implement kprobe function signature parsing and enhance context handling for new signature format. Update code generation to support flexible parameter counts and validate return types for kprobe functions. Refactor include generation to prioritize necessary headers for kprobe programs.

Author: congwang-mk

src/btf_parser.ml

@@ -150,14 +150,9 @@ and extract_types_from_btf btf_path type_names =
 let get_kprobe_program_template target_function btf_path =
   let context_type = "*pt_regs" in
       let return_type = "i32" in
-  let common_types = ["pt_regs"] in
 
-  (* Extract types from BTF - BTF file is required *)
-  let extracted_types = match btf_path with
-    | Some path when Sys.file_exists path -> extract_types_from_btf path common_types
-    | Some path -> failwith (sprintf "BTF file not found: %s" path)
-    | None -> failwith "BTF file path is required. Use --btf-vmlinux-path option."
-  in
+  (* For kprobe, we don't need to extract pt_regs since we're hiding it from users *)
+  let extracted_types = [] in
 
   (* Extract specific function signature for the target *)
   let function_signatures = match btf_path with
@@ -220,6 +215,55 @@ fn main() -> i32 {
     example_usage
     project_name
 
+(** Parse BTF function signature to extract parameter information *)
+let parse_function_signature signature =
+  (* Parse "fn(param1: type1, param2: type2, ...) -> return_type" *)
+  try
+    if String.length signature < 3 || not (String.sub signature 0 3 = "fn(") then
+      failwith "Invalid function signature format"
+    else
+      let paren_start = 3 in
+      let paren_end = String.index signature ')' in
+      let params_str = String.sub signature paren_start (paren_end - paren_start) in
+      
+      (* Parse return type *)
+      let arrow_pos = try Some (String.index signature '>') with Not_found -> None in
+      let return_type = match arrow_pos with
+        | Some pos when pos > paren_end + 2 ->
+            String.trim (String.sub signature (pos + 1) (String.length signature - pos - 1))
+        | _ -> "i32"  (* Default return type for kprobe *)
+      in
+      
+      (* Parse parameters *)
+      let params = if String.trim params_str = "" then
+        []
+      else
+        let param_list = String.split_on_char ',' params_str in
+        List.map (fun param_str ->
+          let trimmed = String.trim param_str in
+          let colon_pos = String.index trimmed ':' in
+          let param_name = String.trim (String.sub trimmed 0 colon_pos) in
+          let param_type = String.trim (String.sub trimmed (colon_pos + 1) (String.length trimmed - colon_pos - 1)) in
+          (param_name, param_type)
+        ) param_list
+      in
+      
+      (params, return_type)
+  with
+  | exn ->
+      printf "⚠️ Warning: Failed to parse function signature '%s': %s\n" signature (Printexc.to_string exn);
+      ([], "i32")  (* Fallback *)
+
+(** Generate kprobe function definition from BTF signature *)
+let generate_kprobe_function_from_signature func_name signature =
+  let (params, return_type) = parse_function_signature signature in
+  let params_str = if params = [] then
+    ""
+  else
+    String.concat ", " (List.map (fun (name, typ) -> sprintf "%s: %s" name typ) params)
+  in
+  sprintf "fn %s(%s) -> %s" func_name params_str return_type
+
 (** Generate KernelScript source code from template *)
 let generate_kernelscript_source template project_name =
   let context_comment = match template.program_type with
@@ -268,21 +312,22 @@ let generate_kernelscript_source template project_name =
     | [] -> "0"
   in
 
-  (* Generate function signature comments for kprobe programs *)
-  let (function_signatures_comment, target_function_name) = 
+  (* Generate function signature comments and actual function definition for kprobe programs *)
+  let (function_signatures_comment, target_function_name, function_definition) = 
     if template.program_type = "kprobe" && template.function_signatures <> [] then
       let signature_lines = List.map (fun (func_name, signature) ->
         sprintf "// Target function: %s -> %s" func_name signature
       ) template.function_signatures in
       let comment = sprintf "\n// Target kernel function signature:\n%s\n" 
         (String.concat "\n" signature_lines) in
-      let first_func = match template.function_signatures with
-        | (name, _) :: _ -> name
-        | [] -> "target_function"
+      let first_func, first_sig = match template.function_signatures with
+        | (name, sig_str) :: _ -> (name, sig_str)
+        | [] -> ("target_function", "fn() -> i32")
       in
-      (comment, first_func)
+      let func_def = generate_kprobe_function_from_signature first_func first_sig in
+      (comment, first_func, func_def)
     else 
-      ("", "target_function")
+      ("", "target_function", sprintf "fn %s_handler(ctx: %s) -> %s" project_name template.context_type template.return_type)
   in
 
   (* Customize attach call for kprobe *)
@@ -293,20 +338,22 @@ let generate_kernelscript_source template project_name =
     "    // TODO: Update interface name and attachment parameters"
   in
 
+  let function_name = if template.program_type = "kprobe" then target_function_name else sprintf "%s_handler" project_name in
+  
   sprintf {|%s
 // Generated by KernelScript compiler with direct BTF parsing
 %s
 %s
 
 @%s
-fn %s_handler(ctx: %s) -> %s {
+%s {
     // TODO: Implement your %s logic here
 
     return %s
 }
 
 fn main() -> i32 {
-    var prog = load(%s_handler)
+    var prog = load(%s)
 
 %s
     var result = attach(prog, "%s", 0)
@@ -320,4 +367,4 @@ fn main() -> i32 {
 
     return 0
 }
-|} context_comment function_signatures_comment type_definitions template.program_type project_name template.context_type template.return_type template.program_type sample_return project_name attach_comment attach_target template.program_type template.program_type
+|} context_comment function_signatures_comment type_definitions template.program_type function_definition template.program_type sample_return function_name attach_comment attach_target template.program_type template.program_type

src/btf_stubs.c

@@ -256,7 +256,7 @@ static char* resolve_type_to_string(struct btf *btf, int type_id) {
         if (kind == BTF_KIND_PTR) {
             const struct btf_type *target = btf__type_by_id(btf, t->type);
             if (target && btf_kind(target) == BTF_KIND_INT && target->size == 1) {
-                return strdup("str");
+                return strdup("*u8");  /* Use *u8 for char* to avoid str parsing issues */
             }
             return strdup("*u8");
         }
@@ -371,7 +371,7 @@ value btf_resolve_type_stub(value btf_handle, value type_id) {
             }
             /* Check if this points to char (string) */
             if (target && btf_kind(target) == BTF_KIND_INT && target->size == 1) {
-                CAMLreturn(caml_copy_string("str"));
+                CAMLreturn(caml_copy_string("*u8"));  /* Use *u8 for char* to avoid str parsing issues */
             }
             /* Other pointer types */
             CAMLreturn(caml_copy_string("*u8"));
@@ -520,7 +520,7 @@ value btf_extract_function_signatures_stub(value btf_handle, value function_name
         temp = Field(temp, 1);
     }
 
-    char **target_functions = malloc(func_count * sizeof(char*));
+    const char **target_functions = malloc(func_count * sizeof(const char*));
     temp = function_names;
     for (int i = 0; i < func_count; i++) {
         target_functions[i] = String_val(Field(temp, 0));

src/context/kprobe_codegen.ml

@@ -94,22 +94,59 @@ let kprobe_field_mappings = [
   });
 ]
 
-(** Generate kprobe-specific includes *)
+(** Dynamic kprobe parameter mappings - populated during compilation based on function signature *)
+let kprobe_parameter_mappings = ref []
+
+(** Register kprobe parameter mappings for a specific function *)
+let register_kprobe_parameter_mappings func_name parameters =
+  let param_mappings = List.mapi (fun i (param_name, param_type) ->
+    let parm_macro = match i with
+      | 0 -> "PT_REGS_PARM1"
+      | 1 -> "PT_REGS_PARM2" 
+      | 2 -> "PT_REGS_PARM3"
+      | 3 -> "PT_REGS_PARM4"
+      | 4 -> "PT_REGS_PARM5"
+      | 5 -> "PT_REGS_PARM6"
+      | _ -> failwith (sprintf "Too many parameters for kprobe function %s (max 6)" func_name)
+    in
+    (param_name, {
+      field_name = param_name;
+      c_expression = (fun ctx_var -> sprintf "%s(%s)" parm_macro ctx_var);
+      requires_cast = false;
+      field_type = param_type;
+    })
+  ) parameters in
+  kprobe_parameter_mappings := param_mappings
+
+(** Clear kprobe parameter mappings *)
+let clear_kprobe_parameter_mappings () =
+  kprobe_parameter_mappings := []
+
+(** Generate kprobe-specific includes with architecture definition at the top *)
 let generate_kprobe_includes () = [
+  "/* Target architecture definition required for PT_REGS_PARM* macros */";
+  "#ifndef __TARGET_ARCH_x86";
+  "#define __TARGET_ARCH_x86";
+  "#endif";
+  "";
   "#include <linux/types.h>";
-  "#include <linux/bpf.h>";
-  "#include <linux/ptrace.h>";
   "#include <bpf/bpf_tracing.h>";
-  "#include <bpf/bpf_helpers.h>";
+  "#include <linux/ptrace.h>";
 ]
 
 (** Generate field access for kprobe context *)
 let generate_kprobe_field_access ctx_var field_name =
   try
-    let (_, field_access) = List.find (fun (name, _) -> name = field_name) kprobe_field_mappings in
+    (* First check dynamic parameter mappings *)
+    let (_, field_access) = List.find (fun (name, _) -> name = field_name) !kprobe_parameter_mappings in
     field_access.c_expression ctx_var
   with Not_found ->
-    failwith ("Unknown kprobe context field: " ^ field_name)
+    try
+      (* Fallback to static register mappings *)
+      let (_, field_access) = List.find (fun (name, _) -> name = field_name) kprobe_field_mappings in
+      field_access.c_expression ctx_var
+    with Not_found ->
+      failwith ("Unknown kprobe context field: " ^ field_name)
 
 (** Map kprobe return constants *)
 let map_kprobe_action_constant = function