various fixes

6 files changed, 61 insertions, 55 deletions

diff --git a/lib/arm/arch-dis.h b/lib/arm/arch-dis.h
index 8f2400b..1c63b12 100644
--- a/lib/arm/arch-dis.h
+++ b/lib/arm/arch-dis.h
@@ -16,6 +16,7 @@
  * since only the actually used space will be taken up in the final output
  */
 #define TD_MAX_REWRITTEN_SIZE (7*4*7 + 4) /* 196 */
+#define ARCH_MAX_CODE_ALIGNMENT 4
 
 struct arch_pcrel_info {
     unsigned reg;
diff --git a/lib/arm64/arch-dis.h b/lib/arm64/arch-dis.h
index f5b0518..142186f 100644
--- a/lib/arm64/arch-dis.h
+++ b/lib/arm64/arch-dis.h
@@ -1,6 +1,7 @@
 #pragma once
 #define MIN_INSN_SIZE 4
 #define TD_MAX_REWRITTEN_SIZE (7 * 2 * 4) /* also conservative */
+#define ARCH_MAX_CODE_ALIGNMENT 4
 
 struct arch_pcrel_info {
     unsigned reg;
@@ -26,10 +27,6 @@ static inline void arch_dis_ctx_init(struct arch_dis_ctx *ctx) {
     ctx->regs_possibly_written = 0;
 }
 
-static inline int arch_code_alignment(UNUSED struct arch_dis_ctx ctx) {
-   return 4;
-}
-
 static inline int arm64_get_unwritten_temp_reg(struct arch_dis_ctx *ctx) {
     uint32_t avail = ~ctx->regs_possibly_written & ((1 << 19) - (1 << 9));
     if (!avail)
diff --git a/lib/darwin/execmem.c b/lib/darwin/execmem.c
index 419d7b5..d99e947 100644
--- a/lib/darwin/execmem.c
+++ b/lib/darwin/execmem.c
@@ -16,6 +16,7 @@
 #include <stdlib.h>
 #include <ucontext.h>
 #include <signal.h>
+#include <pthread.h>
 
 #define port_hash(portp) (*(portp))
 #define port_eq(port1p, port2p) (*(port1p) == *(port2p))
@@ -198,6 +199,12 @@ static void resume_other_threads() {
 
 static void segfault_handler(UNUSED int sig, UNUSED siginfo_t *info,
                              void *uap_) {
+    if (pthread_main_np()) {
+        /* The patcher itself segfaulted.  Oops.  Reset the signal so the
+         * process exits rather than going into an infinite loop. */
+        signal(sig, SIG_DFL);
+        return;
+    }
     /* We didn't catch it before it segfaulted so have to fix it up here. */
     ucontext_t *uap = uap_;
     apply_one_pcp_with_state(&uap->uc_mcontext->__ss, g_pc_patch_callback,
@@ -311,8 +318,9 @@ int execmem_foreign_write_with_pc_patch(struct execmem_foreign_write *writes,
 
         last = first;
         while (last + 1 < nwrites) {
-            uintptr_t this_start = (uintptr_t) first_write->dst & ~PAGE_MASK;
-            uintptr_t this_end = ((uintptr_t) first_write->dst +
+            const struct execmem_foreign_write *write = &writes[last + 1];
+            uintptr_t this_start = (uintptr_t) write->dst & ~PAGE_MASK;
+            uintptr_t this_end = ((uintptr_t) write->dst +
                                   first_write->len - 1) & ~PAGE_MASK;
             if (page_start <= this_start && this_start <= page_end) {
                 if (this_end > page_end)
@@ -408,7 +416,7 @@ int execmem_foreign_write_with_pc_patch(struct execmem_foreign_write *writes,
         /* This is probably useless, since the original page is gone
          * forever (intentionally, see above).  May as well arrange the
          * deck chairs, though. */
-        munmap(new, PAGE_SIZE);
+        munmap(new, len);
         goto fail;
     }
 
diff --git a/lib/hook-functions.c b/lib/hook-functions.c
index 5d1f1d5..faadd82 100644
--- a/lib/hook-functions.c
+++ b/lib/hook-functions.c
@@ -80,46 +80,48 @@ static int check_intro_trampoline(void **trampoline_ptr_p,
 
     *need_intro_trampoline_p = true;
 
-    /* Try existing trampoline */
-    *patch_size_p = jump_patch_size(pc, (uintptr_t) trampoline_ptr, arch, false);
+    if (trampoline_ptr) {
+        /* Try existing trampoline */
+        *patch_size_p = jump_patch_size(pc, (uintptr_t) trampoline_ptr, arch, false);
 
-    if (*patch_size_p != -1 && (size_t) *patch_size_p <= *trampoline_size_left_p)
-        return SUBSTITUTE_OK;
+        if (*patch_size_p != -1 && (size_t) *patch_size_p <= *trampoline_size_left_p)
+            return SUBSTITUTE_OK;
+    }
 
     /* Allocate new trampoline - try after pc.  If this fails, we can try
      * before pc before giving up. */
     int ret = execmem_alloc_unsealed(pc, &trampoline_ptr, &trampoline_size_left);
-    if (ret)
-        goto skip_after;
-
-    *patch_size_p = jump_patch_size(pc, (uintptr_t) trampoline_ptr, arch, false);
-    if (*patch_size_p != -1) {
-        *trampoline_ptr_p = trampoline_ptr;
-        *trampoline_size_left_p = trampoline_size_left;
-        *trampoline_page_p = trampoline_ptr;
-        return SUBSTITUTE_OK;
-    }
+    if (!ret) {
+        *patch_size_p = jump_patch_size(pc, (uintptr_t) trampoline_ptr, arch, false);
+        if (*patch_size_p != -1) {
+            ret = SUBSTITUTE_OK;
+            goto end;
+        }
 
-    execmem_free(trampoline_ptr);
+        execmem_free(trampoline_ptr);
+    }
 
-skip_after:;
     /* Allocate new trampoline - try before pc (xxx only meaningful on arm64) */
-    uintptr_t start_address = pc - 0xffff0000;
+    uintptr_t start_address = pc - 0x80000000;
     ret = execmem_alloc_unsealed(start_address, &trampoline_ptr, &trampoline_size_left);
-    if (ret)
-        return ret;
-
-    *patch_size_p = jump_patch_size(pc, (uintptr_t) trampoline_ptr, arch, false);
-    if (*patch_size_p != -1) {
-        *trampoline_ptr_p = trampoline_ptr;
-        *trampoline_size_left_p = trampoline_size_left;
-        *trampoline_page_p = trampoline_ptr;
-        return SUBSTITUTE_OK;
+    if (!ret) {
+        *patch_size_p = jump_patch_size(pc, (uintptr_t) trampoline_ptr, arch, false);
+        if (*patch_size_p != -1) {
+            *trampoline_ptr_p = trampoline_ptr;
+            *trampoline_size_left_p = trampoline_size_left;
+            *trampoline_page_p = trampoline_ptr;
+            return SUBSTITUTE_OK;
+        }
+
+        execmem_free(trampoline_ptr);
+        ret = SUBSTITUTE_ERR_OUT_OF_RANGE;
     }
 
-    /* I give up... */
-    execmem_free(trampoline_ptr);
-    return SUBSTITUTE_ERR_OUT_OF_RANGE;
+end:
+    *trampoline_ptr_p = trampoline_ptr;
+    *trampoline_size_left_p = trampoline_size_left;
+    *trampoline_page_p = trampoline_ptr;
+    return ret;
 }
 
 
@@ -185,10 +187,7 @@ int substitute_hook_functions(const struct substitute_function_hook *hooks,
         make_jump_patch(&jp, pc_patch_start, initial_target, arch);
         hi->jump_patch_size = (uint8_t *) jp - hi->jump_patch;
 
-        size_t align_bytes = (-(uintptr_t) trampoline_ptr)
-                             & (arch_code_alignment(arch) - 1);
-        size_t outro_est = align_bytes +
-                           TD_MAX_REWRITTEN_SIZE + MAX_JUMP_PATCH_SIZE;
+        size_t outro_est = TD_MAX_REWRITTEN_SIZE + MAX_JUMP_PATCH_SIZE;
 
         if (outro_est > trampoline_size_left) {
             /* Not enough space left in our existing block... */
@@ -196,20 +195,12 @@ int substitute_hook_functions(const struct substitute_function_hook *hooks,
                                               &trampoline_size_left)))
                 goto end;
             hi->trampoline_page = trampoline_ptr;
-        } else {
-            trampoline_ptr += align_bytes;
-            trampoline_size_left -= align_bytes;
         }
 
         hi->outro_trampoline = trampoline_ptr;
-        *(void **) hook->old_ptr = hi->outro_trampoline;
-        uintptr_t dpc = pc_patch_end;
-#ifdef __arm__
-        if (arch.pc_low_bit) {
-            hi->outro_trampoline++;
-            dpc++;
-        }
-#endif
+        if (hook->old_ptr)
+            *(void **) hook->old_ptr = hi->outro_trampoline;
+
         /* Generate the rewritten start of the function for the outro
          * trampoline (complaining if any bad instructions are found)
          * (on arm64, this modifies regs_possibly_written, which is used by the
@@ -218,6 +209,15 @@ int substitute_hook_functions(const struct substitute_function_hook *hooks,
                                       &pc_patch_end, (uintptr_t) trampoline_ptr,
                                       &arch, hi->offset_by_pcdiff)))
             goto end;
+
+        uintptr_t dpc = pc_patch_end;
+#ifdef __arm__
+        if (arch.pc_low_bit) {
+            hi->outro_trampoline++;
+            dpc++;
+        }
+#endif
+
         /* Now that transform_dis_main has given us the final pc_patch_end,
          * check some of the rest of the function for jumps back into the
          * patched region. */
@@ -225,6 +225,7 @@ int substitute_hook_functions(const struct substitute_function_hook *hooks,
             goto end;
         /* Okay, continue with the outro. */
         make_jump_patch(&trampoline_ptr, (uintptr_t) trampoline_ptr, dpc, arch);
+        trampoline_ptr += -(uintptr_t) trampoline_ptr % ARCH_MAX_CODE_ALIGNMENT;
         trampoline_size_left -= (uint8_t *) trampoline_ptr
                               - (uint8_t *) hi->outro_trampoline;
     }
@@ -251,6 +252,7 @@ int substitute_hook_functions(const struct substitute_function_hook *hooks,
         goto end_dont_free;
     }
 
+    goto end_dont_free;
 end:
     /* if we failed, get rid of the trampolines. */
     for (size_t i = 0; i < nhooks; i++) {
diff --git a/lib/x86/arch-dis.h b/lib/x86/arch-dis.h
index a76884f..c52561a 100644
--- a/lib/x86/arch-dis.h
+++ b/lib/x86/arch-dis.h
@@ -6,6 +6,7 @@
  * See transform_dis_* for size figures.  Technically unsafe, since we don't
  * check for overlong x86 instructions. */
 #define TD_MAX_REWRITTEN_SIZE 55
+#define ARCH_MAX_CODE_ALIGNMENT 1
 
 struct arch_pcrel_info {
     int reg;
@@ -14,6 +15,3 @@ 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/jump-patch.h b/lib/x86/jump-patch.h
index ca7b16c..569eb74 100644
--- a/lib/x86/jump-patch.h
+++ b/lib/x86/jump-patch.h
@@ -5,7 +5,7 @@
 static inline int jump_patch_size(uint_tptr pc, uint_tptr dpc,
                                   UNUSED struct arch_dis_ctx arch,
                                   bool force) {
-    uint_tptr diff = pc - (dpc + 5);
+    uint_tptr diff = dpc - (pc + 5);
     /* fits in 32? */
     if (diff == (uint_tptr) (int32_t) diff)
         return 5;