Merge pull request #78 from yosefe/topic/ucm-bistro-use-relative-jump…

…-int3

UCM/BISTRO: Use relative jump instead of absolute jump

src/ucm/bistro/bistro.c

@@ -61,11 +61,13 @@ ucs_status_t ucm_bistro_apply_patch(void *dst, void *patch, size_t len)
 
 #if defined(__x86_64__) || defined (__aarch64__)
 struct ucm_bistro_restore_point {
-    void               *addr; /* address of function to restore */
-    ucm_bistro_patch_t patch; /* original function body */
+    void               *addr;     /* address of function to restore */
+    size_t             patch_len; /* patch length */
+    char               orig[0];   /* orig func code */
 };
 
-ucs_status_t ucm_bistro_create_restore_point(void *addr, ucm_bistro_restore_point_t **rp)
+ucs_status_t ucm_bistro_create_restore_point(void *addr, size_t len,
+                                             ucm_bistro_restore_point_t **rp)
 {
     ucm_bistro_restore_point_t *point;
 
@@ -74,22 +76,24 @@ ucs_status_t ucm_bistro_create_restore_point(void *addr, ucm_bistro_restore_poin
         return UCS_OK;
     }
 
-    point = malloc(sizeof(*point));
-    if (!point) {
+    point = malloc(sizeof(*point) + len);
+    if (point == NULL) {
         return UCS_ERR_NO_MEMORY;
     }
 
-    point->addr  = addr;
-    point->patch = *(ucm_bistro_patch_t*)addr;
-    *rp          = point;
+    *rp              = point;
+    point->addr      = addr;
+    point->patch_len = len;
+    memcpy(point->orig, addr, len);
+
     return UCS_OK;
 }
 
 ucs_status_t ucm_bistro_restore(ucm_bistro_restore_point_t *rp)
 {
     ucs_status_t status;
 
-    status = ucm_bistro_apply_patch(rp->addr, &rp->patch, sizeof(rp->patch));
+    status = ucm_bistro_apply_patch(rp->addr, rp->orig, rp->patch_len);
     if (!UCS_STATUS_IS_ERR(status)) {
         ucm_bistro_remove_restore_point(rp);
     }

src/ucm/bistro/bistro_aarch64.c

@@ -75,7 +75,7 @@ ucs_status_t ucm_bistro_patch(const char *symbol, void *hook,
 
     UCM_LOOKUP_SYMBOL(func, symbol);
 
-    status = ucm_bistro_create_restore_point(func, rp);
+    status = ucm_bistro_create_restore_point(func, sizeof(patch), rp);
     if (UCS_STATUS_IS_ERR(status)) {
         return status;
     }

src/ucm/bistro/bistro_int.h

@@ -30,7 +30,8 @@
 
 ucs_status_t ucm_bistro_apply_patch(void *dst, void *patch, size_t len);
 
-ucs_status_t ucm_bistro_create_restore_point(void *addr, ucm_bistro_restore_point_t **rp);
+ucs_status_t ucm_bistro_create_restore_point(void *addr, size_t len,
+                                             ucm_bistro_restore_point_t **rp);
 
 static inline void *ucm_bistro_lookup(const char *symbol)
 {

src/ucm/bistro/bistro_x86_64.c

@@ -25,27 +25,44 @@
 #include <ucs/arch/cpu.h>
 #include <ucs/debug/assert.h>
 
-static const ucm_bistro_patch_t patch_tmpl = {
-    .mov_r11 = {0x49, 0xbb},
-    .jmp_r11 = {0x41, 0xff, 0xe3}
-};
 
 ucs_status_t ucm_bistro_patch(const char *symbol, void *hook,
                               ucm_bistro_restore_point_t **rp)
 {
-    ucm_bistro_patch_t patch = patch_tmpl;
+    ucm_bistro_jmp_r11_patch_t patch_jmp_r11   = {
+        .mov_r11 = {0x49, 0xbb},
+        .jmp_r11 = {0x41, 0xff, 0xe3}
+    };
+    ucm_bistro_jmp_near_patch_t patch_jmp_near = {
+        .jmp_rel = 0xe9
+    };
+    void *func, *patch, *jmp_base;
     ucs_status_t status;
-    void *func;
+    ptrdiff_t jmp_disp;
+    size_t patch_len;
 
     UCM_LOOKUP_SYMBOL(func, symbol);
 
-    patch.ptr = hook;
+    jmp_base = UCS_PTR_BYTE_OFFSET(func, sizeof(patch_jmp_near));
+    jmp_disp = UCS_PTR_BYTE_DIFF(jmp_base, hook);
+    if (labs(jmp_disp) < INT32_MAX) {
+        /* if 32-bit near jump is possible, use it, since it's a short 5-byte
+         * instruction which reduces the chances of racing with other thread
+         */
+        patch_jmp_near.disp = jmp_disp;
+        patch               = &patch_jmp_near;
+        patch_len           = sizeof(patch_jmp_near);
+    } else {
+        patch_jmp_r11.ptr   = hook;
+        patch               = &patch_jmp_r11;
+        patch_len           = sizeof(patch_jmp_r11);
+    }
 
-    status = ucm_bistro_create_restore_point(func, rp);
+    status = ucm_bistro_create_restore_point(func, patch_len, rp);
     if (UCS_STATUS_IS_ERR(status)) {
         return status;
     }
 
-    return ucm_bistro_apply_patch(func, &patch, sizeof(patch));
+    return ucm_bistro_apply_patch(func, patch, patch_len);
 }
 #endif

src/ucm/bistro/bistro_x86_64.h

@@ -16,11 +16,20 @@
 #define UCM_BISTRO_PROLOGUE
 #define UCM_BISTRO_EPILOGUE
 
-typedef struct ucm_bistro_patch {
+/* Patch by jumping to absolute address loaded from register */
+typedef struct ucm_bistro_jmp_r11_patch {
     uint8_t mov_r11[2];  /* mov %r11, addr */
     void    *ptr;
     uint8_t jmp_r11[3];  /* jmp r11        */
-} UCS_S_PACKED ucm_bistro_patch_t;
+} UCS_S_PACKED ucm_bistro_jmp_r11_patch_t;
+
+
+/* Patch by jumping to relative address by immediate displacement */
+typedef struct ucm_bistro_jmp_near_patch {
+    uint8_t jmp_rel; /* opcode:  JMP rel32          */
+    int32_t disp;    /* operand: jump displacement */
+} UCS_S_PACKED ucm_bistro_jmp_near_patch_t;
+
 
 /**
  * Set library function call hook using Binary Instrumentation