Add tracepoint support. Implement tracepoint program template and enhance BTF extraction.

Author: congwang-mk

SPEC.md

@@ -35,7 +35,7 @@ var flows : hash<IpAddress, PacketStats>(1024)
 KernelScript uses a simple and clear scoping model that eliminates ambiguity:
 
 - **`@helper` functions**: Kernel-shared functions - accessible by all eBPF programs, compile to eBPF bytecode
-- **Attributed functions** (e.g., `@xdp`, `@tc`): eBPF program entry points - compile to eBPF bytecode
+- **Attributed functions** (e.g., `@xdp`, `@tc`, `@tracepoint`): eBPF program entry points - compile to eBPF bytecode
 - **Regular functions**: User space - functions and data structures compile to native executable
 - **Maps and global configs**: Shared resources accessible from both kernel and user space
 - **No wrapper syntax**: Direct, flat structure without unnecessary nesting
@@ -234,6 +234,121 @@ fn trace_write(fd: u32, buf: *u8, count: usize) -> i32 {
 - Validates parameter count (maximum 6 on x86_64)
 - Provides meaningful error messages for unknown functions
 
+#### 3.1.2 Tracepoint Functions with BTF Event Structure Extraction
+
+KernelScript automatically extracts tracepoint event structures from BTF (BPF Type Format) for tracepoint functions, providing type-safe access to tracepoint event data through the appropriate `trace_event_raw_*` structures.
+
+```kernelscript
+@tracepoint("sched/sched_switch")
+fn sched_switch_handler(ctx: *trace_event_raw_sched_switch) -> i32 {
+    // Direct access to tracepoint event fields with correct types
+    // Compiler automatically extracts structure from BTF:
+    // struct trace_event_raw_sched_switch {
+    //     struct trace_entry ent;
+    //     char prev_comm[16];
+    //     pid_t prev_pid;
+    //     int prev_prio;
+    //     long prev_state;
+    //     char next_comm[16];
+    //     pid_t next_pid;
+    //     int next_prio;
+    //     ...
+    // }
+    
+    print("Task switch: %s[%d] -> %s[%d]", 
+          ctx.prev_comm, ctx.prev_pid, 
+          ctx.next_comm, ctx.next_pid)
+    return 0
+}
+
+@tracepoint("syscalls/sys_enter_read")
+fn sys_enter_read_handler(ctx: *trace_event_raw_sys_enter) -> i32 {
+    // Syscall tracepoints use generic sys_enter structure
+    // struct trace_event_raw_sys_enter {
+    //     struct trace_entry ent;
+    //     long id;
+    //     unsigned long args[6];
+    // }
+    
+    var fd = ctx.args[0]
+    var count = ctx.args[2]
+    print("sys_read: fd=%d, count=%d", fd, count)
+    return 0
+}
+
+@tracepoint("net/netif_rx") 
+fn netif_rx_handler(ctx: *trace_event_raw_netif_rx) -> i32 {
+    // Network tracepoint with packet information
+    print("Network packet received")
+    return 0
+}
+```
+
+**Key Benefits:**
+- **Event Structure Access**: Direct access to tracepoint event fields with correct types
+- **Category/Event Organization**: Clear separation using `category/event` format
+- **BTF Integration**: Automatic extraction of `trace_event_raw_*` structures from kernel BTF
+- **Compile-Time Safety**: Type checking for tracepoint context structures
+- **Flexible Event Types**: Support for scheduler, syscall, network, and custom tracepoints
+
+**BTF Structure Mapping:**
+```kernelscript
+// Scheduler tracepoints: trace_event_raw_<event_name>
+@tracepoint("sched/sched_wakeup")
+fn wakeup_handler(ctx: *trace_event_raw_sched_wakeup) -> i32 {
+    // Access scheduler-specific fields
+    print("Waking up PID %d", ctx.pid)
+    return 0
+}
+
+// Syscall enter tracepoints: trace_event_raw_sys_enter (generic)
+@tracepoint("syscalls/sys_enter_open") 
+fn open_handler(ctx: *trace_event_raw_sys_enter) -> i32 {
+    // Access syscall arguments through args array
+    var filename_ptr = ctx.args[0]
+    var flags = ctx.args[1]
+    print("Opening file with flags %d", flags)
+    return 0
+}
+
+// Syscall exit tracepoints: trace_event_raw_sys_exit (generic)
+@tracepoint("syscalls/sys_exit_read")
+fn read_exit_handler(ctx: *trace_event_raw_sys_exit) -> i32 {
+    // Access return value
+    print("sys_read returned %d", ctx.ret)
+    return 0
+}
+
+// Custom subsystem tracepoints: trace_event_raw_<event_name>
+@tracepoint("block/block_rq_complete")
+fn block_complete_handler(ctx: *trace_event_raw_block_rq_complete) -> i32 {
+    // Access block layer specific fields
+    return 0
+}
+```
+
+**Compiler Implementation:**
+- Automatically determines BTF structure name based on category/event:
+  - `syscalls/sys_enter_*` → `trace_event_raw_sys_enter` 
+  - `syscalls/sys_exit_*` → `trace_event_raw_sys_exit`
+  - `<category>/<event>` → `trace_event_raw_<event>`
+- Extracts tracepoint structure definitions from kernel BTF information
+- Generates appropriate `SEC("tracepoint")` section for eBPF programs
+- Validates tracepoint context parameter types at compile time
+- Provides meaningful error messages for unknown tracepoints
+
+**Project Initialization:**
+```bash
+# Initialize project with specific tracepoint
+kernelscript init tracepoint/sched/sched_switch my_scheduler_tracer
+
+# Initialize project with syscall tracepoint  
+kernelscript init tracepoint/syscalls/sys_enter_read my_syscall_tracer
+
+# The init command automatically extracts BTF structures and generates
+# appropriate KernelScript templates with correct context types
+```
+
 ### 3.2 Named Configuration Blocks
 ```kernelscript
 // Named configuration blocks - globally accessible
@@ -545,7 +660,7 @@ fn main() -> i32 {
 ### 3.4 Kernel-Userspace Scoping Model
 
 KernelScript uses a simple and intuitive scoping model:
-- **Attributed functions** (e.g., `@xdp`, `@tc`): Kernel space (eBPF) - compiles to eBPF bytecode
+- **Attributed functions** (e.g., `@xdp`, `@tc`, `@tracepoint`): Kernel space (eBPF) - compiles to eBPF bytecode
 - **`@kfunc` functions**: Kernel modules (full privileges) - exposed to eBPF programs via BTF
 - **`@private` functions**: Kernel modules (full privileges) - internal helpers for kfuncs
 - **Regular functions**: User space - compiles to native executable
@@ -2426,7 +2541,7 @@ value = (value % modulus);                  // From: value %= modulus
 KernelScript functions support both traditional unnamed return types and modern named return values. The complete grammar is defined in Section 15 (Complete Formal Grammar).
 
 Key function types:
-- **eBPF program functions**: Attributed with `@xdp`, `@tc`, etc. - compile to eBPF bytecode
+- **eBPF program functions**: Attributed with `@xdp`, `@tc`, `@tracepoint`, etc. - compile to eBPF bytecode
 - **Helper functions**: Attributed with `@helper` - shared across all eBPF programs
 - **Userspace functions**: No attributes - compile to native executable
 
@@ -3837,12 +3952,6 @@ mod test {
     //   test(program_function, test_context) -> return_value
     //   var result = test(packet_filter, xdp_test_ctx)  // Returns XDP action
     //   var retval = test(kprobe_handler, kprobe_test_ctx)  // Returns kprobe return value
-    // 
-    // The test context struct must match the program type being tested:
-    //   - XdpTestContext for @xdp programs
-    //   - TcTestContext for @tc programs  
-    //   - KprobeTestContext for @kprobe programs
-    //   - etc.
     pub fn test(program_function: FunctionRef, test_context: TestContext) -> i32
 }
 ```
@@ -4316,7 +4425,7 @@ global_variable_declaration = [ "pin" ] [ "local" ] "var" identifier [ ":" type_
 *) 
 
 (* Scoping rules for KernelScript:
-   - Attributed functions (e.g., @xdp, @tc): Kernel space (eBPF) - compiles to eBPF bytecode
+   - Attributed functions (e.g., @xdp, @tc, @tracepoint): Kernel space (eBPF) - compiles to eBPF bytecode
    - Regular functions: User space - compiles to native executable  
    - Maps, global configs, and global variables: Shared between both kernel and user space
    
@@ -4507,7 +4616,7 @@ whitespace = " " | "\t" | "\n" | "\r"
 
 **Function Structure:**
 - `function_declaration` defines functions with optional attributes
-- Functions with attributes (e.g., `@xdp`, `@tc`) are eBPF programs
+- Functions with attributes (e.g., `@xdp`, `@tc`, `@tracepoint`) are eBPF programs
 - Functions without attributes are userspace functions
 - `@helper` functions are shared across all eBPF programs
 

src/btf_parser.ml

@@ -61,7 +61,7 @@ let create_hardcoded_tc_action_enum () = {
 }
 
 (** Get program template based on eBPF program type *)
-let rec get_program_template prog_type btf_path = 
+let get_program_template prog_type btf_path = 
   let (context_type, return_type, common_types) = match prog_type with
     | "xdp" -> ("xdp_md", "xdp_action", [
         "xdp_md"; "xdp_action"
@@ -89,7 +89,16 @@ let rec get_program_template prog_type btf_path =
 
   (* 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 when Sys.file_exists path -> 
+        let binary_types = Btf_binary_parser.parse_btf_file path common_types in
+        (* Convert binary parser types to btf_type_info *)
+        List.map (fun bt -> {
+          name = bt.Btf_binary_parser.name;
+          kind = bt.Btf_binary_parser.kind;
+          size = bt.Btf_binary_parser.size;
+          members = bt.Btf_binary_parser.members;
+          kernel_defined = is_well_known_kernel_type bt.Btf_binary_parser.name;
+        }) binary_types
     | Some path -> failwith (sprintf "BTF file not found: %s" path)
     | None -> failwith "BTF file path is required. Use --btf-vmlinux-path option."
   in
@@ -119,32 +128,57 @@ let rec get_program_template prog_type btf_path =
     function_signatures = function_signatures;
   }
 
-(** Extract specific types from BTF file using binary parser *)
-and extract_types_from_btf btf_path type_names =
-  try
-    printf "Extracting types from BTF file: %s\n" btf_path;
-    let binary_types = Btf_binary_parser.parse_btf_file btf_path type_names in
-    
-    (* Convert binary parser types to btf_type_info *)
-    let converted_types = List.map (fun bt ->
-      {
-        name = bt.Btf_binary_parser.name;
-        kind = bt.Btf_binary_parser.kind;
-        size = bt.Btf_binary_parser.size;
-        members = bt.Btf_binary_parser.members;
-        kernel_defined = is_well_known_kernel_type bt.Btf_binary_parser.name;
-      }
-    ) binary_types in
-    
-    if List.length converted_types > 0 then (
-      printf "Successfully extracted %d types from BTF\n" (List.length converted_types);
-      converted_types
-    ) else (
-      failwith "No types extracted from BTF - requested types not found"
-    )
-  with
-  | exn ->
-      failwith (sprintf "BTF extraction failed: %s" (Printexc.to_string exn))
+(** Get tracepoint program template for a specific category/event *)
+let get_tracepoint_program_template category_event btf_path =
+  (* Parse category and event from the category_event string *)
+  let (category, event) = 
+    if String.contains category_event '/' then
+      let parts = String.split_on_char '/' category_event in
+      match parts with
+      | [cat; evt] -> (cat, evt)
+      | _ -> failwith (sprintf "Invalid tracepoint format '%s'. Use 'category/event'" category_event)
+    else
+      failwith (sprintf "Invalid tracepoint format '%s'. Use 'category/event'" category_event)
+  in
+  
+  (* Determine typedef_name and raw_name based on the user's logic *)
+  let (typedef_name, raw_name) = 
+    if category = "syscalls" && String.starts_with event ~prefix:"sys_enter_" then
+      ("btf_trace_sys_enter", "trace_event_raw_sys_enter")
+    else if category = "syscalls" && String.starts_with event ~prefix:"sys_exit_" then
+      ("btf_trace_sys_exit", "trace_event_raw_sys_exit")
+    else
+      (sprintf "btf_trace_%s" event, sprintf "trace_event_raw_%s" event)
+  in
+  
+  (* Extract the tracepoint structure from BTF *)
+  let common_types = [raw_name; typedef_name] in
+  let extracted_types = match btf_path with
+    | Some path when Sys.file_exists path -> 
+        let binary_types = Btf_binary_parser.parse_btf_file path common_types in
+        (* Convert binary parser types to btf_type_info *)
+        List.map (fun bt -> {
+          name = bt.Btf_binary_parser.name;
+          kind = bt.Btf_binary_parser.kind;
+          size = bt.Btf_binary_parser.size;
+          members = bt.Btf_binary_parser.members;
+          kernel_defined = is_well_known_kernel_type bt.Btf_binary_parser.name;
+        }) binary_types
+    | Some path -> failwith (sprintf "BTF file not found: %s" path)
+    | None -> failwith "BTF file path is required for tracepoint extraction. Use --btf-vmlinux-path option."
+  in
+  
+  (* Create the context type from the extracted struct *)
+  let context_type = sprintf "*%s" raw_name in
+  
+  {
+    program_type = "tracepoint";
+    context_type = context_type;
+    return_type = "i32";
+    includes = ["linux/bpf.h"; "bpf/bpf_helpers.h"; "bpf/bpf_tracing.h"; "linux/trace_events.h"];
+    types = extracted_types;
+    function_signatures = [(sprintf "%s/%s" category event, sprintf "fn(%s) -> i32" context_type)];
+  }
 
 (** Get kprobe program template for a specific target function *)
 let get_kprobe_program_template target_function btf_path =
@@ -312,8 +346,8 @@ let generate_kernelscript_source template project_name =
     | [] -> "0"
   in
 
-  (* Generate function signature comments and actual function definition for kprobe programs *)
-  let (function_signatures_comment, target_function_name, function_definition) = 
+  (* Generate function signature comments and actual function definition for specific program types *)
+  let (function_signatures_comment, target_function_name, function_definition, custom_attribute) = 
     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
@@ -325,27 +359,58 @@ let generate_kernelscript_source template project_name =
         | [] -> ("target_function", "fn() -> i32")
       in
       let func_def = generate_kprobe_function_from_signature first_func first_sig in
-      (comment, first_func, func_def)
+      (comment, first_func, func_def, Some (sprintf "@kprobe(\"%s\")" first_func))
+    else if template.program_type = "tracepoint" && template.function_signatures <> [] then
+      let signature_lines = List.map (fun (event_name, signature) ->
+        sprintf "// Tracepoint event: %s -> %s" event_name signature
+      ) template.function_signatures in
+      let comment = sprintf "\n// Tracepoint event signature:\n%s\n" 
+        (String.concat "\n" signature_lines) in
+      let first_event, _first_sig = match template.function_signatures with
+        | (name, sig_str) :: _ -> (name, sig_str)
+        | [] -> ("category/event", "fn(void*) -> i32")
+      in
+      let func_def = sprintf "fn %s_handler(ctx: %s) -> %s" 
+        (String.map (function '/' -> '_' | c -> c) first_event) template.context_type template.return_type in
+      (comment, first_event, func_def, Some (sprintf "@tracepoint(\"%s\")" first_event))
     else 
-      ("", "target_function", sprintf "fn %s_handler(ctx: %s) -> %s" project_name template.context_type template.return_type)
+      ("", "target_function", sprintf "fn %s_handler(ctx: %s) -> %s" project_name template.context_type template.return_type, None)
   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"
+  (* Use custom attribute if available, otherwise use generic program type attribute *)
+  let attribute_line = match custom_attribute with
+    | Some attr -> attr
+    | None -> "@" ^ template.program_type
   in
 
-  let function_name = if template.program_type = "kprobe" then target_function_name else sprintf "%s_handler" project_name in
+  (* Customize attach call for kprobe/tracepoint *)
+  let attach_target = 
+    if template.program_type = "kprobe" then target_function_name 
+    else if template.program_type = "tracepoint" then target_function_name
+    else "eth0" 
+  in
+  let attach_comment = 
+    if template.program_type = "kprobe" then 
+      "    // Attach kprobe to target kernel function"
+    else if template.program_type = "tracepoint" then
+      "    // Attach tracepoint to target kernel event"
+    else 
+      "    // TODO: Update interface name and attachment parameters"
+  in
+  
+  let function_name = 
+    if template.program_type = "kprobe" then target_function_name 
+    else if template.program_type = "tracepoint" then 
+      String.map (function '/' -> '_' | c -> c) target_function_name ^ "_handler"
+    else sprintf "%s_handler" project_name
+  in
 
   sprintf {|%s
 // Generated by KernelScript compiler with direct BTF parsing
 %s
 %s
 
-@%s
+%s
 %s {
     // TODO: Implement your %s logic here
 
@@ -367,4 +432,4 @@ fn main() -> i32 {
 
     return 0
 }
-|} 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
+|} context_comment function_signatures_comment type_definitions attribute_line function_definition template.program_type sample_return function_name attach_comment attach_target template.program_type template.program_type

src/btf_parser.mli

@@ -39,6 +39,9 @@ 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
 
+(** Get tracepoint program template for a specific target function *)
+val get_tracepoint_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/context/dune

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