more fixes

3 files changed, 48 insertions, 20 deletions

diff --git a/lib/x86/arch-dis.h b/lib/x86/arch-dis.h
index d121549..a76884f 100644
--- a/lib/x86/arch-dis.h
+++ b/lib/x86/arch-dis.h
@@ -9,6 +9,7 @@
 
 struct arch_pcrel_info {
     int reg;
+    bool is_jump;
 };
 
 struct arch_dis_ctx {};
diff --git a/lib/x86/arch-transform-dis.inc.h b/lib/x86/arch-transform-dis.inc.h
index 67390c9..bb18d7d 100644
--- a/lib/x86/arch-transform-dis.inc.h
+++ b/lib/x86/arch-transform-dis.inc.h
@@ -24,17 +24,32 @@ UNUSED
 static void transform_dis_pcrel(struct transform_dis_ctx *ctx, uint64_t dpc,
                                 struct arch_pcrel_info info) {
     /* push %reg; mov $dpc, %reg; <orig but with reg instead>; pop %reg
+     * -or, if jump-
+     * push %reg; mov $dpc, %reg; mov %reg, -8(%rsp); pop %reg;
+     *      <orig but with -0x10(%rsp)>
      * reg is rcx, or rax if the instruction might be using rcx.
      * Max size: 11 + orig + 1
      * Minimum size is 6 bytes, so there could be at most 1 in a patch area. */
-    bool rax = info.reg == 1;
     void *code = *ctx->rewritten_ptr_ptr;
-    push_mov_head(&code, dpc, rax);
-    ctx->write_newop_here = code;
-    code += ctx->base.op_size;
-    push_mov_tail(&code, rax);
+    if (info.is_jump) {
+        push_mov_head(&code, dpc, true);
+        memcpy(code, ((uint8_t[]) {0x48, 0x89, 0x44, 0x24, 0xf8}), 5);
+        code += 5;
+        push_mov_tail(&code, true);
+        ctx->write_newop_here = code;
+        code += ctx->base.op_size - 2;
+        ctx->base.newval[0] = 4; /* esp */
+        ctx->base.newval[1] = -0x10;
+    } else {
+        bool rax = info.reg == 1;
+        push_mov_head(&code, dpc, rax);
+        ctx->write_newop_here = code;
+        code += ctx->base.op_size - 4 /* see dis-main.inc.h */;
+        push_mov_tail(&code, rax);
+        ctx->base.newval[0] = rax ? 0 /* rcx */ : 1 /* rax */;
+        ctx->base.newval[1] = 0;
+    }
     *ctx->rewritten_ptr_ptr = code;
-    ctx->base.newval[0] = rax ? 0 : 1;
     ctx->base.modify = true;
 }
 
diff --git a/lib/x86/dis-main.inc.h b/lib/x86/dis-main.inc.h
index f948026..4be45d8 100644
--- a/lib/x86/dis-main.inc.h
+++ b/lib/x86/dis-main.inc.h
@@ -131,13 +131,11 @@ restart:;
             bits = I_MODA;
         } else if (byte1 == 0xff) {
             uint8_t modrm = *ptr;
-            if (modrm >> 6 == 3) {
-                int subop = modrm >> 3 & 7;
-                if (subop == 4 || subop == 5) /* JMP */
-                    bits = I_JMP | I_MODA;
-                else
-                    bits = I_MODA;
-            }
+            int subop = modrm >> 3 & 7;
+            if (subop == 4 || subop == 5) /* JMP */
+                bits = I_JMP | I_MODA;
+            else
+                bits = I_MODA;
         } else {
             __builtin_abort();
         }
@@ -287,22 +285,36 @@ got_bits: UNUSED
         int32_t disp = *(int32_t *) (orig + modrm_off + 1);
         /* unlike ARM, we can always switch to non-pcrel without making the
          * instruction from scratch, so we don't have 'reg' and 'lm' */
-        struct arch_pcrel_info info = {modrm >> 3 & 7};
+        struct arch_pcrel_info info = {
+            .reg = modrm >> 3 & 7,
+            .is_jump = !!(bits & I_JMP),
+        };
         P(pcrel)(ctx, ctx->base.pc + ctx->base.op_size + disp, info);
         if (DIS_MAY_MODIFY && ctx->base.modify) {
             uint8_t *new_op = ctx->base.newop;
             memcpy(new_op, orig, ctx->base.op_size);
-            /* newval[0] should be the new register, which should be one that
-             * fits in r/m directly since that's all I need;
-             * displacement is removed */
+            /* newval[0] should be the new register;
+             * newval[1] should be the new displacement */
+            int new_reg = ctx->base.newval[0];
+            uint32_t new_disp = ctx->base.newval[1];
             uint8_t *new_modrm_ptr = new_op + modrm_off;
 
+            int new_disp_size = new_disp ? 1 : 0;
+
             *new_modrm_ptr = (*new_modrm_ptr & ~0xc7) |
-                             0 << 6 |
+                             (new_disp_size ? 1 : 0) << 6 |
                              ctx->base.newval[0];
-            memmove(new_modrm_ptr + 1, new_modrm_ptr + 5,
+            uint8_t *memspec_end = new_modrm_ptr + 1;
+            if (new_reg == 4) {
+                /* rsp - need SIB */
+                *memspec_end++ = 0x24;
+            }
+            if (new_disp_size)
+                *memspec_end++ = new_disp;
+
+            memmove(memspec_end, new_modrm_ptr + 5,
                     ctx->base.op_size - modrm_off - 1);
-            ctx->base.newop_size -= 4;
+            ctx->base.newop_size -= 5 - (memspec_end - new_modrm_ptr);
         }
 #endif
     } else if ((bits & I_TYPE_MASK) == I_JMP) {