stmt_desc = Return (Some {

expr_desc = Literal (IntLit (value, None));

expr_type = Some I32;

expr_pos = dummy_loc;

type_checked = false;

program_context = None;

map_scope = None;

});

stmt_pos = dummy_loc;

(* Simple mock kernel function signatures for testing *)

type mock_kernel_function = {

name: string;

parameters: (string * string) list;

return_type: string;

}

let mock_kernel_functions = [

{

name = "sys_read";

parameters = [("fd", "u32"); ("buf", "*u8"); ("count", "size_t")];

return_type = "isize";

};

{

name = "vfs_write";

parameters = [("file", "*file"); ("buf", "*u8"); ("count", "size_t"); ("pos", "*i64")];

return_type = "isize";

};

{

name = "tcp_sendmsg";

parameters = [("sk", "*sock"); ("msg", "*msghdr"); ("size", "size_t")];

return_type = "i32";

};

]

let source = "@probe(\"sys_read\")

}" in

let ast = parse_string source in

check int "AST should have one declaration" 1 (List.length ast);

match List.hd ast with

| AttributedFunction attr_func ->

check int "Should have one attribute" 1 (List.length attr_func.attr_list);

(match List.hd attr_func.attr_list with

| AttributeWithArg (name, arg) ->

check string "Attribute name" "probe" name;

check string "Attribute argument" "sys_read" arg

| _ -> fail "Expected AttributeWithArg")

let source = "@probe(\"vfs_write\")

}" in

let ast = parse_string source in

check int "AST should have one declaration" 1 (List.length ast);

match List.hd ast with

| AttributedFunction attr_func ->

check int "Should have one attribute" 1 (List.length attr_func.attr_list);

check int "Should have four parameters" 4 (List.length attr_func.attr_function.func_params);

(match List.hd attr_func.attr_list with

| AttributeWithArg (name, arg) ->

check string "Attribute name" "probe" name;

check string "Attribute argument" "vfs_write" arg

| _ -> fail "Expected AttributeWithArg")

(* Test invalid format without target function *)

let source = "@probe

}" in

(* Check that parsing/type checking fails for old format *)

try

let ast = parse_string source in

let _ = type_check_ast ast in

fail "Should have failed parsing old format"

with

(* Test @probe without target function specification *)

let source = "@probe(\"\")

}" in

try

let ast = parse_string source in

let _ = type_check_ast ast in

fail "Should have failed with empty target function"

with

