Fixed defination of CC_CBXZ collides with other bits

Fixed incorrect pc value in tdctx_to_actx
Fixed ARM32 handling of CC_CBXZ (uses incorrect field in ctx->base)
Fixed ARM32 transform_dis_branch incorrectly trashing LR for non-call based jumps
Fixed ARM32 make_jump_patch not updated to use new assemble_ctx
Added new option SUBSTITUTE_RELAXED to relax the disassembly engine
Currently SUBSTITUTE_RELAXED only disables TRANSFORM_DIS_REL_JUMPS so jumps at the beginning of functions are allowed

7 files changed, 41 insertions, 13 deletions

diff --git a/lib/arm/arch-dis.h b/lib/arm/arch-dis.h
index 1c63b12..157acfb 100644
--- a/lib/arm/arch-dis.h
+++ b/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)
diff --git a/lib/arm/arch-transform-dis.inc.h b/lib/arm/arch-transform-dis.inc.h
index 1e9579d..b23c2c3 100644
--- a/lib/arm/arch-transform-dis.inc.h
+++ b/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,23 +166,41 @@ 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;
+    }
+    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;
     }
-    /* If it's a call, we should jump back after the call */
+
     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);
 }
 
diff --git a/lib/arm/jump-patch.h b/lib/arm/jump-patch.h
index 8ba4c3a..361e357 100644
--- a/lib/arm/jump-patch.h
+++ b/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);
 }
diff --git a/lib/hook-functions.c b/lib/hook-functions.c
index 674d424..d2fabe9 100644
--- a/lib/hook-functions.c
+++ b/lib/hook-functions.c
@@ -147,6 +147,7 @@ int substitute_hook_functions(const struct substitute_function_hook *hooks,
 #else
     bool thread_safe = false;
 #endif
+    bool relaxed = !!(options & SUBSTITUTE_RELAXED);
 
     if (recordp)
         *recordp = NULL;
@@ -248,7 +249,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, trampoline_addr,
                                       &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;
         trampoline_addr += (trampoline_ptr - trampoline_prev);
 
diff --git a/lib/substitute.h b/lib/substitute.h
index 8cc2a4d..023a728 100644
--- a/lib/substitute.h
+++ b/lib/substitute.h
@@ -119,6 +119,7 @@ struct substitute_function_hook_record {
 #ifndef NO_PTHREADS
 enum {
     SUBSTITUTE_NO_THREAD_SAFETY = 1,
+    SUBSTITUTE_RELAXED = 2,
 };
 #endif
 
@@ -141,6 +142,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
diff --git a/lib/transform-dis.c b/lib/transform-dis.c
index 4d69da1..20e2dfe 100644
--- a/lib/transform-dis.c
+++ b/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;
diff --git a/lib/transform-dis.h b/lib/transform-dis.h
index 90156af..18e0b17 100644
--- a/lib/transform-dis.h
+++ b/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,