Fix critical memory allocation error handling

This adds proper NULL checks for all critical malloc/calloc/mpool_alloc
calls that previously relied solely on assert(), which compiles out in
production builds with -DNDEBUG.

It also fixes critical JIT initialization error handling and syscall
memory safety issues.

Author: jserv

src/cache.c

@@ -285,6 +285,17 @@ void *cache_put(cache_t *cache, uint32_t key, void *value)
     }
 
     cache_entry_t *new_entry = calloc(1, sizeof(cache_entry_t));
+    if (unlikely(!new_entry)) {
+        /* Allocation failed - restore replaced entry if exists */
+        if (replaced) {
+            replaced->alive = true;
+            list_del_init(&replaced->list);
+            list_add(&replaced->list, &cache->list);
+            cache->size++;
+            cache->ghost_list_size--;
+        }
+        return NULL;
+    }
     assert(new_entry);
 
     INIT_LIST_HEAD(&new_entry->list);

src/emulate.c

@@ -273,6 +273,8 @@ HASH_FUNC_IMPL(map_hash, BLOCK_MAP_CAPACITY_BITS, 1 << BLOCK_MAP_CAPACITY_BITS)
 static block_t *block_alloc(riscv_t *rv)
 {
     block_t *block = mpool_alloc(rv->block_mp);
+    if (unlikely(!block))
+        return NULL;
     assert(block);
     block->n_insn = 0;
 #if RV32_HAS(JIT)
@@ -619,13 +621,15 @@ FORCE_INLINE bool insn_is_indirect_branch(uint8_t opcode)
     }
 }
 
-static void block_translate(riscv_t *rv, block_t *block)
+static bool block_translate(riscv_t *rv, block_t *block)
 {
 retranslate:
     block->pc_start = block->pc_end = rv->PC;
 
     rv_insn_t *prev_ir = NULL;
     rv_insn_t *ir = mpool_calloc(rv->block_ir_mp);
+    if (unlikely(!ir))
+        return false;
     block->ir_head = ir;
 
     /* translate the basic block */
@@ -665,6 +669,8 @@ static void block_translate(riscv_t *rv, block_t *block)
         if (insn_is_branch(ir->opcode)) {
             if (insn_is_indirect_branch(ir->opcode)) {
                 ir->branch_table = calloc(1, sizeof(branch_history_table_t));
+                if (unlikely(!ir->branch_table))
+                    return false;
                 assert(ir->branch_table);
                 memset(ir->branch_table->PC, -1,
                        sizeof(uint32_t) * HISTORY_SIZE);
@@ -673,36 +679,44 @@ static void block_translate(riscv_t *rv, block_t *block)
         }
 
         ir = mpool_calloc(rv->block_ir_mp);
+        if (unlikely(!ir))
+            return false;
     }
 
     assert(prev_ir);
     block->ir_tail = prev_ir;
     block->ir_tail->next = NULL;
+    return true;
 }
 
 #if RV32_HAS(MOP_FUSION)
