1 files changed, 17 insertions, 8 deletions

diff --git a/lib/hook-functions.c b/lib/hook-functions.c
index 996e781..11fa9ac 100644
--- a/lib/hook-functions.c
+++ b/lib/hook-functions.c
@@ -19,6 +19,7 @@ struct hook_internal {
     /* page allocated with execmem_alloc_unsealed - only if we had to allocate
      * one when processing this hook */
     void *trampoline_page;
+    uintptr_t trampoline_addr;
     struct arch_dis_ctx arch_dis_ctx;
 };
 
@@ -72,6 +73,7 @@ static int check_intro_trampoline(void **trampoline_ptr_p,
                                   int *patch_size_p,
                                   bool *need_intro_trampoline_p,
                                   void **trampoline_page_p,
+                                  uintptr_t *trampoline_page_addr_p,
                                   struct arch_dis_ctx arch,
                                   void *opt) {
     void *trampoline_ptr = *trampoline_ptr_p;
@@ -107,7 +109,7 @@ static int check_intro_trampoline(void **trampoline_ptr_p,
             goto end;
         }
 
-        execmem_free(trampoline_ptr, opt);
+        execmem_free(trampoline_ptr, trampoline_addr, opt);
     }
 
     /* Allocate new trampoline - try before pc (xxx only meaningful on arm64) */
@@ -121,7 +123,7 @@ static int check_intro_trampoline(void **trampoline_ptr_p,
             goto end;
         }
 
-        execmem_free(trampoline_ptr, opt);
+        execmem_free(trampoline_ptr, trampoline_addr, opt);
         ret = SUBSTITUTE_ERR_OUT_OF_RANGE;
     }
 
@@ -130,6 +132,7 @@ end:
     *trampoline_addr_p = trampoline_addr;
     *trampoline_size_left_p = trampoline_size_left;
     *trampoline_page_p = trampoline_ptr;
+    *trampoline_page_addr_p = trampoline_addr;
     return ret;
 }
 
@@ -191,7 +194,8 @@ int substitute_hook_functions(const struct substitute_function_hook *hooks,
                                           &trampoline_size_left, pc_patch_start,
                                           (uintptr_t) hook->replacement,
                                           &patch_size, &need_intro_trampoline,
-                                          &hi->trampoline_page, arch, 
+                                          &hi->trampoline_page, 
+                                          &hi->trampoline_addr, arch, 
                                           hook->opt)))
             goto end;
 
@@ -222,10 +226,13 @@ int substitute_hook_functions(const struct substitute_function_hook *hooks,
                                               &trampoline_size_left, 
                                               hook->opt)))
                 goto end;
-            /* NOTE: We assume that each page is large enough (min 0x1000)
-             * so we don't lose a reference by having one hook allocate two 
-             * pages. */
+            /* NOTE: We assume that each page is large enough (min
+             * TD_MAX_REWRITTEN_SIZE + 2 * MAX_JUMP_PATCH_SIZE) so we don't lose
+             * a reference by having one hook allocate two pages. Also must
+             * ensure this size is aligned to ARCH_MAX_CODE_ALIGNMENT otherwise
+             * MAX_JUMP_PATCH_SIZE might be wrong. */
             hi->trampoline_page = trampoline_ptr;
+            hi->trampoline_addr = trampoline_addr;
         }
 
         void *outro_trampoline_real = trampoline_ptr;
@@ -275,8 +282,9 @@ int substitute_hook_functions(const struct substitute_function_hook *hooks,
     for (size_t i = 0; i < nhooks; i++) {
         struct hook_internal *hi = &his[i];
         void *page = hi->trampoline_page;
+        uintptr_t addr = hi->trampoline_addr;
         if (page)
-            execmem_seal(page, hooks[i].opt);
+            execmem_seal(page, addr, hooks[i].opt);
         fws[i].dst = hi->code;
         fws[i].src = hi->jump_patch;
         fws[i].len = hi->jump_patch_size;
@@ -299,8 +307,9 @@ end:
     /* if we failed, get rid of the trampolines. */
     for (size_t i = 0; i < nhooks; i++) {
         void *page = his[i].trampoline_page;
+        uintptr_t addr = his[i].trampoline_addr;
         if (page)
-            execmem_free(page, hooks[i].opt);
+            execmem_free(page, addr, hooks[i].opt);
     }
 end_dont_free:
     return ret;