Add kfunc extraction functionality to BTF parser.

Signed-off-by: Cong Wang <cwang@multikernel.io>

Author: congwang-mk

README.md

@@ -255,6 +255,9 @@ kernelscript init probe/sys_read my_tracer
 # Create project with custom BTF path
 kernelscript init --btf-vmlinux-path /custom/path/vmlinux xdp my_project
 
+# Create XDP project with kfuncs extracted
+kernelscript init --kfuncs xdp my_packet_filter
+
 # Create struct_ops project
 kernelscript init tcp_congestion_ops my_congestion_control
 ```

src/btf_binary_parser.ml

@@ -39,6 +39,7 @@ external btf_type_get_members : btf_handle -> int -> (string * int) array = "btf
 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_extract_kernel_struct_names : btf_handle -> string list = "btf_extract_kernel_struct_names_stub"
+external btf_extract_kfuncs : btf_handle -> (string * string) list = "btf_extract_kfuncs_stub"
 external btf_free : btf_handle -> unit = "btf_free_stub"
 
 (** Parse BTF file and extract requested types using libbpf *)
@@ -195,4 +196,30 @@ let extract_all_kernel_struct_names btf_path =
         btf_free btf_handle;
         struct_names
   with
-  | _ -> [] 
+  | _ -> []
+
+(** Extract kfuncs from BTF file using DECL_TAG annotations *)
+let extract_kfuncs_from_btf btf_path =
+  try
+    printf "Extracting kfuncs from BTF file: %s\n" btf_path;
+    
+    match btf_new_from_file btf_path with
+    | None -> (
+      printf "Error: Failed to open BTF file %s\n" btf_path;
+      []
+    )
+    | Some btf_handle -> (
+      let kfuncs = btf_extract_kfuncs btf_handle in
+      btf_free btf_handle;
+      
+      printf "Successfully extracted %d kfuncs\n" (List.length kfuncs);
+      List.iter (fun (name, sig_str) ->
+        printf "  Kfunc: %s -> %s\n" name sig_str
+      ) kfuncs;
+      
+      kfuncs
+    )
+  with
+  | exn ->
+      printf "Error extracting kfuncs from BTF file %s: %s\n" btf_path (Printexc.to_string exn);
+      [] 

src/btf_binary_parser.mli

@@ -39,4 +39,9 @@ val extract_kernel_function_signatures : string -> string list -> (string * stri
 (** Extract all kernel-defined struct names from BTF file.
     @param btf_path Path to the binary BTF file
     @return List of kernel struct names *)
-val extract_all_kernel_struct_names : string -> string list 
+val extract_all_kernel_struct_names : string -> string list
+
+(** Extract kfuncs from BTF file using DECL_TAG annotations.
+    @param btf_path Path to the binary BTF file
+    @return List of (function_name, signature) pairs for functions tagged with "bpf_kfunc" *)
+val extract_kfuncs_from_btf : string -> (string * string) list 

src/btf_parser.ml

@@ -207,16 +207,21 @@ let extract_struct_ops_definitions btf_path struct_ops_names =
       failwith "BTF file path is required for struct_ops extraction. Use --btf-vmlinux-path option."
 
 (** Generate struct_ops template with BTF extraction *)
-let generate_struct_ops_template btf_path struct_ops_names project_name =
+let generate_struct_ops_template ?include_kfuncs btf_path struct_ops_names project_name =
   let struct_ops_definitions = extract_struct_ops_definitions btf_path struct_ops_names in
   let struct_ops_code = String.concat "\n\n" struct_ops_definitions in
 
   let example_usage = List.map (fun name ->
     Struct_ops_registry.generate_struct_ops_usage_example name
   ) struct_ops_names |> String.concat "\n\n" in
 
+  let include_line = match include_kfuncs with
+    | Some kh_filename -> sprintf "\ninclude \"%s\"\n" kh_filename
+    | None -> ""
+  in
+  
   sprintf {|// Generated struct_ops template for %s
-// Extracted from BTF: %s
+// Extracted from BTF: %s%s
 
 %s
 
@@ -230,6 +235,7 @@ fn main() -> i32 {
     (match btf_path with 
      | Some path -> sprintf "definitions from %s" path 
      | None -> "placeholder definitions")
+    include_line
     struct_ops_code
     example_usage
     project_name
@@ -284,7 +290,7 @@ let generate_kprobe_function_from_signature func_name signature =
   sprintf "fn %s(%s) -> %s" func_name params_str return_type
 
 (** Generate KernelScript source code from template *)
-let generate_kernelscript_source ?extra_param template project_name =
+let generate_kernelscript_source ?extra_param ?include_kfuncs template project_name =
   let context_comment = match template.program_type with
     | "xdp" -> "// XDP (eXpress Data Path) program for high-performance packet processing"
     | "tc" -> "// TC (Traffic Control) program for network traffic shaping and filtering"
@@ -410,8 +416,13 @@ let generate_kernelscript_source ?extra_param template project_name =
     else template.program_type
   in
 
+  let include_line = match include_kfuncs with
+    | Some kh_filename -> sprintf "\ninclude \"%s\"\n" kh_filename
+    | None -> ""
+  in
+  
   sprintf {|%s
-// Generated by KernelScript compiler with direct BTF parsing
+// Generated by KernelScript compiler with direct BTF parsing%s
 %s
 %s
 
@@ -437,4 +448,4 @@ fn main() -> i32 {
 
     return 0
 }
-|} context_comment function_signatures_comment type_definitions attribute_line function_definition template.program_type sample_return function_name attach_comment attach_target program_description program_description
+|} context_comment include_line function_signatures_comment type_definitions attribute_line function_definition template.program_type sample_return function_name attach_comment attach_target program_description program_description

src/btf_parser.mli

@@ -49,7 +49,7 @@ val is_well_known_kernel_type : ?btf_path:string -> string -> bool
 val extract_struct_ops_definitions : string option -> string list -> string list
 
 (** Generate struct_ops template with BTF extraction *)
-val generate_struct_ops_template : string option -> string list -> string -> string
+val generate_struct_ops_template : ?include_kfuncs:string -> string option -> string list -> string -> string
 
 (** Generate KernelScript source code from template *)
-val generate_kernelscript_source : ?extra_param:string -> program_template -> string -> string 
+val generate_kernelscript_source : ?extra_param:string -> ?include_kfuncs:string -> program_template -> string -> string 

src/btf_stubs.c

@@ -631,6 +631,64 @@ value btf_extract_kernel_struct_names_stub(value btf_handle) {
     CAMLreturn(result);
 }
 
+/* Extract kfuncs from BTF using DECL_TAG annotations */
+value btf_extract_kfuncs_stub(value btf_handle) {
+    CAMLparam1(btf_handle);
+    CAMLlocal3(result_list, current, tuple);
+    
+    struct btf *btf = btf_of_value(btf_handle);
+    if (!btf) {
+        CAMLreturn(Val_emptylist);
+    }
+    
+    result_list = Val_emptylist;
+    int nr_types = btf__type_cnt(btf);
+    
+    /* First pass: find all DECL_TAG types that reference "bpf_kfunc" */
+    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_DECL_TAG) {
+            const char *tag_name = btf__name_by_offset(btf, t->name_off);
+            if (tag_name && strcmp(tag_name, "bpf_kfunc") == 0) {
+                /* This is a bpf_kfunc tag, get the function it references */
+                int target_id = t->type;
+                const struct btf_type *target_func = btf__type_by_id(btf, target_id);
+                
+                if (target_func && btf_kind(target_func) == BTF_KIND_FUNC) {
+                    const char *func_name = btf__name_by_offset(btf, target_func->name_off);
+                    if (!func_name) continue;
+                    
+                    /* Get the function prototype */
+                    const struct btf_type *func_proto = btf__type_by_id(btf, target_func->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);
+                    }
+                }
+            }
+        }
+    }
+    
+    CAMLreturn(result_list);
+}
+
 /* Free BTF handle */
 value btf_free_stub(value btf_handle) {
     CAMLparam1(btf_handle);