-#define COMBINE_MEM_OPS(RW)                                       \
-    next_ir = ir->next;                                           \
-    count = 1;                                                    \
-    while (1) {                                                   \
-        if (next_ir->opcode != IIF(RW)(rv_insn_lw, rv_insn_sw))   \
-            break;                                                \
-        count++;                                                  \
-        if (!next_ir->next)                                       \
-            break;                                                \
-        next_ir = next_ir->next;                                  \
-    }                                                             \
-    if (count > 1) {                                              \
-        ir->opcode = IIF(RW)(rv_insn_fuse4, rv_insn_fuse3);       \
-        ir->fuse = malloc(count * sizeof(opcode_fuse_t));         \
-        assert(ir->fuse);                                         \
-        ir->imm2 = count;                                         \
-        memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));              \
-        ir->impl = dispatch_table[ir->opcode];                    \
-        next_ir = ir->next;                                       \
-        for (int j = 1; j < count; j++, next_ir = next_ir->next)  \
-            memcpy(ir->fuse + j, next_ir, sizeof(opcode_fuse_t)); \
-        remove_next_nth_ir(rv, ir, block, count - 1);             \
+#define COMBINE_MEM_OPS(RW)                                           \
+    next_ir = ir->next;                                               \
+    count = 1;                                                        \
+    while (1) {                                                       \
+        if (next_ir->opcode != IIF(RW)(rv_insn_lw, rv_insn_sw))       \
+            break;                                                    \
+        count++;                                                      \
+        if (!next_ir->next)                                           \
+            break;                                                    \
+        next_ir = next_ir->next;                                      \
+    }                                                                 \
+    if (count > 1) {                                                  \
+        ir->opcode = IIF(RW)(rv_insn_fuse4, rv_insn_fuse3);           \
+        ir->fuse = malloc(count * sizeof(opcode_fuse_t));             \
+        if (unlikely(!ir->fuse)) {                                    \
+            ir->opcode = IIF(RW)(rv_insn_lw, rv_insn_sw);             \
+            count = 1; /* Degrade to non-fused operation */           \
+        } else {                                                      \
+            assert(ir->fuse);                                         \
+            ir->imm2 = count;                                         \
+            memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));              \
+            ir->impl = dispatch_table[ir->opcode];                    \
+            next_ir = ir->next;                                       \
+            for (int j = 1; j < count; j++, next_ir = next_ir->next)  \
+                memcpy(ir->fuse + j, next_ir, sizeof(opcode_fuse_t)); \
+            remove_next_nth_ir(rv, ir, block, count - 1);             \
+        }                                                             \
     }
 
 static inline void remove_next_nth_ir(const riscv_t *rv,
@@ -762,16 +776,20 @@ static void match_pattern(riscv_t *rv, block_t *block)
                     next_ir = next_ir->next;
                 }
                 if (count > 1) {
-                    ir->opcode = rv_insn_fuse1;
                     ir->fuse = malloc(count * sizeof(opcode_fuse_t));
-                    assert(ir->fuse);
-                    ir->imm2 = count;
-                    memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));
-                    ir->impl = dispatch_table[ir->opcode];
-                    next_ir = ir->next;
-                    for (int j = 1; j < count; j++, next_ir = next_ir->next)
-                        memcpy(ir->fuse + j, next_ir, sizeof(opcode_fuse_t));
-                    remove_next_nth_ir(rv, ir, block, count - 1);
+                    if (likely(ir->fuse)) {
+                        ir->opcode = rv_insn_fuse1;
+                        assert(ir->fuse);
+                        ir->imm2 = count;
+                        memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));
+                        ir->impl = dispatch_table[ir->opcode];
+                        next_ir = ir->next;
+                        for (int j = 1; j < count; j++, next_ir = next_ir->next)
+                            memcpy(ir->fuse + j, next_ir,
+                                   sizeof(opcode_fuse_t));
+                        remove_next_nth_ir(rv, ir, block, count - 1);
+                    }
+                    /* If malloc failed, degrade gracefully to non-fused ops */
                 }
                 break;
             }
@@ -803,15 +821,18 @@ static void match_pattern(riscv_t *rv, block_t *block)
             }
             if (count > 1) {
                 ir->fuse = malloc(count * sizeof(opcode_fuse_t));
-                assert(ir->fuse);
-                memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));
-                ir->opcode = rv_insn_fuse5;
-                ir->imm2 = count;
-                ir->impl = dispatch_table[ir->opcode];
-                next_ir = ir->next;
-                for (int j = 1; j < count; j++, next_ir = next_ir->next)
-                    memcpy(ir->fuse + j, next_ir, sizeof(opcode_fuse_t));
-                remove_next_nth_ir(rv, ir, block, count - 1);
+                if (likely(ir->fuse)) {
+                    assert(ir->fuse);
+                    memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));
+                    ir->opcode = rv_insn_fuse5;
+                    ir->imm2 = count;
+                    ir->impl = dispatch_table[ir->opcode];
+                    next_ir = ir->next;
+                    for (int j = 1; j < count; j++, next_ir = next_ir->next)
+                        memcpy(ir->fuse + j, next_ir, sizeof(opcode_fuse_t));
+                    remove_next_nth_ir(rv, ir, block, count - 1);
+                }
+                /* If malloc failed, degrade gracefully to non-fused ops */
             }
             break;
         }
