Trampoline fixes.

The transformed code was incorrect because it assumed the pointer it was
writing to was where the code would execute, but it was actually
'rewritten_temp'.  Changed transform_dis_main to take a pc_trampoline
pointer, which also helps the test harness.  However, this means that it
has to be called after the trampoline has been allocated, while before
the trampoline allocation depended on the generated size; this change
doesn't bother to use two passes or anything, but just allocates a new
code buffer if the maximum possible size isn't available - not the end
of the world, since trampoline_ptr will still only be increased by the
actual size before the next hook in the series (if any).

4 files changed, 32 insertions, 16 deletions

diff --git a/lib/x86/arch-dis.h b/lib/x86/arch-dis.h
index 6447f38..d121549 100644
--- a/lib/x86/arch-dis.h
+++ b/lib/x86/arch-dis.h
@@ -1,6 +1,11 @@
 #pragma once
 #define MIN_INSN_SIZE 1
-#define TD_MAX_REWRITTEN_SIZE 100 /* XXX */
+/* min([18 * 3,
+ *      4 + 18 + 15 + 18,
+ *      6 + 12])
+ * See transform_dis_* for size figures.  Technically unsafe, since we don't
+ * check for overlong x86 instructions. */
+#define TD_MAX_REWRITTEN_SIZE 55
 
 struct arch_pcrel_info {
     int reg;
@@ -8,3 +13,6 @@ struct arch_pcrel_info {
 
 struct arch_dis_ctx {};
 static inline void arch_dis_ctx_init(UNUSED struct arch_dis_ctx *ctx) {}
+static inline int arch_code_alignment(UNUSED struct arch_dis_ctx ctx) {
+   return 4;
+}
diff --git a/lib/x86/arch-transform-dis.inc.h b/lib/x86/arch-transform-dis.inc.h
index 1e5d872..6ec6c16 100644
--- a/lib/x86/arch-transform-dis.inc.h
+++ b/lib/x86/arch-transform-dis.inc.h
@@ -23,8 +23,10 @@ static inline void push_mov_tail(void **code, bool rax) {
 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 */
-    /* reg is rcx, or rax if the instruction might be using rcx. */
+    /* push %reg; mov $dpc, %reg; <orig but with reg instead>; pop %reg
+     * 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);
@@ -32,7 +34,7 @@ static void transform_dis_pcrel(struct transform_dis_ctx *ctx, uint64_t dpc,
     code += ctx->base.op_size;
     push_mov_tail(&code, rax);
     *ctx->rewritten_ptr_ptr = code;
-    ctx->base.newop[0] = rax ? 0 : 1;
+    ctx->base.newval[0] = rax ? 0 : 1;
     ctx->base.modify = true;
 }
 
@@ -41,7 +43,9 @@ static void transform_dis_branch(struct transform_dis_ctx *ctx, uint_tptr dpc,
     if (dpc >= ctx->pc_patch_start && dpc < ctx->pc_patch_end) {
         if (dpc == ctx->base.pc + ctx->base.op_size && (cc & CC_CALL)) {
             /* Probably a poor man's PC-rel - 'call .; pop %some'.
-             * Push the original address. */
+             * Push the original address.
+             * Max size: orig + 1 + 11 + 5 + 1
+             * Minimum call size is 4 bytes; at most 2. */
             void *code = *ctx->rewritten_ptr_ptr;
             ctx->write_newop_here = NULL;
 
@@ -72,15 +76,18 @@ static void transform_dis_branch(struct transform_dis_ctx *ctx, uint_tptr dpc,
     code += ctx->base.op_size;
 
     struct arch_dis_ctx arch;
-    uintptr_t source = (uintptr_t) code + 2;
+    uintptr_t source = ctx->pc_trampoline + 2;
     int size = jump_patch_size(source, dpc, arch, true);
-    /* if not taken, jmp past the big jump - this is a bit suboptimal but not that bad */
+    /* If not taken, jmp past the big jump - this is a bit suboptimal but not
+     * that bad.
+     * Max size: orig + 2 + 14
+     * Minimum jump size is 2 bytes; at most 3. */
     op8(&code, 0xeb);
     op8(&code, size);
     make_jump_patch(&code, source, dpc, arch);
 
     *ctx->rewritten_ptr_ptr = code;
-    ctx->base.newop[0] = 2;
+    ctx->base.newval[0] = 2;
     ctx->base.modify = true;
 
     if (!cc)
diff --git a/lib/x86/dis-main.inc.h b/lib/x86/dis-main.inc.h
index 45a0947..f948026 100644
--- a/lib/x86/dis-main.inc.h
+++ b/lib/x86/dis-main.inc.h
@@ -97,7 +97,7 @@ static const uint8_t _0f_bits[] = {
 /*60*/ REP16(I_MODA),
 /*70*/ I_MODA, I_MOD|I_8, I_MOD|I_8, I_MOD|I_8, I_MODA, I_MODA, I_MODA, 0,
      /*78*/ I_MODA, I_MODA, I_BAD, I_BAD, REP4(I_MODA),
-/*80*/ REP16(I_z),
+/*80*/ REP16(I_z|I_JIMM),
 /*90*/ REP16(I_MODA),
 /*Ax*/ 0, 0, 0, 0, 0, 0, I_BAD, I_BAD,
      /*A8*/ 0, 0, 0, I_MODA, I_MODA|I_8, I_MODA, I_MODA, I_MODA,
@@ -279,6 +279,7 @@ got_bits: UNUSED
                 case 1: *(int8_t *)  new_imm_ptr = new_imm; break;
                 case 2: *(int16_t *) new_imm_ptr = new_imm; break;
                 case 4: *(int32_t *) new_imm_ptr = new_imm; break;
+                default: __builtin_abort();
             }
         }
 #ifdef TARGET_x86_64
diff --git a/lib/x86/jump-patch.h b/lib/x86/jump-patch.h
index 4c0172d..8cd7d6f 100644
--- a/lib/x86/jump-patch.h
+++ b/lib/x86/jump-patch.h
@@ -1,23 +1,23 @@
 #pragma once
-#define MAX_JUMP_PATCH_SIZE 5
+#define MAX_JUMP_PATCH_SIZE 14
 #include "dis.h"
 
-static inline int jump_patch_size(uintptr_t pc, uintptr_t dpc,
+static inline int jump_patch_size(uint_tptr pc, uint_tptr dpc,
                                   UNUSED struct arch_dis_ctx arch,
                                   bool force) {
-    uintptr_t diff = pc - (dpc + 5);
+    uint_tptr diff = pc - (dpc + 5);
     /* fits in 32? */
-    if (diff == (uintptr_t) (int32_t) diff)
+    if (diff == (uint_tptr) (int32_t) diff)
         return 5;
     else
         return force ? (2+4+8) : -1;
 }
 
-static inline void make_jump_patch(void **codep, uintptr_t pc, uintptr_t dpc,
+static inline void make_jump_patch(void **codep, uint_tptr pc, uint_tptr dpc,
                                    UNUSED struct arch_dis_ctx arch) {
-    uintptr_t diff = pc - (dpc + 5);
+    uint_tptr diff = pc - (dpc + 5);
     void *code = *codep;
-    if (diff == (uintptr_t) (int32_t) diff) {
+    if (diff == (uint_tptr) (int32_t) diff) {
         op8(&code, 0xe9);
         op32(&code, diff);
     } else {