Update eBPF C code generation to utilize vmlinux.h for kernel types, simplifying include management and removing unnecessary context-specific headers.

Author: congwang-mk

src/ebpf_c_codegen.ml

@@ -1109,72 +1109,39 @@ let generate_declarations_in_source_order ctx _ir_multi_program type_aliases =
 (** Generate standard eBPF includes *)
 
 let generate_includes ctx ?(program_types=[]) () =
-  (* Essential kernel types must come first *)
-  let base_type_includes = [
-    "#include <linux/types.h>";
+  (* Use vmlinux.h which contains all kernel types from BTF *)
+  let vmlinux_includes = [
+    "#include \"vmlinux.h\"";
   ] in
 
+  (* Only include essential eBPF helpers, vmlinux.h provides all kernel types *)
   let standard_includes = [
-    "#include <linux/bpf.h>";
     "#include <bpf/bpf_helpers.h>";
-    "#include <linux/if_ether.h>";
-    "#include <linux/ip.h>";
-    "#include <linux/in.h>";
-    "#include <linux/if_xdp.h>";
   ] in
 
-  (* Get context-specific includes *)
-  let context_includes = List.fold_left (fun acc prog_type ->
-    let context_type = match prog_type with
-      | Ast.Xdp -> Some "xdp"
-      | Ast.Tc -> Some "tc"
-      | Ast.Kprobe -> Some "kprobe"
-      | Ast.Tracepoint -> Some "tracepoint"
-      | _ -> None
-    in
-    match context_type with
-    | Some ctx_type -> 
-        let includes = Kernelscript_context.Context_codegen.get_context_includes ctx_type in
-        acc @ includes
-    | None -> acc
-  ) [] program_types in
+  (* When using vmlinux.h, we don't need context-specific kernel headers *)
+  (* vmlinux.h contains all kernel types, so we only need eBPF helpers *)
+  let context_includes = [] in (* No additional kernel headers needed with vmlinux.h *)
 
   (* Remove duplicates between all include sets *)
-  let all_base_includes = base_type_includes @ standard_includes in
+  let all_base_includes = vmlinux_includes @ standard_includes in
   let unique_context_includes = List.filter (fun inc -> 
     not (List.mem inc all_base_includes)) context_includes in
 
   (* For kprobe programs, only emit kprobe includes which contain everything needed *)
   let has_kprobe = List.exists (function Ast.Kprobe -> true | _ -> false) program_types in
   if has_kprobe then (
-    (* Emit ONLY kprobe includes which contain architecture definition and all needed headers *)
-    let kprobe_includes = Kernelscript_context.Context_codegen.get_context_includes "kprobe" in
-    List.iter (emit_line ctx) kprobe_includes;
-    emit_blank_line ctx;
-    
-    (* Add additional standard includes that aren't in kprobe includes *)
-    let all_additional_includes = [
-      "#include <linux/types.h>";
-      "#include <linux/bpf.h>";
+    (* With vmlinux.h, we only need it and eBPF helpers, no additional kernel headers *)
+    let vmlinux_and_helpers = [
+      "#include \"vmlinux.h\"";
       "#include <bpf/bpf_helpers.h>";
-      "#include <linux/if_ether.h>";
-      "#include <linux/ip.h>";
-      "#include <linux/in.h>";
-      "#include <linux/if_xdp.h>";
     ] in
-    let additional_includes = List.filter (fun inc -> 
-      not (List.mem inc kprobe_includes)) all_additional_includes in
-    
-    (* Add context-specific includes for non-kprobe programs *)
-    let filtered_context_includes = List.filter (fun inc -> 
-      not (List.mem inc kprobe_includes) && not (List.mem inc additional_includes)) unique_context_includes in
 
-    List.iter (emit_line ctx) additional_includes;
-    List.iter (emit_line ctx) filtered_context_includes;
+    List.iter (emit_line ctx) vmlinux_and_helpers;
     emit_blank_line ctx
   ) else (
-    (* For non-kprobe programs, use standard processing *)
-    let all_includes = standard_includes @ unique_context_includes @ base_type_includes in
+    (* For non-kprobe programs, use vmlinux.h and standard processing *)
+    let all_includes = vmlinux_includes @ standard_includes @ unique_context_includes in
     List.iter (emit_line ctx) all_includes;
     emit_blank_line ctx;
 

src/main.ml

@@ -1067,6 +1067,12 @@ uninstall-module:
 
       let kernel_module_clean = if has_kernel_module then " $(KERNEL_MODULE_OBJ) modules.order Module.symvers .*.cmd" else "" in
 
+      (* Determine BTF path for vmlinux.h generation *)
+      let btf_path_for_makefile = match btf_vmlinux_path with
+        | Some path -> path
+        | None -> "/sys/kernel/btf/vmlinux" (* Default fallback *)
+      in
+      
       let makefile_content = Printf.sprintf {|# Multi-Program eBPF Makefile
 # Generated by KernelScript compiler
 
@@ -1076,7 +1082,7 @@ CC = gcc
 
 # BPF compilation flags
 BPF_CFLAGS = -target bpf -O2 -Wall -Wextra -g
-BPF_INCLUDES = -I/usr/include -I/usr/include/x86_64-linux-gnu
+BPF_INCLUDES = -I.
 
 # Userspace compilation flags
 CFLAGS = -Wall -Wextra -O2 -fPIC
@@ -1086,6 +1092,7 @@ LIBS = -lbpf -lelf -lz%s
 BPF_OBJ = %s.ebpf.o
 USERSPACE_BIN = %s
 SKELETON_H = %s.skel.h
+VMLINUX_H = vmlinux.h
 
 # Source files
 BPF_SRC = %s.ebpf.c
@@ -1094,8 +1101,15 @@ USERSPACE_SRC = %s.c
 # Default target
 all:%s $(BPF_OBJ) $(SKELETON_H) $(USERSPACE_BIN)%s
 
-# Compile eBPF C to object file
-$(BPF_OBJ): $(BPF_SRC)
+# Generate vmlinux.h from BTF (requires bpftool)
+$(VMLINUX_H):
+	@echo "Generating vmlinux.h from kernel BTF..."
+	@which bpftool > /dev/null || (echo "Error: bpftool not found. Please install bpftool package." && exit 1)
+	bpftool btf dump file %s format c > $@
+	@echo "Generated vmlinux.h successfully"
+
+# Compile eBPF C to object file (depends on vmlinux.h)
+$(BPF_OBJ): $(BPF_SRC) $(VMLINUX_H)
 	$(BPF_CC) $(BPF_CFLAGS) $(BPF_INCLUDES) -c $< -o $@
 
 # Generate skeleton header
@@ -1110,7 +1124,7 @@ $(USERSPACE_BIN): $(USERSPACE_SRC) $(SKELETON_H)%s
 %s
 # Clean generated files
 clean:
-	rm -f $(BPF_OBJ) $(SKELETON_H) $(USERSPACE_BIN)%s%s
+	rm -f $(BPF_OBJ) $(SKELETON_H) $(USERSPACE_BIN) $(VMLINUX_H)%s%s
 
 # Build just the eBPF object without skeleton (for testing)
 ebpf-only: $(BPF_OBJ)
@@ -1129,7 +1143,7 @@ help:
 
 .PHONY: all clean run ebpf-only help%s
 |} kernel_module_vars base_name base_name base_name base_name base_name 
-       shared_lib_deps kernel_module_target shared_lib_deps
+       shared_lib_deps kernel_module_target btf_path_for_makefile shared_lib_deps
        (if shared_lib_targets = "" then "" else (" " ^ String.concat " " (List.map (fun import -> "./" ^ import.Import_resolver.module_name ^ ".so") ks_imports)))
        python_flags
        kernel_module_rules 

tests/test_ebpf_c_codegen.ml

@@ -231,7 +231,7 @@ let test_complete_program () =
   let c_code = compile_multi_to_c multi_ir in
 
   (* Verify the generated C code contains expected elements *)
-  check bool "program contains linux/bpf.h include" true (contains_substr c_code "#include <linux/bpf.h>");
+  check bool "program contains vmlinux.h include" true (contains_substr c_code "#include \"vmlinux.h\"");
   check bool "program contains map name" true (contains_substr c_code "packet_count");
   check bool "program contains maps section" true (contains_substr c_code "SEC(\".maps\")");
   check bool "program contains xdp section" true (contains_substr c_code "SEC(\"xdp\")");

tests/test_ebpf_string_generation.ml

@@ -115,7 +115,7 @@ let test_ebpf_compilation () =
     check bool "contains SEC xdp" true
       (contains_pattern ebpf_code "SEC(\"xdp\")");
     check bool "contains includes" true
-      (contains_pattern ebpf_code "#include.*linux/bpf.h");
+      (contains_pattern ebpf_code "#include.*vmlinux.h");
     check bool "contains license" true
       (contains_pattern ebpf_code "SEC(\"license\")");