@@ -881,8 +902,11 @@ static block_t *block_find_or_translate(riscv_t *rv)
 #endif
     /* allocate a new block */
     next_blk = block_alloc(rv);
+    if (unlikely(!next_blk))
+        return NULL;
 
-    block_translate(rv, next_blk);
+    if (unlikely(!block_translate(rv, next_blk)))
+        return NULL;
 
 #if RV32_HAS(JIT) && RV32_HAS(SYSTEM)
     /*
@@ -1129,6 +1153,10 @@ void rv_step(void *arg)
         else if (!block->compiled && block->n_invoke >= THRESHOLD) {
             block->compiled = true;
             queue_entry_t *entry = malloc(sizeof(queue_entry_t));
+            if (unlikely(!entry)) {
+                /* Malloc failed - degrade to tier-1 JIT, don't queue for T2C */
+                continue;
+            }
             entry->block = block;
             pthread_mutex_lock(&rv->wait_queue_lock);
             list_add(&entry->list, &rv->wait_queue);

src/io.c

@@ -23,6 +23,8 @@ memory_t *memory_new(uint32_t size)
         return NULL;
 
     memory_t *mem = malloc(sizeof(memory_t));
+    if (!mem)
+        return NULL;
     assert(mem);
 #if HAVE_MMAP
     data_memory_base = mmap(NULL, size, PROT_READ | PROT_WRITE,

src/jit.c

@@ -2330,7 +2330,10 @@ void jit_translate(riscv_t *rv, block_t *block)
 struct jit_state *jit_state_init(size_t size)
 {
     struct jit_state *state = malloc(sizeof(struct jit_state));
+    if (!state)
+        return NULL;
     assert(state);
+
     state->offset = 0;
     state->size = size;
     state->buf = mmap(0, size, PROT_READ | PROT_WRITE | PROT_EXEC,
@@ -2340,13 +2343,32 @@ struct jit_state *jit_state_init(size_t size)
 #endif
                       ,
                       -1, 0);
-    state->n_blocks = 0;
+    if (state->buf == MAP_FAILED) {
+        free(state);
+        return NULL;
+    }
     assert(state->buf != MAP_FAILED);
+
+    state->n_blocks = 0;
     set_reset(&state->set);
     reset_reg();
     prepare_translate(state);
+
     state->offset_map = calloc(MAX_BLOCKS, sizeof(struct offset_map));
+    if (!state->offset_map) {
+        munmap(state->buf, state->size);
+        free(state);
+        return NULL;
+    }
+
     state->jumps = calloc(MAX_JUMPS, sizeof(struct jump));
+    if (!state->jumps) {
+        free(state->offset_map);
+        munmap(state->buf, state->size);
+        free(state);
+        return NULL;
+    }
+
     return state;
 }
 

src/main.c

@@ -190,7 +190,7 @@ static bool parse_args(int argc, char **args)
     return true;
 }
 
