Implement kprobe support by extracting kernel function signatures and updating context handling. Remove KprobeContext references and replace with pt_regs.

Author: congwang-mk

src/ast.ml

@@ -74,7 +74,7 @@ and bpf_type =
   | Function of bpf_type list * bpf_type
   | Map of bpf_type * bpf_type * map_type
   (* Built-in context types *)
-  | Xdp_md | KprobeContext | UprobeContext 
+  | Xdp_md | UprobeContext 
   | TracepointContext | LsmContext | CgroupSkbContext
   | Xdp_action
   (* Program reference types *)
@@ -527,7 +527,6 @@ let rec string_of_bpf_type = function
         (string_of_bpf_type value_type)
         (string_of_map_type map_type)
   | Xdp_md -> "xdp_md"
-  | KprobeContext -> "KprobeContext"
   | UprobeContext -> "UprobeContext"
   | TracepointContext -> "TracepointContext"
   | LsmContext -> "LsmContext"

src/btf_binary_parser.ml

@@ -37,6 +37,7 @@ external btf_type_by_id : btf_handle -> int -> (int * string * int * int * int)
 
 external btf_type_get_members : btf_handle -> int -> (string * int) array = "btf_type_get_members_stub"
 external btf_resolve_type : btf_handle -> int -> string = "btf_resolve_type_stub"
+external btf_extract_function_signatures : btf_handle -> string list -> (string * string) list = "btf_extract_function_signatures_stub"
 external btf_free : btf_handle -> unit = "btf_free_stub"
 
 (** Parse BTF file and extract requested types using libbpf *)
@@ -152,4 +153,31 @@ let parse_btf_file btf_path target_types =
   with
   | exn ->
       printf "Error parsing BTF file %s: %s\n" btf_path (Printexc.to_string exn);
+      []
+
+(** Extract kernel function signatures for kprobe targets *)
+let extract_kernel_function_signatures btf_path function_names =
+  try
+    printf "Extracting function signatures from BTF file: %s\n" btf_path;
+    printf "Target functions: %s\n" (String.concat ", " function_names);
+    
+    match btf_new_from_file btf_path with
+    | None -> (
+      printf "Error: Failed to open BTF file %s\n" btf_path;
+      []
+    )
+    | Some btf_handle -> (
+      let signatures = btf_extract_function_signatures btf_handle function_names in
+      btf_free btf_handle;
+      
+      printf "Successfully extracted %d function signatures\n" (List.length signatures);
+      List.iter (fun (name, sig_str) ->
+        printf "  Function: %s -> %s\n" name sig_str
+      ) signatures;
+      
+      signatures
+    )
+  with
+  | exn ->
+      printf "Error extracting function signatures from BTF file %s: %s\n" btf_path (Printexc.to_string exn);
       [] 

src/btf_binary_parser.mli

@@ -28,4 +28,10 @@ type btf_type_info = {
     @param btf_path Path to the binary BTF file
     @param target_types List of type names to extract
     @return List of extracted type definitions in KernelScript format *)
-val parse_btf_file : string -> string list -> btf_type_info list 
+val parse_btf_file : string -> string list -> btf_type_info list
+
+(** Extract kernel function signatures for kprobe targets.
+    @param btf_path Path to the binary BTF file
+    @param function_names List of kernel function names to extract signatures for
+    @return List of (function_name, signature) pairs *)
+val extract_kernel_function_signatures : string -> string list -> (string * string) list 

src/btf_parser.ml

@@ -32,6 +32,7 @@ type program_template = {
   return_type: string;
   includes: string list;
   types: btf_type_info list;
+  function_signatures: (string * string) list; (* Function name and signature for kprobe targets *)
 }
 
 
@@ -106,12 +107,16 @@ let rec get_program_template prog_type btf_path =
 
   let all_types = filtered_types @ hardcoded_types in
 
+  (* No function signatures for generic program templates - kprobe uses specific template *)
+  let function_signatures = [] in
+  
   {
     program_type = prog_type;
     context_type = context_type;
     return_type = return_type;
     includes = ["linux/bpf.h"; "linux/pkt_cls.h"; "linux/if_ether.h"; "linux/ip.h"; "linux/tcp.h"; "linux/udp.h"];
     types = all_types;
+    function_signatures = function_signatures;
   }
 
 (** Extract specific types from BTF file using binary parser *)
@@ -141,7 +146,40 @@ and extract_types_from_btf btf_path type_names =
   | exn ->
       failwith (sprintf "BTF extraction failed: %s" (Printexc.to_string exn))
 
-
+(** Get kprobe program template for a specific target function *)
+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
+  
+  (* Extract specific function signature for the target *)
+  let function_signatures = match btf_path with
+    | Some path when Sys.file_exists path ->
+        printf "🔧 Extracting function signature for %s...\n" target_function;
+        let signatures = Btf_binary_parser.extract_kernel_function_signatures path [target_function] in
+        if signatures = [] then
+          printf "⚠️ Function '%s' not found in BTF - proceeding without signature\n" target_function
+        else
+          printf "✅ Extracted signature for %s\n" target_function;
+        signatures
+    | _ -> []
+  in
+  
+  {
+    program_type = "kprobe";
+    context_type = context_type;
+    return_type = return_type;
+    includes = ["linux/bpf.h"; "linux/pkt_cls.h"; "linux/if_ether.h"; "linux/ip.h"; "linux/tcp.h"; "linux/udp.h"];
+    types = extracted_types;
+    function_signatures = function_signatures;
+  }
 
 (** Extract struct_ops definitions from BTF and generate KernelScript code *)
 let extract_struct_ops_definitions btf_path struct_ops_names =
