Custom code generation for ret imm (C2)

gen/x86_table.js

@@ -246,7 +246,7 @@ const encodings = [
     { opcode: 0xF2AF, block_boundary: 1, is_string: 1, os: 1, },
     { opcode: 0xF3AF, block_boundary: 1, is_string: 1, os: 1, },
 
-    { opcode: 0xC2, block_boundary: 1, no_next_instruction: 1, os: 1, imm16: 1, skip: 1, }, // ret
+    { opcode: 0xC2, custom: 1, block_boundary: 1, no_next_instruction: 1, os: 1, imm16: 1, skip: 1, }, // ret
     { opcode: 0xC3, custom: 1, block_boundary: 1, no_next_instruction: 1, os: 1, skip: 1, },
 
     { opcode: 0xC4, block_boundary: 1, os: 1, e: 1, skip: 1, }, // les

src/rust/codegen.rs

@@ -880,6 +880,27 @@ pub fn gen_pop32s(ctx: &mut JitContext) {
     }
 }
 
+pub fn gen_adjust_stack_reg(ctx: &mut JitContext, offset: u32) {
+    if ctx.cpu.ssize_32() {
+        ctx.builder
+            .instruction_body
+            .const_i32(global_pointers::get_reg32_offset(regs::ESP) as i32);
+        gen_get_reg32(ctx.builder, regs::ESP);
+        ctx.builder.instruction_body.const_i32(offset as i32);
+        ctx.builder.instruction_body.add_i32();
+        ctx.builder.instruction_body.store_aligned_i32(0);
+    }
+    else {
+        ctx.builder
+            .instruction_body
+            .const_i32(global_pointers::get_reg16_offset(regs::SP) as i32);
+        gen_get_reg16(ctx.builder, regs::SP);
+        ctx.builder.instruction_body.const_i32(offset as i32);
+        ctx.builder.instruction_body.add_i32();
+        ctx.builder.instruction_body.store_aligned_u16(0);
+    }
+}
+
 pub fn gen_task_switch_test(ctx: &mut JitContext) {
     // generate if(cr[0] & (CR0_EM | CR0_TS)) { task_switch_test_void(); return; }
     let cr0_offset = global_pointers::get_creg_offset(0);

src/rust/jit_instructions.rs

@@ -1802,6 +1802,31 @@ pub fn instr32_E9_jit(ctx: &mut JitContext, imm: u32) {
     codegen::gen_relative_jump(ctx.builder, imm as i32);
 }
 
+pub fn instr16_C2_jit(ctx: &mut JitContext, imm16: u32) {
+    let cs_addr = global_pointers::get_seg_offset(CS);
+    ctx.builder
+        .instruction_body
+        .const_i32(global_pointers::INSTRUCTION_POINTER as i32);
+    ctx.builder.instruction_body.load_aligned_i32(cs_addr);
+    codegen::gen_pop16(ctx);
+    codegen::gen_adjust_stack_reg(ctx, imm16);
+    ctx.builder.instruction_body.add_i32();
+    ctx.builder.instruction_body.store_aligned_i32(0);
+}
+
+pub fn instr32_C2_jit(ctx: &mut JitContext, imm16: u32) {
+    ctx.builder
+        .instruction_body
+        .const_i32(global_pointers::INSTRUCTION_POINTER as i32);
+    ctx.builder
+        .instruction_body
+        .load_aligned_i32(global_pointers::get_seg_offset(CS));
+    codegen::gen_pop32s(ctx);
+    codegen::gen_adjust_stack_reg(ctx, imm16);
+    ctx.builder.instruction_body.add_i32();
+    ctx.builder.instruction_body.store_aligned_i32(0);
+}
+
 pub fn instr16_C3_jit(ctx: &mut JitContext) {
     let cs_addr = global_pointers::get_seg_offset(CS);
 

tests/nasm/ret-imm.asm

@@ -0,0 +1,15 @@
+global _start
+
+%include "header.inc"
+
+    jmp start
+foo:
+    mov eax, esp
+    ret 123
+
+start:
+    call foo
+    mov dword [eax], 0 ; clear the address pushed by the call instruction
+
+
+%include "footer.inc"