Merge pull request #617 from sysprog21/memory-safety

Fix critical memory safety issues

Author: jserv

src/cache.c

@@ -147,6 +147,10 @@ static inline void hlist_del_init(struct hlist_node *n)
 
 cache_t *cache_create(uint32_t size_bits)
 {
+    /* Prevent integer overflow in 1 << size_bits */
+    if (size_bits >= 32)
+        return NULL;
+
     cache_t *cache = malloc(sizeof(cache_t));
     if (!cache)
         return NULL;
@@ -160,16 +164,23 @@ cache_t *cache_create(uint32_t size_bits)
     cache->ghost_list_size = 0;
     cache->capacity = cache_size;
 
-    cache->map.ht_list_head = malloc(cache_size * sizeof(struct hlist_head));
-    if (!cache->map.ht_list_head) {
-        free(cache);
-        return NULL;
-    }
+    /* Check for overflow in size calculation */
+    size_t alloc_size = cache_size * sizeof(struct hlist_head);
+    if (alloc_size / sizeof(struct hlist_head) != cache_size)
+        goto fail_cache;
+
+    cache->map.ht_list_head = malloc(alloc_size);
+    if (!cache->map.ht_list_head)
+        goto fail_cache;
 
     for (uint32_t i = 0; i < cache_size; i++)
         INIT_HLIST_HEAD(&cache->map.ht_list_head[i]);
 
     return cache;
+
+fail_cache:
+    free(cache);
+    return NULL;
 }
 
 void *cache_get(const cache_t *cache, uint32_t key, bool update)
@@ -285,6 +296,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/elf.c

@@ -274,14 +274,14 @@ bool elf_load(elf_t *e, memory_t *mem)
 
         /* memcpy required range */
         const int to_copy = min(phdr->p_memsz, phdr->p_filesz);
-        if (to_copy)
-            memory_write(mem, phdr->p_vaddr, e->raw_data + phdr->p_offset,
-                         to_copy);
+        if (to_copy && !memory_write(mem, phdr->p_vaddr,
+                                     e->raw_data + phdr->p_offset, to_copy))
+            return false;
 
         /* zero fill required range */
         const int to_zero = max(phdr->p_memsz, phdr->p_filesz) - to_copy;
-        if (to_zero)
-            memory_fill(mem, phdr->p_vaddr + to_copy, to_zero, 0);
+        if (to_zero && !memory_fill(mem, phdr->p_vaddr + to_copy, to_zero, 0))
+            return false;
     }
 
     return true;

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)
     /*
@@ -1078,6 +1102,12 @@ void rv_step(void *arg)
          */
         block_t *block = block_find_or_translate(rv);
         /* by now, a block should be available */
+        if (unlikely(!block)) {
+            rv_log_fatal("Failed to allocate or translate block at PC=0x%08x",
+                         rv->PC);
+            rv->halt = true;
+            return;
+        }
         assert(block);
 
 #if RV32_HAS(JIT) && RV32_HAS(SYSTEM)
@@ -1129,6 +1159,11 @@ 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 - reset compiled flag to allow retry later */
+                block->compiled = false;
+                continue;
+            }
             entry->block = block;
             pthread_mutex_lock(&rv->wait_queue_lock);
             list_add(&entry->list, &rv->wait_queue);
@@ -1378,16 +1413,38 @@ void ecall_handler(riscv_t *rv)
 void memset_handler(riscv_t *rv)
 {
     memory_t *m = PRIV(rv)->mem;
-    memset((char *) m->mem_base + rv->X[rv_reg_a0], rv->X[rv_reg_a1],
-           rv->X[rv_reg_a2]);
+    uint32_t dest = rv->X[rv_reg_a0];
+    uint32_t value = rv->X[rv_reg_a1];
+    uint32_t count = rv->X[rv_reg_a2];
+
+    /* Bounds checking to prevent buffer overflow */
+    if (dest >= m->mem_size || count > m->mem_size - dest) {
+        SET_CAUSE_AND_TVAL_THEN_TRAP(rv, STORE_MISALIGNED, dest);
+        return;
+    }
+
+    memset((char *) m->mem_base + dest, value, count);
     rv->PC = rv->X[rv_reg_ra] & ~1U;
 }
 
 void memcpy_handler(riscv_t *rv)
 {
     memory_t *m = PRIV(rv)->mem;
-    memcpy((char *) m->mem_base + rv->X[rv_reg_a0],
-           (char *) m->mem_base + rv->X[rv_reg_a1], rv->X[rv_reg_a2]);
+    uint32_t dest = rv->X[rv_reg_a0];
+    uint32_t src = rv->X[rv_reg_a1];
+    uint32_t count = rv->X[rv_reg_a2];
+
+    /* Bounds checking to prevent buffer overflow */
+    if (dest >= m->mem_size || count > m->mem_size - dest) {
+        SET_CAUSE_AND_TVAL_THEN_TRAP(rv, STORE_MISALIGNED, dest);
+        return;
+    }
+    if (src >= m->mem_size || count > m->mem_size - src) {
+        SET_CAUSE_AND_TVAL_THEN_TRAP(rv, LOAD_MISALIGNED, src);
+        return;
+    }
+
+    memcpy((char *) m->mem_base + dest, (char *) m->mem_base + src, count);
     rv->PC = rv->X[rv_reg_ra] & ~1U;
 }
 

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/io.h

@@ -36,12 +36,16 @@ uint8_t memory_read_b(uint32_t addr);
 void memory_read(const memory_t *m, uint8_t *dst, uint32_t addr, uint32_t size);
 
 /* write a length of data to memory */
-static inline void memory_write(memory_t *m,
+static inline bool memory_write(memory_t *m,
                                 uint32_t addr,
                                 const uint8_t *src,
                                 uint32_t size)
 {
+    /* Bounds checking to prevent buffer overflow */
+    if (addr >= m->mem_size || size > m->mem_size - addr)
+        return false;
     memcpy(m->mem_base + addr, src, size);
+    return true;
 }
 
 /* write a word to memory */
@@ -54,10 +58,14 @@ void memory_write_s(uint32_t addr, const uint8_t *src);
 void memory_write_b(uint32_t addr, const uint8_t *src);
 
 /* write a length of certain value to memory */
-static inline void memory_fill(memory_t *m,
+static inline bool memory_fill(memory_t *m,
                                uint32_t addr,
                                uint32_t size,
                                uint8_t val)
 {
+    /* Bounds checking to prevent buffer overflow */
+    if (addr >= m->mem_size || size > m->mem_size - addr)
+        return false;
     memset(m->mem_base + addr, val, size);
+    return true;
 }