-static void dump_test_signature(const char *prog_name)
+static void dump_test_signature(const char UNUSED *prog_name)
 {
     elf_t *elf = elf_new();
     assert(elf && elf_open(elf, prog_name));

src/riscv.c

@@ -510,6 +510,8 @@ riscv_t *rv_create(riscv_user_t rv_attr)
     assert(rv_attr);
 
     riscv_t *rv = calloc(1, sizeof(riscv_t));
+    if (!rv)
+        return NULL;
     assert(rv);
 
 #if RV32_HAS(SYSTEM) && !RV32_HAS(ELF_LOADER)
@@ -578,7 +580,7 @@ riscv_t *rv_create(riscv_user_t rv_attr)
     assert(elf_load(elf, attr->mem));
 
     /* set the entry pc */
-    const struct Elf32_Ehdr *hdr = get_elf_header(elf);
+    const struct Elf32_Ehdr UNUSED *hdr = get_elf_header(elf);
     assert(rv_set_pc(rv, hdr->e_entry));
 
     elf_delete(elf);
@@ -724,11 +726,22 @@ riscv_t *rv_create(riscv_user_t rv_attr)
 #else
     INIT_LIST_HEAD(&rv->block_list);
     rv->jit_state = jit_state_init(CODE_CACHE_SIZE);
+    if (!rv->jit_state) {
+        rv_log_fatal("Failed to initialize JIT state");
+        goto fail_jit_state;
+    }
     rv->block_cache = cache_create(BLOCK_MAP_CAPACITY_BITS);
-    assert(rv->block_cache);
+    if (!rv->block_cache) {
+        rv_log_fatal("Failed to create block cache");
+        goto fail_block_cache;
+    }
 #if RV32_HAS(T2C)
     rv->quit = false;
     rv->jit_cache = jit_cache_init();
+    if (!rv->jit_cache) {
+        rv_log_fatal("Failed to initialize JIT cache");
+        goto fail_jit_cache;
+    }
     /* prepare wait queue. */
     pthread_mutex_init(&rv->wait_queue_lock, NULL);
     pthread_mutex_init(&rv->cache_lock, NULL);
@@ -739,6 +752,22 @@ riscv_t *rv_create(riscv_user_t rv_attr)
 #endif
 
     return rv;
+
+#if RV32_HAS(JIT)
+#if RV32_HAS(T2C)
+fail_jit_cache:
+    cache_free(rv->block_cache);
+#endif
+fail_block_cache:
+    jit_state_exit(rv->jit_state);
+fail_jit_state:
+    mpool_destroy(rv->block_ir_mp);
+    mpool_destroy(rv->block_mp);
+    map_delete(attr->fd_map);
+    memory_delete(attr->mem);
+    free(rv);
+    return NULL;
+#endif
 }
 
 #if !RV32_HAS(SYSTEM) || (RV32_HAS(SYSTEM) && RV32_HAS(ELF_LOADER))
@@ -753,7 +782,7 @@ static void rv_run_and_trace(riscv_t *rv)
     assert(attr && attr->data.user.elf_program);
     attr->cycle_per_step = 1;
 
-    const char *prog_name = attr->data.user.elf_program;
+    const char UNUSED *prog_name = attr->data.user.elf_program;
     elf_t *elf = elf_new();
     assert(elf && elf_open(elf, prog_name));
 

src/syscall.c

@@ -362,22 +362,43 @@ static void syscall_open(riscv_t *rv)
     uint32_t flags = rv_get_reg(rv, rv_reg_a1);
     uint32_t mode = rv_get_reg(rv, rv_reg_a2);
 
-    /* read name from runtime memory */
-    const size_t name_len = strlen((char *) attr->mem->mem_base + name);
+    /* read name from runtime memory with bounds checking */
+    if (name >= attr->mem->mem_size) {
+        rv_set_reg(rv, rv_reg_a0, -1);
+        return;
+    }
+
+    /* Calculate safe maximum length to prevent reading beyond memory */
+    const size_t max_len = attr->mem->mem_size - name;
+    const char *name_ptr = (char *) attr->mem->mem_base + name;
+
+    /* Use strnlen to safely find string length within bounds */
+    const size_t name_len = strnlen(name_ptr, max_len);
+    if (name_len == max_len) {
+        /* No null terminator found within bounds */
+        rv_set_reg(rv, rv_reg_a0, -1);
+        return;
+    }
     char *name_str = malloc(name_len + 1);
+    if (!name_str) {
+        rv_set_reg(rv, rv_reg_a0, -1);
+        return;
+    }
     assert(name_str);
     name_str[name_len] = '\0';
     memory_read(attr->mem, (uint8_t *) name_str, name, name_len);
 
     /* open the file */
     const char *mode_str = get_mode_str(flags, mode);
     if (!mode_str) {
+        free(name_str);
         rv_set_reg(rv, rv_reg_a0, -1);
         return;
     }
 
     FILE *handle = fopen(name_str, mode_str);
     if (!handle) {
+        free(name_str);
         rv_set_reg(rv, rv_reg_a0, -1);
         return;
     }