let source = "@probe(\"sys_read\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let source = "@probe(\"sys_read\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

match List.hd typed_ast with

| AttributedFunction attr_func ->

check string "Function name" "sys_read_handler" attr_func.attr_function.func_name;

check int "Parameter count" 3 (List.length attr_func.attr_function.func_params);

(* Verify parameter types *)

(match attr_func.attr_function.func_params with

| [(fd_name, fd_type); (buf_name, buf_type); (count_name, count_type)] ->

check string "First parameter name" "fd" fd_name;

check string "Second parameter name" "buf" buf_name;

check string "Third parameter name" "count" count_name;

check bool "First parameter type should be U32" true

(match fd_type with U32 -> true | _ -> false);

check bool "Second parameter type should be Pointer" true

(match buf_type with Pointer _ -> true | _ -> false);

check bool "Third parameter type should be UserType size_t" true

(match count_type with UserType "size_t" -> true | _ -> false)

| _ -> fail "Expected exactly three parameters")

(* Test valid return types for kprobe - only i32 allowed due to BPF_PROG() constraint *)

let test_cases = [

("i32", "fn handler() -> i32 { return 0 }");

] in

List.iter (fun (ret_type, func_def) ->

let source = "@probe(\"sys_read\")\n" ^ func_def in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

check int (Printf.sprintf "Type checking should succeed for %s return type" ret_type) 1 (List.length typed_ast)

(* 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

| _ -> ()

(* Test rejection of functions with more than 6 parameters *)

let source = "@probe(\"invalid_function\")

}" in

try

let ast = parse_string source in

let _ = type_check_ast ast in

fail "Should have failed with too many parameters"

with

(* Test rejection of direct pt_regs parameter usage *)

let source = "@probe(\"sys_read\")

}" in

try

let ast = parse_string source in

let _ = type_check_ast ast in

fail "Should have failed with pt_regs parameter"

with

let source = "@probe(\"sys_read\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_probe" in

check int "Should generate one program" 1 (List.length (Kernelscript.Ir.get_programs ir_multi_prog));

let program = List.hd (Kernelscript.Ir.get_programs ir_multi_prog) in

check string "Program name" "sys_read_handler" program.name;

check bool "Program type should be Kprobe" true

let source = "@probe(\"tcp_sendmsg\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_probe_complex" in

check int "Should generate one program" 1 (List.length (Kernelscript.Ir.get_programs ir_multi_prog));

let program = List.hd (Kernelscript.Ir.get_programs ir_multi_prog) in

check string "Program name" "tcp_sendmsg_handler" program.name;

check bool "Program type should be Kprobe" true

let source = "@probe(\"vfs_write\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_probe" in

let program = List.hd (Kernelscript.Ir.get_programs ir_multi_prog) in

let main_func = program.entry_function in

(* Test that the function has the correct properties *)

check bool "Function should be marked as main" true main_func.is_main;

let source = "@probe(\"sys_read\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_probe" in

let c_code = generate_c_multi_program ir_multi_prog in

(* Check for fentry section with target function *)

check bool "Should contain SEC(\"fentry/sys_read\")" true

(try ignore (Str.search_forward (Str.regexp_string "SEC(\"fentry/sys_read\")") c_code 0); true with Not_found -> false);

(* Should NOT contain pt_regs parameter for fprobe *)

check bool "Should NOT contain struct pt_regs *ctx" false

let source = "@probe(\"vfs_read+0x10\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_probe" in

let c_code = generate_c_multi_program ir_multi_prog in

(* Check for kprobe section with target function and offset *)

check bool "Should contain SEC(\"kprobe/vfs_read+0x10\")" true

(try ignore (Str.search_forward (Str.regexp_string "SEC(\"kprobe/vfs_read+0x10\")") c_code 0); true with Not_found -> false);

(* Should contain pt_regs parameter for kprobe *)

check bool "Should contain struct pt_regs *ctx" true

let source = "@probe(\"tcp_sendmsg\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_complex_probe" in

let c_code = generate_c_multi_program ir_multi_prog in

(* Check for fentry section with target function *)

check bool "Should contain SEC(\"fentry/tcp_sendmsg\")" true

(try ignore (Str.search_forward (Str.regexp_string "SEC(\"fentry/tcp_sendmsg\")") c_code 0); true with Not_found -> false);

check bool "Should contain tcp_sendmsg_handler function" true

(try ignore (Str.search_forward (Str.regexp_string "tcp_sendmsg_handler") c_code 0); true with Not_found -> false);

(* Should have direct parameters, not pt_regs *)

check bool "Should contain direct parameters" true

(try ignore (Str.search_forward (Str.regexp_string "struct sock* sk") c_code 0); true

with Not_found ->

try ignore (Str.search_forward (Str.regexp_string "struct sock *sk") c_code 0); true

let source = "@probe(\"sys_read\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_probe" in

let c_code = generate_c_multi_program ir_multi_prog in

(* Check for fprobe-specific C code elements *)

check bool "Should contain SEC(\"fentry/sys_read\")" true

(try ignore (Str.search_forward (Str.regexp_string "SEC(\"fentry/sys_read\")") c_code 0); true with Not_found -> false);

check bool "Should contain function definition" true

(try ignore (Str.search_forward (Str.regexp_string "sys_read_handler") c_code 0); true with Not_found -> false);

check bool "Should contain direct parameters" true

(try ignore (Str.search_forward (Str.regexp_string "__u32 fd") c_code 0); true

with Not_found ->

try ignore (Str.search_forward (Str.regexp_string "u32 fd") c_code 0); true

let source = "@probe(\"vfs_read+0x20\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_probe" in

let c_code = generate_c_multi_program ir_multi_prog in

(* Check for kprobe-specific C code elements *)

check bool "Should contain SEC(\"kprobe/vfs_read+0x20\")" true

(try ignore (Str.search_forward (Str.regexp_string "SEC(\"kprobe/vfs_read+0x20\")") c_code 0); true with Not_found -> false);

check bool "Should contain function definition" true

(try ignore (Str.search_forward (Str.regexp_string "vfs_read_handler") c_code 0); true with Not_found -> false);

check bool "Should contain pt_regs parameter" true

let source = "@probe(\"sys_read\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_probe" in

let c_code = generate_c_multi_program ir_multi_prog in

(* Check for fprobe-specific includes *)

check bool "Should include bpf/bpf_helpers.h" true

(try ignore (Str.search_forward (Str.regexp_string "bpf/bpf_helpers.h") c_code 0); true with Not_found -> false);

check bool "Should include vmlinux.h" true

(try ignore (Str.search_forward (Str.regexp_string "vmlinux.h") c_code 0); true with Not_found -> false);

(* fprobe should NOT need linux/ptrace.h *)

check bool "Should NOT include linux/ptrace.h for fprobe" false

let source = "@probe(\"vfs_write+0x8\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_probe" in

let c_code = generate_c_multi_program ir_multi_prog in

(* Check for kprobe-specific elements *)

check bool "Should contain SEC(\"kprobe/vfs_write+0x8\")" true

(try ignore (Str.search_forward (Str.regexp_string "SEC(\"kprobe/vfs_write+0x8\")") c_code 0); true with Not_found -> false);

check bool "Should contain struct pt_regs *ctx parameter" true

(try ignore (Str.search_forward (Str.regexp_string "struct pt_regs *ctx") c_code 0); true with Not_found -> false);

check bool "Should include bpf/bpf_tracing.h for kprobe" true

(* Test target function extraction logic *)

let test_cases = [

("sys_read", "sys_read");

("vfs_write", "vfs_write");

("tcp_sendmsg", "tcp_sendmsg");

("schedule", "schedule");

] in

List.iter (fun (input, expected) ->

(* This tests the internal logic that extracts target function names *)

check string (Printf.sprintf "Target function for %s" input) expected input

(* Test parameter mapping to PT_REGS_PARM macros *)

let test_cases = [

(0, "PT_REGS_PARM1");

(1, "PT_REGS_PARM2");

(2, "PT_REGS_PARM3");

(3, "PT_REGS_PARM4");

(4, "PT_REGS_PARM5");

(5, "PT_REGS_PARM6");

] in

List.iter (fun (index, expected_macro) ->

let actual_macro = match index with

| 0 -> "PT_REGS_PARM1"

| 1 -> "PT_REGS_PARM2"

| 2 -> "PT_REGS_PARM3"

| 3 -> "PT_REGS_PARM4"

| 4 -> "PT_REGS_PARM5"

| 5 -> "PT_REGS_PARM6"

| _ -> "INVALID"

in

check string (Printf.sprintf "Parameter mapping for index %d" index) expected_macro actual_macro

let source = "@probe(\"sys_read\")

}" in

(* Check that compilation fails for invalid return type *)

try

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let _ = generate_ir typed_ast symbol_table "test" in

fail "Should have failed with invalid return type"

with

let source = "@probe(\"invalid_function\")

}" in

(* Check that compilation fails for too many parameters *)

try

let ast = parse_string source in

let _ = type_check_ast ast in

fail "Should have failed with too many parameters"

with

let source = "@probe(\"\")

}" in

try

let ast = parse_string source in

let _ = type_check_ast ast in

fail "Should have failed with empty target function"

with

let source = "@probe(\"sys_open\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_syscall" in

let c_code = generate_c_multi_program ir_multi_prog in

(* Comprehensive end-to-end validation for fprobe *)

check bool "Contains fentry section" true

(try ignore (Str.search_forward (Str.regexp_string "SEC(\"fentry/sys_open\")") c_code 0); true with Not_found -> false);

check bool "Contains function name" true

(try ignore (Str.search_forward (Str.regexp_string "sys_open_handler") c_code 0); true with Not_found -> false);

check bool "Should NOT contain pt_regs parameter for fprobe" false

(try ignore (Str.search_forward (Str.regexp_string "struct pt_regs *ctx") c_code 0); true with Not_found -> false);

check bool "Contains return statement" true

let source = "@probe(\"sys_open+0x4\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_syscall" in

let c_code = generate_c_multi_program ir_multi_prog in

(* Comprehensive end-to-end validation for kprobe *)

check bool "Contains kprobe section" true

(try ignore (Str.search_forward (Str.regexp_string "SEC(\"kprobe/sys_open+0x4\")") c_code 0); true with Not_found -> false);

check bool "Contains function name" true

(try ignore (Str.search_forward (Str.regexp_string "sys_open_handler") c_code 0); true with Not_found -> false);

check bool "Contains pt_regs parameter" true

(try ignore (Str.search_forward (Str.regexp_string "struct pt_regs *ctx") c_code 0); true with Not_found -> false);

check bool "Contains return statement" true

let source = "@probe(\"tcp_sendmsg\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_network" in

let c_code = generate_c_multi_program ir_multi_prog in

check bool "End-to-end fprobe works for network functions" true

(try ignore (Str.search_forward (Str.regexp_string "tcp_sendmsg_handler") c_code 0); true with Not_found -> false);

check bool "Contains fentry section" true

(try ignore (Str.search_forward (Str.regexp_string "SEC(\"fentry/tcp_sendmsg\")") c_code 0); true with Not_found -> false);

check bool "Should NOT contain struct pt_regs parameter for fprobe" false

let source = "@probe(\"sys_read\")

}" in

let ast = parse_string source in

let typed_ast = type_check_ast ast in

let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in

let ir_multi_prog = generate_ir typed_ast symbol_table "test_multiple" in

let c_code = generate_c_multi_program ir_multi_prog in

check int "Should generate two programs" 2 (List.length (Kernelscript.Ir.get_programs ir_multi_prog));

check bool "Contains both function names" true

(try ignore (Str.search_forward (Str.regexp_string "sys_read_handler") c_code 0);

ignore (Str.search_forward (Str.regexp_string "sys_write_handler") c_code 0); true

with Not_found -> false);

(* Check for both fprobe and kprobe sections *)

check bool "Contains fentry section for sys_read" true

(try ignore (Str.search_forward (Str.regexp_string "SEC(\"fentry/sys_read\")") c_code 0); true with Not_found -> false);

check bool "Contains kprobe section for sys_write" true