handle oom and silly machos and stuff

2 files changed, 31 insertions, 10 deletions

diff --git a/lib/interpose.c b/lib/interpose.c
index 9718813..4857648 100644
--- a/lib/interpose.c
+++ b/lib/interpose.c
@@ -15,7 +15,8 @@
 enum { MAX_SEGMENTS = 32 };
 struct interpose_state {
     int nsegments;
-    segment_command_x *segments[MAX_SEGMENTS];
+    segment_command_x **segments;
+    segment_command_x *stack_segments[MAX_SEGMENTS];
     uintptr_t slide;
     const struct substitute_import_hook *hooks;
     size_t nhooks;
@@ -54,7 +55,6 @@ static int try_bind_section(void *bind, size_t size, const struct interpose_stat
     uint8_t type = lazy ? BIND_TYPE_POINTER : 0;
     intptr_t addend = 0;
     size_t offset = 0;
-    int n = 0;
     void *segment = NULL;
     while (ptr < end) {
         uint8_t byte = *(uint8_t *) ptr;
@@ -135,7 +135,7 @@ static int try_bind_section(void *bind, size_t size, const struct interpose_stat
             break;
         }
     }
-    return n;
+    return SUBSTITUTE_OK;
 }
 
 static void *off_to_addr(const struct interpose_state *st, uint32_t off) {
@@ -151,19 +151,29 @@ EXPORT
 int substitute_interpose_imports(const struct substitute_image *image,
                                  const struct substitute_import_hook *hooks,
                                  size_t nhooks, UNUSED int options) {
+    int ret = SUBSTITUTE_OK;
     struct interpose_state st;
     st.slide = image->slide;
     st.nsegments = 0;
     st.hooks = hooks;
     st.nhooks = nhooks;
+    st.segments = st.stack_segments;
 
     const mach_header_x *mh = image->image_header;
     const struct load_command *lc = (void *) (mh + 1);
     for (uint32_t i = 0; i < mh->ncmds; i++) {
         if (lc->cmd == LC_SEGMENT_X) {
             segment_command_x *sc = (void *) lc;
-            if (st.nsegments < MAX_SEGMENTS)
-                st.segments[st.nsegments++] = sc;
+            if (st.nsegments == MAX_SEGMENTS) {
+                segment_command_x **new = calloc(st.nsegments * 2, sizeof(*st.segments));
+                if (!new)
+                    substitute_panic("%s: out of memory\n", __func__);
+                memcpy(new, st.segments, st.nsegments * sizeof(*st.segments));
+                if (st.segments != st.stack_segments)
+                    free(st.segments);
+                st.segments = new;
+            }
+            st.segments[st.nsegments++] = sc;
         }
         lc = (void *) lc + lc->cmdsize;
     }
@@ -176,13 +186,16 @@ int substitute_interpose_imports(const struct substitute_image *image,
             if ((ret = try_bind_section(off_to_addr(&st, dc->bind_off), dc->bind_size, &st, false)) ||
                 (ret = try_bind_section(off_to_addr(&st, dc->weak_bind_off), dc->weak_bind_size, &st, false)) ||
                 (ret = try_bind_section(off_to_addr(&st, dc->lazy_bind_off), dc->lazy_bind_size, &st, true)))
-                return ret;
+                goto fail;
 
             break;
         }
         lc = (void *) lc + lc->cmdsize;
     }
-    return SUBSTITUTE_OK;
+fail:
+    if (st.segments != st.stack_segments)
+        free(st.segments);
+    return ret;
 }
 
 #endif /* __APPLE__ */
diff --git a/lib/jump-dis.inc.h b/lib/jump-dis.inc.h
index 0bfa495..a070a74 100644
--- a/lib/jump-dis.inc.h
+++ b/lib/jump-dis.inc.h
@@ -60,6 +60,8 @@ static void P(add_to_queue)(struct jump_dis_ctx *ctx, uintptr_t pc) {
     if (ctx->queue_write_off == ctx->queue_read_off && (ctx->queue_count || !ctx->queue_size)) {
         size_t new_size = ctx->queue_size * 2 + 5;
         ctx->queue = realloc(ctx->queue, new_size * sizeof(*ctx->queue));
+        if (!ctx->queue)
+            substitute_panic("%s: out of memory\n", __func__);
         size_t new_read_off = new_size - (ctx->queue_size - ctx->queue_read_off);
         memmove(ctx->queue + new_read_off, ctx->queue + ctx->queue_read_off, (ctx->queue_size - ctx->queue_read_off) * sizeof(*ctx->queue));
         ctx->queue_read_off = new_read_off % new_size;
@@ -109,6 +111,7 @@ static void P(bad)(struct jump_dis_ctx *ctx) {
 static void P(dis)(tdis_ctx ctx);
 
 bool P(main)(void *code_ptr, uintptr_t pc_patch_start, uintptr_t pc_patch_end, bool pc_low_bit) {
+    bool ret;
     struct jump_dis_ctx ctx;
     memset(&ctx, 0, sizeof(ctx));
     ctx.pc_patch_start = pc_patch_start;
@@ -128,8 +131,10 @@ bool P(main)(void *code_ptr, uintptr_t pc_patch_start, uintptr_t pc_patch_end, b
             ctx.bad_insn,
             ctx.continue_after_this_insn);
 #endif
-        if (ctx.bad_insn)
-            return true;
+        if (ctx.bad_insn) {
+            ret = true;
+            goto fail;
+        }
         if (ctx.continue_after_this_insn)
             P(add_to_queue)(&ctx, ctx.pc + ctx.op_size);
 
@@ -141,5 +146,8 @@ bool P(main)(void *code_ptr, uintptr_t pc_patch_start, uintptr_t pc_patch_end, b
         ctx.queue_count--;
     }
     /* no bad instructions! */
-    return false;
+    ret = false;
+fail:
+    free(ctx.queue);
+    return ret;
 }