@@ -273,6 +273,8 @@ HASH_FUNC_IMPL(map_hash, BLOCK_MAP_CAPACITY_BITS, 1 << BLOCK_MAP_CAPACITY_BITS)
273273static block_t *block_alloc(riscv_t *rv)
274274{
275275 block_t *block = mpool_alloc(rv->block_mp);
276+ if (unlikely(!block))
277+ return NULL;
276278 assert(block);
277279 block->n_insn = 0;
278280#if RV32_HAS(JIT)
@@ -619,13 +621,15 @@ FORCE_INLINE bool insn_is_indirect_branch(uint8_t opcode)
619621 }
620622}
621623
622- static void block_translate(riscv_t *rv, block_t *block)
624+ static bool block_translate(riscv_t *rv, block_t *block)
623625{
624626retranslate:
625627 block->pc_start = block->pc_end = rv->PC;
626628
627629 rv_insn_t *prev_ir = NULL;
628630 rv_insn_t *ir = mpool_calloc(rv->block_ir_mp);
631+ if (unlikely(!ir))
632+ return false;
629633 block->ir_head = ir;
630634
631635 /* translate the basic block */
@@ -665,6 +669,8 @@ static void block_translate(riscv_t *rv, block_t *block)
665669 if (insn_is_branch(ir->opcode)) {
666670 if (insn_is_indirect_branch(ir->opcode)) {
667671 ir->branch_table = calloc(1, sizeof(branch_history_table_t));
672+ if (unlikely(!ir->branch_table))
673+ return false;
668674 assert(ir->branch_table);
669675 memset(ir->branch_table->PC, -1,
670676 sizeof(uint32_t) * HISTORY_SIZE);
@@ -673,36 +679,44 @@ static void block_translate(riscv_t *rv, block_t *block)
673679 }
674680
675681 ir = mpool_calloc(rv->block_ir_mp);
682+ if (unlikely(!ir))
683+ return false;
676684 }
677685
678686 assert(prev_ir);
679687 block->ir_tail = prev_ir;
680688 block->ir_tail->next = NULL;
689+ return true;
681690}
682691
683692#if RV32_HAS(MOP_FUSION)
684- #define COMBINE_MEM_OPS(RW) \
685- next_ir = ir->next; \
686- count = 1; \
687- while (1) { \
688- if (next_ir->opcode != IIF(RW)(rv_insn_lw, rv_insn_sw)) \
689- break; \
690- count++; \
691- if (!next_ir->next) \
692- break; \
693- next_ir = next_ir->next; \
694- } \
695- if (count > 1) { \
696- ir->opcode = IIF(RW)(rv_insn_fuse4, rv_insn_fuse3); \
697- ir->fuse = malloc(count * sizeof(opcode_fuse_t)); \
698- assert(ir->fuse); \
699- ir->imm2 = count; \
700- memcpy(ir->fuse, ir, sizeof(opcode_fuse_t)); \
701- ir->impl = dispatch_table[ir->opcode]; \
702- next_ir = ir->next; \
703- for (int j = 1; j < count; j++, next_ir = next_ir->next) \
704- memcpy(ir->fuse + j, next_ir, sizeof(opcode_fuse_t)); \
705- remove_next_nth_ir(rv, ir, block, count - 1); \
693+ #define COMBINE_MEM_OPS(RW) \
694+ next_ir = ir->next; \
695+ count = 1; \
696+ while (1) { \
697+ if (next_ir->opcode != IIF(RW)(rv_insn_lw, rv_insn_sw)) \
698+ break; \
699+ count++; \
700+ if (!next_ir->next) \
701+ break; \
702+ next_ir = next_ir->next; \
703+ } \
704+ if (count > 1) { \
705+ ir->opcode = IIF(RW)(rv_insn_fuse4, rv_insn_fuse3); \
706+ ir->fuse = malloc(count * sizeof(opcode_fuse_t)); \
707+ if (unlikely(!ir->fuse)) { \
708+ ir->opcode = IIF(RW)(rv_insn_lw, rv_insn_sw); \
709+ count = 1; /* Degrade to non-fused operation */ \
710+ } else { \
711+ assert(ir->fuse); \
712+ ir->imm2 = count; \
713+ memcpy(ir->fuse, ir, sizeof(opcode_fuse_t)); \
714+ ir->impl = dispatch_table[ir->opcode]; \
715+ next_ir = ir->next; \
716+ for (int j = 1; j < count; j++, next_ir = next_ir->next) \
717+ memcpy(ir->fuse + j, next_ir, sizeof(opcode_fuse_t)); \
718+ remove_next_nth_ir(rv, ir, block, count - 1); \
719+ } \
706720 }
707721
708722static inline void remove_next_nth_ir(const riscv_t *rv,
@@ -762,16 +776,20 @@ static void match_pattern(riscv_t *rv, block_t *block)
762776 next_ir = next_ir->next;
763777 }
764778 if (count > 1) {
765- ir->opcode = rv_insn_fuse1;
766779 ir->fuse = malloc(count * sizeof(opcode_fuse_t));
767- assert(ir->fuse);
768- ir->imm2 = count;
769- memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));
770- ir->impl = dispatch_table[ir->opcode];
771- next_ir = ir->next;
772- for (int j = 1; j < count; j++, next_ir = next_ir->next)
773- memcpy(ir->fuse + j, next_ir, sizeof(opcode_fuse_t));
774- remove_next_nth_ir(rv, ir, block, count - 1);
780+ if (likely(ir->fuse)) {
781+ ir->opcode = rv_insn_fuse1;
782+ assert(ir->fuse);
783+ ir->imm2 = count;
784+ memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));
785+ ir->impl = dispatch_table[ir->opcode];
786+ next_ir = ir->next;
787+ for (int j = 1; j < count; j++, next_ir = next_ir->next)
788+ memcpy(ir->fuse + j, next_ir,
789+ sizeof(opcode_fuse_t));
790+ remove_next_nth_ir(rv, ir, block, count - 1);
791+ }
792+ /* If malloc failed, degrade gracefully to non-fused ops */
775793 }
776794 break;
777795 }
@@ -803,15 +821,18 @@ static void match_pattern(riscv_t *rv, block_t *block)
803821 }
804822 if (count > 1) {
805823 ir->fuse = malloc(count * sizeof(opcode_fuse_t));
806- assert(ir->fuse);
807- memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));
808- ir->opcode = rv_insn_fuse5;
809- ir->imm2 = count;
810- ir->impl = dispatch_table[ir->opcode];
811- next_ir = ir->next;
812- for (int j = 1; j < count; j++, next_ir = next_ir->next)
813- memcpy(ir->fuse + j, next_ir, sizeof(opcode_fuse_t));
814- remove_next_nth_ir(rv, ir, block, count - 1);
824+ if (likely(ir->fuse)) {
825+ assert(ir->fuse);
826+ memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));
827+ ir->opcode = rv_insn_fuse5;
828+ ir->imm2 = count;
829+ ir->impl = dispatch_table[ir->opcode];
830+ next_ir = ir->next;
831+ for (int j = 1; j < count; j++, next_ir = next_ir->next)
832+ memcpy(ir->fuse + j, next_ir, sizeof(opcode_fuse_t));
833+ remove_next_nth_ir(rv, ir, block, count - 1);
834+ }
835+ /* If malloc failed, degrade gracefully to non-fused ops */
815836 }
816837 break;
817838 }
@@ -881,8 +902,11 @@ static block_t *block_find_or_translate(riscv_t *rv)
881902#endif
882903 /* allocate a new block */
883904 next_blk = block_alloc(rv);
905+ if (unlikely(!next_blk))
906+ return NULL;
884907
885- block_translate(rv, next_blk);
908+ if (unlikely(!block_translate(rv, next_blk)))
909+ return NULL;
886910
887911#if RV32_HAS(JIT) && RV32_HAS(SYSTEM)
888912 /*
@@ -1078,6 +1102,12 @@ void rv_step(void *arg)
10781102 */
10791103 block_t *block = block_find_or_translate(rv);
10801104 /* by now, a block should be available */
1105+ if (unlikely(!block)) {
1106+ rv_log_fatal("Failed to allocate or translate block at PC=0x%08x",
1107+ rv->PC);
1108+ rv->halt = true;
1109+ return;
1110+ }
10811111 assert(block);
10821112
10831113#if RV32_HAS(JIT) && RV32_HAS(SYSTEM)
@@ -1129,6 +1159,11 @@ void rv_step(void *arg)
11291159 else if (!block->compiled && block->n_invoke >= THRESHOLD) {
11301160 block->compiled = true;
11311161 queue_entry_t *entry = malloc(sizeof(queue_entry_t));
1162+ if (unlikely(!entry)) {
1163+ /* Malloc failed - reset compiled flag to allow retry later */
1164+ block->compiled = false;
1165+ continue;
1166+ }
11321167 entry->block = block;
11331168 pthread_mutex_lock(&rv->wait_queue_lock);
11341169 list_add(&entry->list, &rv->wait_queue);
@@ -1378,16 +1413,38 @@ void ecall_handler(riscv_t *rv)
13781413void memset_handler(riscv_t *rv)
13791414{
13801415 memory_t *m = PRIV(rv)->mem;
1381- memset((char *) m->mem_base + rv->X[rv_reg_a0], rv->X[rv_reg_a1],
1382- rv->X[rv_reg_a2]);
1416+ uint32_t dest = rv->X[rv_reg_a0];
1417+ uint32_t value = rv->X[rv_reg_a1];
1418+ uint32_t count = rv->X[rv_reg_a2];
1419+
1420+ /* Bounds checking to prevent buffer overflow */
1421+ if (dest >= m->mem_size || count > m->mem_size - dest) {
1422+ SET_CAUSE_AND_TVAL_THEN_TRAP(rv, STORE_MISALIGNED, dest);
1423+ return;
1424+ }
1425+
1426+ memset((char *) m->mem_base + dest, value, count);
13831427 rv->PC = rv->X[rv_reg_ra] & ~1U;
13841428}
13851429
13861430void memcpy_handler(riscv_t *rv)
13871431{
13881432 memory_t *m = PRIV(rv)->mem;
1389- memcpy((char *) m->mem_base + rv->X[rv_reg_a0],
1390- (char *) m->mem_base + rv->X[rv_reg_a1], rv->X[rv_reg_a2]);
1433+ uint32_t dest = rv->X[rv_reg_a0];
1434+ uint32_t src = rv->X[rv_reg_a1];
1435+ uint32_t count = rv->X[rv_reg_a2];
1436+
1437+ /* Bounds checking to prevent buffer overflow */
1438+ if (dest >= m->mem_size || count > m->mem_size - dest) {
1439+ SET_CAUSE_AND_TVAL_THEN_TRAP(rv, STORE_MISALIGNED, dest);
1440+ return;
1441+ }
1442+ if (src >= m->mem_size || count > m->mem_size - src) {
1443+ SET_CAUSE_AND_TVAL_THEN_TRAP(rv, LOAD_MISALIGNED, src);
1444+ return;
1445+ }
1446+
1447+ memcpy((char *) m->mem_base + dest, (char *) m->mem_base + src, count);
13911448 rv->PC = rv->X[rv_reg_ra] & ~1U;
13921449}
13931450
0 commit comments