Enforce return type constraint for eBPF probe functions to i32 across the codebase.

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

Author: congwang-mk

SPEC.md

@@ -243,6 +243,12 @@ fn arbitrary_address() -> i32 {
 - **Intelligent Probe Selection**: Automatically chooses fprobe for function entrance (better performance) or kprobe for arbitrary addresses
 - **Type Safety**: Function entrance probes have correct types extracted from kernel BTF information
 
+**Return Type Constraint:**
+- **All probe functions must return `i32`** due to eBPF's `BPF_PROG()` macro constraint
+- The return value controls execution flow: `0` = continue normally, non-zero = may alter behavior
+- This applies regardless of the target kernel function's actual return type (which may be `void`, pointers, etc.)
+- BTF function signature extraction automatically converts all return types to `i32` for consistency
+
 #### 3.1.2 Traffic Control (TC) Programs with Direction Support
 
 TC programs must specify traffic direction for proper kernel attachment point selection.

examples/probe_do_exit.ks

@@ -5,18 +5,21 @@
 // We print the exit code parameter to see why processes are exiting.
 
 // Target kernel function signature:
-// do_exit(code: i64) -> void
+// do_exit(code: i64) -> void (in kernel)
 // 
 // The 'code' parameter contains the exit status/signal that caused
 // the process to exit. In the kernel, it's declared as 'long' (signed 64-bit).
+// 
+// Note: eBPF probe functions must return i32 due to BPF_PROG() constraint,
+// regardless of the target kernel function's return type.
 
 
 @probe("do_exit")
-fn do_exit(code: i64) -> void {
+fn do_exit(code: i64) -> i32 {
     // Print the exit code parameter
     // This will show us the exit status/signal for the exiting process
     print("Process exiting with code: %ld", code)
-    return 0
+    return 0  // Continue normal execution
 }
 
 fn main() -> i32 {

src/btf_stubs.c

@@ -311,13 +311,24 @@ static char* resolve_type_to_string(struct btf *btf, int type_id) {
 }
 
 /* Helper function to format function prototype */
-static char* format_function_prototype(struct btf *btf, const struct btf_type *func_proto) {
+static char* format_function_prototype(struct btf *btf, const struct btf_type *func_proto, int force_i32_return) {
     char result[1024];
     int ret_type_id = func_proto->type;
     int param_count = btf_vlen(func_proto);
 
     /* Get return type */
-    char *ret_type = resolve_type_to_string(btf, ret_type_id);
+    char *original_ret_type = resolve_type_to_string(btf, ret_type_id);
+    
+    /* Use original return type unless forced to i32 for eBPF probe/kfunc functions */
+    const char *ret_type;
+    if (force_i32_return) {
+        /* eBPF probe/kfunc functions must return i32 due to BPF_PROG() constraint,
+         * regardless of the kernel function's actual return type */
+        ret_type = "i32";
+    } else {
+        /* Preserve original return type for struct_ops and other contexts */
+        ret_type = original_ret_type;
+    }
 
     /* Start building the function signature */
     snprintf(result, sizeof(result), "fn(");
@@ -353,7 +364,7 @@ static char* format_function_prototype(struct btf *btf, const struct btf_type *f
     snprintf(closing, sizeof(closing), ") -> %s", ret_type);
     strncat(result, closing, sizeof(result) - strlen(result) - 1);
 
-    free(ret_type);
+    free(original_ret_type);
     return strdup(result);
 }
 
@@ -384,7 +395,7 @@ value btf_resolve_type_stub(value btf_handle, value type_id) {
             /* Check if this points to a function prototype */
             const struct btf_type *target = btf__type_by_id(btf, t->type);
             if (target && btf_kind(target) == BTF_KIND_FUNC_PROTO) {
-                char *func_sig = format_function_prototype(btf, target);
+                char *func_sig = format_function_prototype(btf, target, 0); /* Don't force i32 for general BTF resolution */
                 value result = caml_copy_string(func_sig);
                 free(func_sig);
                 CAMLreturn(result);
@@ -474,7 +485,7 @@ value btf_resolve_type_stub(value btf_handle, value type_id) {
             CAMLreturn(caml_copy_string("fwd"));
         }
         case BTF_KIND_FUNC_PROTO: {
-            char *func_sig = format_function_prototype(btf, t);
+            char *func_sig = format_function_prototype(btf, t, 0); /* Don't force i32 for general BTF resolution */
             value result = caml_copy_string(func_sig);
             free(func_sig);
             CAMLreturn(result);
@@ -573,8 +584,8 @@ value btf_extract_function_signatures_stub(value btf_handle, value function_name
                 /* 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);
+                    /* Extract function signature - force i32 return for probe functions */
+                    char *signature = format_function_prototype(btf, func_proto, 1);
 
                     /* Create tuple (function_name, signature) */
                     tuple = caml_alloc_tuple(2);
@@ -665,8 +676,7 @@ value btf_extract_kfuncs_stub(value btf_handle) {
                     /* 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);
+                        char *signature = format_function_prototype(btf, func_proto, 0);
 
                         /* Create tuple (function_name, signature) */
                         tuple = caml_alloc_tuple(2);

src/type_checker.ml

@@ -3054,16 +3054,14 @@ let rec type_check_and_annotate_ast ?symbol_table:(provided_symbol_table=None) ?
                  failwith (sprintf "Internal error: %s without target function should have been rejected earlier" probe_type_name)
              );
 
-             (* Allow both i32 (standard eBPF probe return) and void (matching kernel function signature) *)
+             (* Require i32 return type for eBPF probe functions - BPF_PROG() always returns int *)
              let valid_return_type = match resolved_return_type with
                | Some I32 -> true   (* Standard eBPF probe return type *)
-               | Some Void -> true  (* Allow void to match kernel function signatures *)
-               | Some U32 -> true   (* Allow u32 as alternative to i32 *)
                | _ -> false
              in
 
              if not valid_return_type then
-               type_error (sprintf "@%s attributed function must return i32, u32, or void" probe_type_name) attr_func.attr_pos
+               type_error (sprintf "@%s attributed function must return i32" probe_type_name) attr_func.attr_pos
            | Some _ -> () (* Other program types - validation can be added later *)
            | None -> type_error ("Invalid or unsupported attribute") attr_func.attr_pos);
 

tests/test_probe.ml

@@ -172,11 +172,9 @@ fn sys_read_handler(fd: u32, buf: *u8, count: usize) -> i32 {
   | _ -> fail "Expected AttributedFunction"
 
 let test_probe_return_type_validation _ =
-  (* Test valid return types for kprobe *)
+  (* Test valid return types for kprobe - only i32 allowed due to BPF_PROG() constraint *)
   let test_cases = [
     ("i32", "fn handler() -> i32 { return 0 }");
-    ("void", "fn handler() -> void { }");
-    ("u32", "fn handler() -> u32 { return 0 }");
   ] in
 
   List.iter (fun (ret_type, func_def) ->
@@ -186,6 +184,25 @@ let test_probe_return_type_validation _ =
     check int (Printf.sprintf "Type checking should succeed for %s return type" ret_type) 1 (List.length typed_ast)
   ) test_cases
 
+let test_probe_invalid_return_types _ =
+  (* Test that void, u32, and other invalid return types are rejected for probe functions *)
+  let invalid_cases = [
+    ("void", "fn handler() -> void { }");
+    ("u32", "fn handler() -> u32 { return 0 }");
+    ("str", "fn handler() -> str(32) { return \"test\" }");
+    ("bool", "fn handler() -> bool { return true }");
+  ] in
+  
+  List.iter (fun (ret_type, func_def) ->
+    let source = "@probe(\"sys_read\")\n" ^ func_def in
+    try
+      let ast = parse_string source in
+      let _ = type_check_ast ast in
+      fail (Printf.sprintf "Should have rejected %s return type for probe function" ret_type)
+    with
+    | _ -> check bool (Printf.sprintf "Correctly rejected %s return type" ret_type) true true
+  ) invalid_cases
+
 let test_probe_too_many_parameters _ =
   (* Test rejection of functions with more than 6 parameters *)
   let source = "@probe(\"invalid_function\")
@@ -579,6 +596,7 @@ let type_checking_tests = [
   "probe type checking", `Quick, test_probe_type_checking;
   "probe parameter validation", `Quick, test_probe_parameter_validation;
   "probe return type validation", `Quick, test_probe_return_type_validation;
+  "probe invalid return types", `Quick, test_probe_invalid_return_types;
   "probe too many parameters", `Quick, test_probe_too_many_parameters;
   "probe pt_regs rejection", `Quick, test_probe_pt_regs_rejection;
 ]