Address PR review comments: add overflow checks and improve large program test

- Add overflow checks for calloc in ubpf_create, ubpf_unload_code, and ubpf_set_max_instructions
- Update ubpf_test_large_program to actually JIT compile and execute large programs
- Initialize r0 to 0 in test program and verify result
- Set appropriate JIT buffer size for large programs
- Update documentation to reflect actual test behavior

Co-authored-by: Alan-Jowett <[email protected]>

Author: Copilot

custom_tests/descrs/ubpf_test_large_program.md

@@ -1,16 +1,20 @@
 # Large Program Test
 
-This test validates that programs with more than 65,536 instructions can be loaded correctly when `ubpf_set_max_instructions()` is used to increase the limit.
+This test validates that programs with more than 65,536 instructions can be loaded, JIT compiled, and executed correctly when `ubpf_set_max_instructions()` is used to increase the limit.
 
 The test performs the following:
 1. Creates a VM and sets max_instructions to 100,000
 2. Generates a program with 66,000 instructions (NOP-like JA instructions with offset 0) plus an EXIT
 3. Loads the program into the VM
+4. JIT compiles the program
+5. Executes the JIT-compiled program
 
 This validates that:
 - The type change from uint16_t to uint32_t for num_insts works correctly
 - Programs beyond the old 65,536 limit can be loaded
 - The validation path handles large instruction counts correctly
+- The JIT compiler can handle large programs (> 65,536 instructions)
+- JIT-compiled large programs execute correctly
 
-Note: The test skips interpreter and JIT execution to keep test runtime reasonable.
-Large programs with many sequential NOPs would take prohibitively long to execute.
+Note: The test skips interpreter execution because executing 66,000 sequential NOP-like
+instructions via the interpreter would take prohibitively long for a test.

custom_tests/srcs/ubpf_test_large_program.cc

@@ -25,8 +25,15 @@ main(int argc, char** argv)
     const uint32_t num_instructions = 66000;
     std::vector<struct ebpf_inst> program(num_instructions);
 
-    // Fill with NOP-like JA instructions (jump with offset 0)
-    for (uint32_t i = 0; i < num_instructions - 1; i++) {
+    // First instruction: set r0 = 0
+    program[0].opcode = EBPF_OP_MOV_IMM;
+    program[0].dst = 0;  // r0
+    program[0].src = 0;
+    program[0].offset = 0;
+    program[0].imm = 0;  // value 0
+    
+    // Fill rest with NOP-like JA instructions (jump with offset 0)
+    for (uint32_t i = 1; i < num_instructions - 1; i++) {
         program[i].opcode = EBPF_OP_JA;
         program[i].dst = 0;
         program[i].src = 0;
@@ -53,6 +60,14 @@ main(int argc, char** argv)
         return 1;
     }
 
+    // Set a larger JIT buffer size for large programs
+    // Estimate: ~50 bytes per instruction + overhead
+    size_t jit_buffer_size = num_instructions * 50 + 4096;
+    if (ubpf_set_jit_code_size(vm.get(), jit_buffer_size) != 0) {
+        std::cerr << "Failed to set JIT buffer size" << std::endl;
+        return 1;
+    }
+    
     // Load the program
     char* errmsg = nullptr;
     int result = ubpf_load(vm.get(), program.data(), num_instructions * sizeof(struct ebpf_inst), &errmsg);
@@ -64,8 +79,29 @@ main(int argc, char** argv)
 
     std::cout << "Successfully loaded program with " << num_instructions << " instructions" << std::endl;
 
-    // Skip JIT and interpreter execution for large programs to keep test fast
-    // The key validation is that we can load programs > 65,536 instructions
+    // Test JIT compilation with large program
+    errmsg = nullptr;
+    ubpf_jit_fn jit_fn = ubpf_compile(vm.get(), &errmsg);
+    if (!jit_fn) {
+        std::cerr << "Failed to JIT compile: " << (errmsg ? errmsg : "unknown error") << std::endl;
+        free(errmsg);
+        return 1;
+    }
+    
+    std::cout << "Successfully JIT compiled program with " << num_instructions << " instructions" << std::endl;
+    
+    // Execute via JIT - the program sets r0 to 0 and exits
+    uint64_t jit_result = jit_fn(nullptr, 0);
+    
+    if (jit_result != 0) {
+        std::cerr << "JIT execution returned unexpected value: " << jit_result << " (expected 0)" << std::endl;
+        return 1;
+    }
+    
+    std::cout << "JIT execution result: " << jit_result << " (correct)" << std::endl;
+    
+    // Note: We skip interpreter execution because executing 66,000 sequential NOP-like
+    // instructions via the interpreter would take too long for a test
 
     std::cout << "Test passed!" << std::endl;
     return 0;

vm/ubpf_jit.c

@@ -115,6 +115,11 @@ ubpf_set_max_instructions(struct ubpf_vm* vm, uint32_t max_insts)
 
     // Reallocate local_func_stack_usage if needed
     if (max_insts != vm->max_insts) {
+        // Check for overflow on 32-bit systems where size_t might be 32-bit
+        if (max_insts > SIZE_MAX / sizeof(struct ubpf_stack_usage)) {
+            return -1;
+        }
+        
         struct ubpf_stack_usage* new_stack_usage = calloc(max_insts, sizeof(struct ubpf_stack_usage));
         if (new_stack_usage == NULL) {
             return -1;

vm/ubpf_vm.c

@@ -115,6 +115,12 @@ ubpf_create(void)
     // Initialize max_insts to compile-time default
     vm->max_insts = UBPF_MAX_INSTS;
 
+    // Check for overflow on 32-bit systems where size_t might be 32-bit
+    if (vm->max_insts > SIZE_MAX / sizeof(struct ubpf_stack_usage)) {
+        ubpf_destroy(vm);
+        return NULL;
+    }
+
     vm->local_func_stack_usage = calloc(vm->max_insts, sizeof(struct ubpf_stack_usage));
     if (vm->local_func_stack_usage == NULL) {
         ubpf_destroy(vm);
@@ -375,7 +381,12 @@ ubpf_unload_code(struct ubpf_vm* vm)
 
     // Reset the stack usage amounts when code is unloaded.
     free(vm->local_func_stack_usage);
-    vm->local_func_stack_usage = calloc(vm->max_insts, sizeof(struct ubpf_stack_usage));
+    vm->local_func_stack_usage = NULL;
+    
+    // Check for overflow on 32-bit systems where size_t might be 32-bit
+    if (vm->max_insts > 0 && vm->max_insts <= SIZE_MAX / sizeof(struct ubpf_stack_usage)) {
+        vm->local_func_stack_usage = calloc(vm->max_insts, sizeof(struct ubpf_stack_usage));
+    }
 
     if (vm->jitted) {
         munmap(vm->jitted, vm->jitted_size);