Added slab allocator for trampoline in Vita platform

Added support for smaller (non-page) allocations for trampoline

6 files changed, 491 insertions, 20 deletions

diff --git a/lib/darwin/execmem.c b/lib/darwin/execmem.c
index 44b26ae..ce3bafd 100644
--- a/lib/darwin/execmem.c
+++ b/lib/darwin/execmem.c
@@ -74,13 +74,13 @@ int execmem_alloc_unsealed(uintptr_t hint, void **page_p, uintptr_t *vma_p,
     return SUBSTITUTE_OK;
 }
 
-int execmem_seal(void *page, UNUSED void *opt) {
+int execmem_seal(void *page, UNUSED uintptr_t vma, UNUSED void *opt) {
     if (mprotect(page, PAGE_SIZE, PROT_READ | PROT_EXEC))
         return SUBSTITUTE_ERR_VM;
     return SUBSTITUTE_OK;
 }
 
-void execmem_free(void *page, UNUSED void *opt) {
+void execmem_free(void *page, UNUSED uintptr_t vma, UNUSED void *opt) {
     munmap(page, PAGE_SIZE);
 }
 
diff --git a/lib/execmem.h b/lib/execmem.h
index 07e9fa4..dbf2dda 100644
--- a/lib/execmem.h
+++ b/lib/execmem.h
@@ -2,8 +2,8 @@
 #include <sys/types.h>
 /* For allocating trampolines - this is just a mmap wrapper. */
 int execmem_alloc_unsealed(uintptr_t hint, void **page_p, uintptr_t *vma_p, size_t *size_p, void *opt);
-int execmem_seal(void *page, void *opt);
-void execmem_free(void *page, void *opt);
+int execmem_seal(void *page, uintptr_t vma, void *opt);
+void execmem_free(void *page, uintptr_t vma, void *opt);
 
 /* Write to "foreign" (i.e. owned by another library, not out-of-process) pages
  * which are already RX or have unknown permissions.
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;
diff --git a/lib/vita/execmem.c b/lib/vita/execmem.c
index 6bfb2ba..526b29b 100644
--- a/lib/vita/execmem.c
+++ b/lib/vita/execmem.c
@@ -1,22 +1,86 @@
-#include <psp2kern/kernel/sysmem.h>
-#include "execmem.h"
 #include "substitute.h"
+#include "dis.h"
+#include "execmem.h"
+#include stringify(TARGET_DIR/jump-patch.h)
+#include <psp2kern/kernel/sysmem.h>
+#include "../../../taihen_internal.h"
 
-int execmem_alloc_unsealed(uintptr_t hint, void **page_p, uintptr_t *vma_p, 
+/** The size of each trampoline allocation. We use it for outro and optional
+ * intro. Realistically, we do not use an intro.
+ */
+#define SLAB_SIZE (TD_MAX_REWRITTEN_SIZE + 2 * MAX_JUMP_PATCH_SIZE)
+#if (SLAB_SIZE % ARCH_MAX_CODE_ALIGNMENT != 0)
+// if not aligned then substitute_hook_functions breaks!
+#error SLAB_SIZE Must be aligned to ARCH_MAX_CODE_ALIGNMENT
+#endif
+
+/**
+ * @file execmem.c
+ *
+ * @brief      Functions for allocating executable memory and writing to RO
+ *             memory.
+ *
+ *             We only consider two arenas for allocating trampoline executable
+ *             memory. One is shared in user space across all processes. The
+ *             other is in kernel space for kernel hooks.
+ */
+
+/**
+ * @brief      Allocate a slab of executable memory.
+ *
+ *             Two pointers will be returned: the executable ro pointer and a
+ *             writable pointer.
+ *
+ * @param[in]  hint       Unused
+ * @param      ptr_p      The writable pointer
+ * @param      vma_p      The executable pointer address
+ * @param      size_p     The size of the allocation. Always `SLAB_SIZE`.
+ * @param      opt        A `tai_substitute_args_t` structure
+ * @param[in]  hint  Unused
+ *
+ * @return     `SUBSTITUTE_OK` or `SUBSTITUTE_ERR_VM` if out of memory
+ */
+int execmem_alloc_unsealed(UNUSED uintptr_t hint, void **ptr_p, uintptr_t *vma_p, 
                            size_t *size_p, void *opt) {
-    return 0;
+    return SUBSTITUTE_OK;
 }
 
-int execmem_seal(void *page, void *opt) {
+/**
+ * @brief      Flushes icache
+ *
+ * @param      ptr   Unused
+ * @param[in]  vma   Pointer to flush
+ * @param      opt   A `tai_substitute_args_t` structure
+ *
+ * @return     `SUBSTITUTE_OK`
+ */
+int execmem_seal(UNUSED void *ptr, uintptr_t vma, void *opt) {
     return SUBSTITUTE_OK;
 }
 
-void execmem_free(void *page, void *opt) {
+/**
+ * @brief      Frees executable memory from slab allocator
+ *
+ * @param      ptr   The writable pointer
+ * @param[in]  vma   The executable address. Must match with ptr.
+ * @param      opt   A `tai_substitute_args_t` structure
+ */
+void execmem_free(void *ptr, uintptr_t vma, void *opt) {
 }
 
+/**
+ * @brief      Write to executable process memory
+ *
+ * @param      writes        List of writes
+ * @param[in]  nwrites       Number of writes
+ * @param[in]  callback      Unused
+ * @param      callback_ctx  Unused
+ *
+ * @return     `SUBSTITUTE_OK` or `SUBSTITUTE_ERR_VM` on failure
+ */
 int execmem_foreign_write_with_pc_patch(struct execmem_foreign_write *writes,
                                         size_t nwrites,
-                                        execmem_pc_patch_callback callback,
-                                        void *callback_ctx, void *opt) {
+                                        UNUSED execmem_pc_patch_callback callback,
+                                        UNUSED void *callback_ctx) {
     return 0;
 }
diff --git a/lib/vita/slab.c b/lib/vita/slab.c
new file mode 100755
index 0000000..0ca3529
--- /dev/null
+++ b/lib/vita/slab.c
@@ -0,0 +1,359 @@
+/* ref: https://github.com/bbu/userland-slab-allocator */
+
+#include "slab.h"
+
+#include <math.h>
+#include <stdint.h>
+#include <stddef.h>
+
+#define assert(x) // turn off asserts
+
+#define SLAB_DUMP_COLOURED
+
+#ifdef SLAB_DUMP_COLOURED
+# define GRAY(s)   "\033[1;30m" s "\033[0m"
+# define RED(s)    "\033[0;31m" s "\033[0m"
+# define GREEN(s)  "\033[0;32m" s "\033[0m"
+# define YELLOW(s) "\033[1;33m" s "\033[0m"
+#else
+# define GRAY(s)   s
+# define RED(s)    s
+# define GREEN(s)  s
+# define YELLOW(s) s
+#endif
+
+#define SLOTS_ALL_ZERO ((uint64_t) 0)
+#define SLOTS_FIRST ((uint64_t) 1)
+#define FIRST_FREE_SLOT(s) ((size_t) __builtin_ctzll(s))
+#define FREE_SLOTS(s) ((size_t) __builtin_popcountll(s))
+#define ONE_USED_SLOT(slots, empty_slotmask) \
+    ( \
+        ( \
+            (~(slots) & (empty_slotmask))       & \
+            ((~(slots) & (empty_slotmask)) - 1)   \
+        ) == SLOTS_ALL_ZERO \
+    )
+
+#define POWEROF2(x) ((x) != 0 && ((x) & ((x) - 1)) == 0)
+
+#define LIKELY(exp) __builtin_expect(exp, 1)
+#define UNLIKELY(exp) __builtin_expect(exp, 0)
+
+const size_t slab_pagesize = 0x1000;
+
+static SceUID sce_exe_alloc(SceUID pid, void **ptr, uintptr_t *exe_addr, SceUID *exe_res, size_t align, size_t size) {
+    return 0;
+}
+
+static int sce_exe_free(SceUID write_res, SceUID exe_res) {
+    return 0;
+}
+
+void slab_init(struct slab_chain *const sch, const size_t itemsize, SceUID pid)
+{
+    assert(sch != NULL);
+    assert(itemsize >= 1 && itemsize <= SIZE_MAX);
+    assert(POWEROF2(slab_pagesize));
+
+    sch->itemsize = itemsize;
+    sch->pid = pid;
+
+    const size_t data_offset = offsetof(struct slab_header, data);
+    const size_t least_slabsize = data_offset + 64 * sch->itemsize;
+    sch->slabsize = (size_t) 1 << (size_t) ceil(log2(least_slabsize));
+    sch->itemcount = 64;
+
+    if (sch->slabsize - least_slabsize != 0) {
+        const size_t shrinked_slabsize = sch->slabsize >> 1;
+
+        if (data_offset < shrinked_slabsize &&
+            shrinked_slabsize - data_offset >= 2 * sch->itemsize) {
+
+            sch->slabsize = shrinked_slabsize;
+            sch->itemcount = (shrinked_slabsize - data_offset) / sch->itemsize;
+        }
+    }
+
+    sch->pages_per_alloc = sch->slabsize > slab_pagesize ?
+        sch->slabsize : slab_pagesize;
+
+    sch->empty_slotmask = ~SLOTS_ALL_ZERO >> (64 - sch->itemcount);
+    sch->initial_slotmask = sch->empty_slotmask ^ SLOTS_FIRST;
+    sch->alignment_mask = ~(sch->slabsize - 1);
+    sch->partial = sch->empty = sch->full = NULL;
+
+    assert(slab_is_valid(sch));
+}
+
+void *slab_alloc(struct slab_chain *const sch, uintptr_t *exe_addr)
+{
+    assert(sch != NULL);
+    assert(slab_is_valid(sch));
+
+    if (LIKELY(sch->partial != NULL)) {
+        /* found a partial slab, locate the first free slot */
+        register const size_t slot = FIRST_FREE_SLOT(sch->partial->slots);
+        sch->partial->slots ^= SLOTS_FIRST << slot;
+
+        if (UNLIKELY(sch->partial->slots == SLOTS_ALL_ZERO)) {
+            /* slab has become full, change state from partial to full */
+            struct slab_header *const tmp = sch->partial;
+
+            /* skip first slab from partial list */
+            if (LIKELY((sch->partial = sch->partial->next) != NULL))
+                sch->partial->prev = NULL;
+
+            if (LIKELY((tmp->next = sch->full) != NULL))
+                sch->full->prev = tmp;
+
+            sch->full = tmp;
+            *exe_addr = sch->full->exe_data + slot * sch->itemsize;
+            return sch->full->data + slot * sch->itemsize;
+        } else {
+            *exe_addr = sch->partial->exe_data + slot * sch->itemsize;
+            return sch->partial->data + slot * sch->itemsize;
+        }
+    } else if (LIKELY((sch->partial = sch->empty) != NULL)) {
+        /* found an empty slab, change state from empty to partial */
+        if (LIKELY((sch->empty = sch->empty->next) != NULL))
+            sch->empty->prev = NULL;
+
+        sch->partial->next = NULL;
+
+        /* slab is located either at the beginning of page, or beyond */
+        UNLIKELY(sch->partial->refcount != 0) ?
+            sch->partial->refcount++ : sch->partial->page->refcount++;
+
+        sch->partial->slots = sch->initial_slotmask;
+        *exe_addr = sch->partial->exe_data;
+        return sch->partial->data;
+    } else {
+        /* no empty or partial slabs available, create a new one */
+        SceUID write_res, exe_res;
+        uintptr_t exe_data;
+        if ((write_res = sce_exe_alloc(sch->pid, (void **)&sch->partial, &exe_data, 
+                          &exe_res, sch->slabsize, sch->pages_per_alloc)) < 0) {
+            *exe_addr = 0;
+            return sch->partial = NULL;
+        }
+        sch->partial->write_res = write_res;
+        sch->partial->exe_res = exe_res;
+        sch->partial->exe_data = exe_data + offsetof(struct slab_header, data);
+        exe_data += sch->slabsize;
+
+        struct slab_header *prev = NULL;
+
+        const char *const page_end =
+            (char *) sch->partial + sch->pages_per_alloc;
+
+        union {
+            const char *c;
+            struct slab_header *const s;
+        } curr = {
+            .c = (const char *) sch->partial + sch->slabsize
+        };
+
+        __builtin_prefetch(sch->partial, 1);
+
+        sch->partial->prev = sch->partial->next = NULL;
+        sch->partial->refcount = 1;
+        sch->partial->slots = sch->initial_slotmask;
+
+        if (LIKELY(curr.c != page_end)) {
+            curr.s->prev = NULL;
+            curr.s->refcount = 0;
+            curr.s->page = sch->partial;
+            curr.s->write_res = write_res;
+            curr.s->exe_res = exe_res;
+            curr.s->exe_data = exe_data + offsetof(struct slab_header, data);
+            exe_data += sch->slabsize;
+            curr.s->slots = sch->empty_slotmask;
+            sch->empty = prev = curr.s;
+
+            while (LIKELY((curr.c += sch->slabsize) != page_end)) {
+                prev->next = curr.s;
+                curr.s->prev = prev;
+                curr.s->refcount = 0;
+                curr.s->page = sch->partial;
+                curr.s->write_res = write_res;
+                curr.s->exe_res = exe_res;
+                curr.s->exe_data = exe_data + offsetof(struct slab_header, data);
+                exe_data += sch->slabsize;
+                curr.s->slots = sch->empty_slotmask;
+                prev = curr.s;
+            }
+
+            prev->next = NULL;
+        }
+
+        *exe_addr = sch->partial->exe_data;
+        return sch->partial->data;
+    }
+
+    /* unreachable */
+}
+
+void slab_free(struct slab_chain *const sch, const void *const addr)
+{
+    assert(sch != NULL);
+    assert(slab_is_valid(sch));
+    assert(addr != NULL);
+
+    struct slab_header *const slab = (void *)
+        ((uintptr_t) addr & sch->alignment_mask);
+
+    register const int slot = ((char *) addr - (char *) slab -
+        offsetof(struct slab_header, data)) / sch->itemsize;
+
+    if (UNLIKELY(slab->slots == SLOTS_ALL_ZERO)) {
+        /* target slab is full, change state to partial */
+        slab->slots = SLOTS_FIRST << slot;
+
+        if (LIKELY(slab != sch->full)) {
+            if (LIKELY((slab->prev->next = slab->next) != NULL))
+                slab->next->prev = slab->prev;
+
+            slab->prev = NULL;
+        } else if (LIKELY((sch->full = sch->full->next) != NULL)) {
+            sch->full->prev = NULL;
+        }
+
+        slab->next = sch->partial;
+
+        if (LIKELY(sch->partial != NULL))
+            sch->partial->prev = slab;
+
+        sch->partial = slab;
+    } else if (UNLIKELY(ONE_USED_SLOT(slab->slots, sch->empty_slotmask))) {
+        /* target slab is partial and has only one filled slot */
+        if (UNLIKELY(slab->refcount == 1 || (slab->refcount == 0 &&
+            slab->page->refcount == 1))) {
+
+            /* unmap the whole page if this slab is the only partial one */
+            if (LIKELY(slab != sch->partial)) {
+                if (LIKELY((slab->prev->next = slab->next) != NULL))
+                    slab->next->prev = slab->prev;
+            } else if (LIKELY((sch->partial = sch->partial->next) != NULL)) {
+                sch->partial->prev = NULL;
+            }
+
+            void *const page = UNLIKELY(slab->refcount != 0) ? slab : slab->page;
+            const char *const page_end = (char *) page + sch->pages_per_alloc;
+            char found_head = 0;
+
+            union {
+                const char *c;
+                const struct slab_header *const s;
+            } s;
+
+            for (s.c = page; s.c != page_end; s.c += sch->slabsize) {
+                if (UNLIKELY(s.s == sch->empty))
+                    found_head = 1;
+                else if (UNLIKELY(s.s == slab))
+                    continue;
+                else if (LIKELY((s.s->prev->next = s.s->next) != NULL))
+                    s.s->next->prev = s.s->prev;
+            }
+
+            if (UNLIKELY(found_head && (sch->empty = sch->empty->next) != NULL))
+                sch->empty->prev = NULL;
+
+            sce_exe_free(slab->write_res, slab->exe_res);
+        } else {
+            slab->slots = sch->empty_slotmask;
+
+            if (LIKELY(slab != sch->partial)) {
+                if (LIKELY((slab->prev->next = slab->next) != NULL))
+                    slab->next->prev = slab->prev;
+
+                slab->prev = NULL;
+            } else if (LIKELY((sch->partial = sch->partial->next) != NULL)) {
+                sch->partial->prev = NULL;
+            }
+
+            slab->next = sch->empty;
+
+            if (LIKELY(sch->empty != NULL))
+                sch->empty->prev = slab;
+
+            sch->empty = slab;
+
+            UNLIKELY(slab->refcount != 0) ?
+                slab->refcount-- : slab->page->refcount--;
+        }
+    } else {
+        /* target slab is partial, no need to change state */
+        slab->slots |= SLOTS_FIRST << slot;
+    }
+}
+
+void slab_traverse(const struct slab_chain *const sch, void (*fn)(const void *))
+{
+    assert(sch != NULL);
+    assert(fn != NULL);
+    assert(slab_is_valid(sch));
+
+    const struct slab_header *slab;
+    const char *item, *end;
+    const size_t data_offset = offsetof(struct slab_header, data);
+
+    for (slab = sch->partial; slab; slab = slab->next) {
+        item = (const char *) slab + data_offset;
+        end = item + sch->itemcount * sch->itemsize;
+        uint64_t mask = SLOTS_FIRST;
+
+        do {
+            if (!(slab->slots & mask))
+                fn(item);
+
+            mask <<= 1;
+        } while ((item += sch->itemsize) != end);
+    }
+
+    for (slab = sch->full; slab; slab = slab->next) {
+        item = (const char *) slab + data_offset;
+        end = item + sch->itemcount * sch->itemsize;
+
+        do fn(item);
+        while ((item += sch->itemsize) != end);
+    }
+}
+
+void slab_destroy(const struct slab_chain *const sch)
+{
+    assert(sch != NULL);
+    assert(slab_is_valid(sch));
+
+    struct slab_header *const heads[] = {sch->partial, sch->empty, sch->full};
+    struct slab_header *pages_head = NULL, *pages_tail;
+
+    for (size_t i = 0; i < 3; ++i) {
+        struct slab_header *slab = heads[i];
+
+        while (slab != NULL) {
+            if (slab->refcount != 0) {
+                struct slab_header *const page = slab;
+                slab = slab->next;
+
+                if (UNLIKELY(pages_head == NULL))
+                    pages_head = page;
+                else
+                    pages_tail->next = page;
+
+                pages_tail = page;
+            } else {
+                slab = slab->next;
+            }
+        }
+    }
+
+    if (LIKELY(pages_head != NULL)) {
+        pages_tail->next = NULL;
+        struct slab_header *page = pages_head;
+
+        do {
+            sce_exe_free(page->write_res, page->exe_res);
+            page = page->next;
+        } while (page != NULL);
+    }
+}
diff --git a/lib/vita/slab.h b/lib/vita/slab.h
new file mode 100755
index 0000000..942eb70
--- /dev/null
+++ b/lib/vita/slab.h
@@ -0,0 +1,39 @@
+/* ref: https://github.com/bbu/userland-slab-allocator */
+
+#ifndef __GNUC__
+# error Can be compiled only with GCC.
+#endif
+
+#pragma once
+
+#include <stdint.h>
+#include <stddef.h>
+#include <psp2kern/types.h>
+
+extern const size_t slab_pagesize;
+
+struct slab_header {
+    struct slab_header *prev, *next;
+    uint64_t slots;
+    uintptr_t refcount;
+    struct slab_header *page;
+    SceUID write_res;
+    SceUID exe_res;
+    uintptr_t exe_data;
+    uint8_t data[] __attribute__((aligned(sizeof(void *))));
+};
+
+struct slab_chain {
+    size_t itemsize, itemcount;
+    size_t slabsize, pages_per_alloc;
+    uint64_t initial_slotmask, empty_slotmask;
+    uintptr_t alignment_mask;
+    struct slab_header *partial, *empty, *full;
+    SceUID pid;
+};
+
+void slab_init(struct slab_chain *, size_t);
+void *slab_alloc(struct slab_chain *, uintptr_t *);
+void slab_free(struct slab_chain *, const void *);
+void slab_traverse(const struct slab_chain *, void (*)(const void *));
+void slab_destroy(const struct slab_chain *);