Added ARM32 support for transforming branches at function beginning

lib/arm/arch-dis.h

@@ -62,4 +62,4 @@ static inline void advance_it_cond(struct arch_dis_ctx *ctx) {
  * to keep going */
 #define CC_ALREADY_IN_IT (CC_CONDITIONAL | 0x800)
 /* CBZ/CBNZ is rewritten */
-#define CC_CBXZ          (CC_CONDITIONAL | 0xc00)
+#define CC_CBXZ          (CC_CONDITIONAL | 0x1000)

lib/arm/arch-transform-dis.inc.h

@@ -1,5 +1,6 @@
 /* TODO fix BL incl MOV LR, PC */
 #include "arm/assemble.h"
+#include "arm/jump-patch.h"
 
 static struct assemble_ctx tdctx_to_actx(const struct transform_dis_ctx *ctx) {
     int cond;
@@ -13,7 +14,7 @@ static struct assemble_ctx tdctx_to_actx(const struct transform_dis_ctx *ctx) {
     return (struct assemble_ctx) {
         ctx->rewritten_ptr_ptr,
         *ctx->rewritten_ptr_ptr,
-        (uint_tptr) (uintptr_t) *ctx->rewritten_ptr_ptr,
+        ctx->base.pc,
         ctx->arch.pc_low_bit,
         cond
     };
@@ -165,22 +166,43 @@ void transform_dis_branch(struct transform_dis_ctx *ctx, uint_tptr dpc, int cc)
     transform_dis_branch_top(ctx, dpc, cc);
     struct assemble_ctx actx = tdctx_to_actx(ctx);
     ctx->write_newop_here = NULL;
-    int replacement_size = 8 + (actx.thumb ? 2 : 4);
+    int replacement_size;
+
+    /* Dry run to get size */
+    if ((cc & CC_ARMCC) == CC_ARMCC) {
+        replacement_size = actx.thumb ? 2 : 4;
+    } else if ((cc & CC_CBXZ) == CC_CBXZ) {
+        replacement_size = 2;
+    } else {
+        replacement_size = 0;
+    }
+    if ((cc & CC_CALL) == CC_CALL) {
+        replacement_size += 8 + (actx.thumb ? 2 : 4);
+    } else {
+        replacement_size += jump_patch_size(actx_pc(actx) + replacement_size, dpc, ctx->arch, 0);
+    }
+
+    /* Actual run */
     if ((cc & CC_ARMCC) == CC_ARMCC) {
-        replacement_size += actx.thumb ? 2 : 4;
         actx.cond = invert_arm_cond(cc & 0xf);
         Bccrel(actx, replacement_size);
     } else if ((cc & CC_CBXZ) == CC_CBXZ) {
-        replacement_size += 2;
         ctx->base.modify = true;
-        ctx->base.newval[0] = actx.pc_of_code_base + replacement_size;
-        ctx->base.newval[1] = 1; /* do invert */
+        ctx->base.newop[0] = actx.pc_of_code_base + replacement_size;
+        ctx->base.newop[1] = 1; /* do invert */
         void **codep = ctx->rewritten_ptr_ptr;
         ctx->write_newop_here = *codep; *codep += 2;
     }
+
     actx.cond = 0xe;
-    MOVW_MOVT(actx, 14, dpc | ctx->arch.pc_low_bit);
-    BLXr(actx, 14);
+    /* If it's a call, we should jump back after the call */
+    if ((cc & CC_CALL) == CC_CALL) {
+        MOVW_MOVT(actx, 14, dpc | ctx->arch.pc_low_bit);
+        BLXr(actx, 14);
+    } else {
+        // otherwise we can't clobber LR
+        LDR_PC(actx, dpc | ctx->arch.pc_low_bit);
+    }
     substitute_assert(*actx.codep - actx.code_base == replacement_size);
 }
 

lib/arm/jump-patch.h

@@ -14,6 +14,6 @@ static inline int jump_patch_size(uint_tptr pc,
 static inline void make_jump_patch(void **codep, uint_tptr pc,
                                    uint_tptr dpc,
                                    struct arch_dis_ctx arch) {
-    struct assemble_ctx actx = {codep, pc, arch.pc_low_bit, 0xe};
+    struct assemble_ctx actx = {codep, *codep, pc, arch.pc_low_bit, 0xe};
     LDR_PC(actx, dpc);
 }

lib/hook-functions.c

@@ -138,6 +138,7 @@ int substitute_hook_functions(const struct substitute_function_hook *hooks,
     bool thread_safe = !(options & SUBSTITUTE_NO_THREAD_SAFETY);
     if (thread_safe && !pthread_main_np())
         return SUBSTITUTE_ERR_NOT_ON_MAIN_THREAD;
+    bool relaxed = !!(options & SUBSTITUTE_RELAXED);
 
     if (recordp)
         *recordp = NULL;
@@ -224,7 +225,8 @@ int substitute_hook_functions(const struct substitute_function_hook *hooks,
         if ((ret = transform_dis_main(code, &trampoline_ptr, pc_patch_start,
                                       &pc_patch_end, (uintptr_t) trampoline_ptr,
                                       &arch, hi->offset_by_pcdiff,
-                                      thread_safe ? TRANSFORM_DIS_BAN_CALLS : 0)))
+                                      (thread_safe ? TRANSFORM_DIS_BAN_CALLS : 0) | 
+                                      (relaxed ? 0 : TRANSFORM_DIS_REL_JUMPS))))
             goto end;
 
         uintptr_t dpc = pc_patch_end;

lib/substitute.h

@@ -103,6 +103,7 @@ struct substitute_function_hook {
 /* substitute_hook_functions options */
 enum {
     SUBSTITUTE_NO_THREAD_SAFETY = 1,
+    SUBSTITUTE_RELAXED = 2,
 };
 
 /* Patch the machine code of the specified functions to redirect them to the
@@ -124,6 +125,9 @@ enum {
  *
  * You can disable the main thread check and all synchronization by passing
  * SUBSTITUTE_NO_THREAD_SAFETY.
+ * 
+ * You can relax the disassembly engine (at the risk of possible incorrect 
+ * results) to be compatible with more functions by passing SUBSTITUTE_RELAXED.
  *
  * Why not just use a mutex to prevent deadlock?  That would work between
  * multiple calls into libsubstitute, but there may be other libraries that

lib/transform-dis.c

@@ -24,6 +24,7 @@ struct transform_dis_ctx {
     bool force_keep_transforming;
 
     bool ban_calls; /* i.e. trying to be thread safe */
+    bool ban_jumps; /* allow transforming rel branches at beginning */
 
     void **rewritten_ptr_ptr;
     void *write_newop_here;
@@ -78,7 +79,7 @@ static void transform_dis_branch_top(struct transform_dis_ctx *ctx,
     }
     if (cc & CC_CALL) {
         transform_dis_indirect_call(ctx);
-    } else {
+    } else if (ctx->ban_jumps) {
         transform_dis_ret(ctx);
     }
 }
@@ -102,6 +103,7 @@ int transform_dis_main(const void *restrict code_ptr,
     ctx.base.pc = pc_patch_start;
     ctx.arch = *arch_ctx_p;
     ctx.ban_calls = options & TRANSFORM_DIS_BAN_CALLS;
+    ctx.ban_jumps = options & TRANSFORM_DIS_REL_JUMPS;
     /* data is written to rewritten both by this function directly and, in case
      * additional scaffolding is needed, by arch-specific transform_dis_* */
     ctx.rewritten_ptr_ptr = rewritten_ptr_ptr;

lib/transform-dis.h

@@ -4,6 +4,7 @@
 #include "dis.h"
 
 #define TRANSFORM_DIS_BAN_CALLS 1
+#define TRANSFORM_DIS_REL_JUMPS 2
 
 int transform_dis_main(const void *restrict code_ptr,
                        void **restrict rewritten_ptr_ptr,