libunwind_working.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/mman.h>

#define UNW_LOCAL_ONLY
#include <libunwind.h>

// AArch64 registers
#define AARCH64_FP 29   // x29 - Frame Pointer
#define AARCH64_LR 30   // x30 - Link Register
#define AARCH64_SP 31   // sp  - Stack Pointer

typedef struct {
    unw_dyn_info_t* di;
    unw_dyn_region_info_t* region;
} unwind_registration_t;

int print_hello(void) {
    printf("Hello from nested function!\n");
    
    unw_context_t context;
    unw_cursor_t cursor;
    char name[128], off[32];
    unw_getcontext(&context);
    unw_init_local(&cursor, &context);
    
    printf("\nBacktrace:\n");
    while (unw_step(&cursor) > 0) {
        unw_word_t ip, offset;
        unw_get_reg(&cursor, UNW_REG_IP, &ip);
      if (unw_get_proc_name (&cursor, name, sizeof (name), &offset) == 0) {
            printf ("ip = %lx <%s%s>\n", (long) ip, name, off);
      } else {
        printf("  IP: 0x%lx\n", (unsigned long)ip);
      }
    }
    printf("\n");
    
    return 42;
}

typedef int (*trampoline)(void);
typedef int (*call_func_t)(trampoline);

__attribute__((noinline))
int call_func(trampoline function) {
    return function();
}

static unwind_registration_t* register_jit_unwind(void* code_start, size_t code_size) {

    unwind_registration_t* reg = malloc(sizeof(unwind_registration_t));
    if (!reg) return NULL;

    // Create prologue region
    int op_count = 4;
    reg->region = malloc(_U_dyn_region_info_size(op_count));
    if (!reg->region) {
        free(reg);
        return NULL;
    }

    memset(reg->region, 0, _U_dyn_region_info_size(op_count));
    reg->region->op_count = op_count;
    reg->region->next = NULL;
    reg->region->insn_count = 2;  // Two instructions in prologue
    
    int i = 0;

    // Create frame setup operations
    _U_dyn_op_spill_sp_rel(&reg->region->op[i++], 
                          _U_QP_TRUE,    // Always true predicate
                          0,             // At first instruction
                          AARCH64_FP,    // Save FP
                          -32);          // At SP-32

    _U_dyn_op_spill_sp_rel(&reg->region->op[i++],
                          _U_QP_TRUE,    // Always true predicate
                          0,             // At first instruction 
                          AARCH64_LR,    // Save LR
                          -24);          // At SP-24

    _U_dyn_op_save_reg(&reg->region->op[i++],
                       _U_QP_TRUE,      // Always true predicate
                       1,               // At second instruction
                       AARCH64_FP,      // Set FP
                       AARCH64_SP);     // To SP's value

    _U_dyn_op_stop(&reg->region->op[i++]);

    //  up the procedure info
    unw_dyn_proc_info_t proc_info = {
        .name_ptr = (unw_word_t) "my_jit_function blech",
        .handler = 0,
        .flags = 0,
        .regions = reg->region
    };

    // Set up dynamic info
    reg->di = malloc(sizeof(unw_dyn_info_t));
    if (!reg->di) {
        free(reg->region);
        free(reg);
        return NULL;
    }

    memset(reg->di, 0, sizeof(*reg->di));
    reg->di->start_ip = (unw_word_t)code_start;
    reg->di->end_ip = (unw_word_t)code_start + code_size;
    reg->di->gp = 0;
    reg->di->format = UNW_INFO_FORMAT_DYNAMIC;
    reg->di->u.pi = proc_info;

    _U_dyn_register(reg->di);
    
    printf("Registered unwind info for JIT code at %p-%p\n", 
           (void*)reg->di->start_ip, (void*)reg->di->end_ip);
    return reg;
}

// Cleanup unwind registration
static void cleanup_unwind_registration(unwind_registration_t* reg) {
    if (!reg) return;
    
    if (reg->di) {
        _U_dyn_cancel(reg->di);
        free(reg->di);
    }
    
    if (reg->region) {
        free(reg->region);
    }
    
    free(reg);
}

int main() {
    // Get page size for alignment
    size_t page_size = sysconf(_SC_PAGESIZE);
    
    // Get source function details
    void* func_addr = (void*)call_func;
    size_t func_size = 64;  // Increased size since we're calling another function
    
    // Allocate executable memory for JIT code
    void* jit_memory = mmap(NULL, page_size,
                           PROT_READ | PROT_WRITE,
                           MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    
    if (jit_memory == MAP_FAILED) {
        perror("mmap failed");
        return 1;
    }
    
    // Copy the function
    memcpy(jit_memory, func_addr, func_size);
    
    // Make JIT memory executable
    if (mprotect(jit_memory, page_size, PROT_READ | PROT_EXEC) != 0) {
        perror("mprotect failed");
        munmap(jit_memory, page_size);
        return 1;
    }
    
    // Register unwind info
    unwind_registration_t* reg = register_jit_unwind(jit_memory, func_size);
    if (!reg) {
        fprintf(stderr, "Failed to register unwind info\n");
        munmap(jit_memory, page_size);
        return 1;
    }
    
    // Ensure instruction cache coherency
    __builtin___clear_cache(jit_memory, (char*)jit_memory + func_size);
    
    // Cast and call the copied function
    call_func_t jit_func = (call_func_t)jit_memory;
    printf("JIT function returned: %d\n", jit_func(&print_hello));
    
    // Cleanup
    cleanup_unwind_registration(reg);
    munmap(jit_memory, page_size);
    
    return 0;
}