@@ -230,13 +268,38 @@ let generate_kernelscript_source template project_name =
     | [] -> "0"
   in
 
+  (* Generate function signature comments for kprobe programs *)
+  let (function_signatures_comment, target_function_name) = 
+    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"
+      in
+      (comment, first_func)
+    else 
+      ("", "target_function")
+  in
+  
+  (* Customize attach call for kprobe *)
+  let attach_target = if template.program_type = "kprobe" then target_function_name else "eth0" in
+  let attach_comment = if template.program_type = "kprobe" then 
+    "    // Attach kprobe to target kernel function"
+  else 
+    "    // TODO: Update interface name and attachment parameters"
+  in
+  
   sprintf {|%s
 // Generated by KernelScript compiler with direct BTF parsing
-
+%s
 %s
 
 @%s
-fn %s_handler(ctx: *%s) -> %s {
+fn %s_handler(ctx: %s) -> %s {
     // TODO: Implement your %s logic here
 
     return %s
@@ -245,8 +308,8 @@ fn %s_handler(ctx: *%s) -> %s {
 fn main() -> i32 {
     var prog = load(%s_handler)
 
-    // TODO: Update interface name and attachment parameters
-    var result = attach(prog, "eth0", 0)
+%s
+    var result = attach(prog, "%s", 0)
 
     if (result == 0) {
         print("%s program loaded successfully")
@@ -257,4 +320,4 @@ fn main() -> i32 {
 
     return 0
 }
-|} context_comment type_definitions template.program_type project_name template.context_type template.return_type template.program_type sample_return project_name template.program_type template.program_type
+|} 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

src/btf_parser.mli

@@ -30,11 +30,15 @@ type program_template = {
   return_type: string;
   includes: string list;
   types: btf_type_info list;
+  function_signatures: (string * string) list; (* Function name and signature for kprobe targets *)
 }
 
 (** Get program template based on eBPF program type with optional BTF extraction *)
 val get_program_template : string -> string option -> program_template
 
+(** Get kprobe program template for a specific target function *)
+val get_kprobe_program_template : string -> string option -> program_template
+
 (** Check if a type name is a well-known eBPF kernel type *)
 val is_well_known_kernel_type : string -> bool
 

src/btf_stubs.c

@@ -500,6 +500,83 @@ value btf_resolve_type_stub(value btf_handle, value type_id) {
     CAMLreturn(caml_copy_string("unknown"));
 }
 
+/* Extract kernel function signatures for kprobe targets */
+value btf_extract_function_signatures_stub(value btf_handle, value function_names) {
+    CAMLparam2(btf_handle, function_names);
+    CAMLlocal3(result_list, current, tuple);
+    
+    struct btf *btf = btf_of_value(btf_handle);
+    if (!btf) {
+        CAMLreturn(Val_emptylist);
+    }
+    
+    result_list = Val_emptylist;
+    
+    /* Convert OCaml list to C array */
+    int func_count = 0;
+    value temp = function_names;
+    while (temp != Val_emptylist) {
+        func_count++;
+        temp = Field(temp, 1);
+    }
+    
+    char **target_functions = malloc(func_count * sizeof(char*));
+    temp = function_names;
+    for (int i = 0; i < func_count; i++) {
+        target_functions[i] = String_val(Field(temp, 0));
+        temp = Field(temp, 1);
+    }
+    
+    int nr_types = btf__type_cnt(btf);
+    
+    /* Search for function prototypes */
+    for (int i = 1; i < nr_types; i++) {
+        const struct btf_type *t = btf__type_by_id(btf, i);
+        if (!t) continue;
+        
+        int kind = btf_kind(t);
+        
+        if (kind == BTF_KIND_FUNC) {
+            const char *func_name = btf__name_by_offset(btf, t->name_off);
+            if (!func_name) continue;
+            
+            /* Check if this is one of our target functions */
+            int is_target = 0;
+            for (int j = 0; j < func_count; j++) {
+                if (strcmp(func_name, target_functions[j]) == 0) {
+                    is_target = 1;
+                    break;
+                }
+            }
+            
+            if (is_target) {
+                /* Get the function prototype */
+                const struct btf_type *func_proto = btf__type_by_id(btf, t->type);
+                if (func_proto && btf_kind(func_proto) == BTF_KIND_FUNC_PROTO) {
+                    /* Extract function signature */
+                    char *signature = format_function_prototype(btf, func_proto);
+                    
+                    /* Create tuple (function_name, signature) */
+                    tuple = caml_alloc_tuple(2);
+                    Store_field(tuple, 0, caml_copy_string(func_name));
+                    Store_field(tuple, 1, caml_copy_string(signature));
+                    
+                    /* Add to result list */
+                    current = caml_alloc(2, 0);
+                    Store_field(current, 0, tuple);
+                    Store_field(current, 1, result_list);
+                    result_list = current;
+                    
+                    free(signature);
+                }
+            }
+        }
+    }
+    
+    free(target_functions);
+    CAMLreturn(result_list);
+}
+
 /* Free BTF handle */
 value btf_free_stub(value btf_handle) {
     CAMLparam1(btf_handle);

src/context/dune

@@ -1,5 +1,5 @@
 (library
  (public_name kernelscript.context)
  (name kernelscript_context)
- (modules context_codegen xdp_codegen tc_codegen)
+ (modules context_codegen xdp_codegen tc_codegen kprobe_codegen)
  (libraries unix str))