patches.c

/* patches.c -- main patch system
 *
 * Copyright (C) 2016 Yifan Lu
 *
 * This software may be modified and distributed under the terms
 * of the MIT license.  See the LICENSE file for details.
 */
#include <psp2kern/types.h>
#include <psp2kern/kernel/cpu.h>
#include <psp2kern/kernel/sysmem.h>
#include <psp2kern/kernel/threadmgr.h>
#include <string.h>
#include "error.h"
#include "taihen_internal.h"
#include "patches.h"
#include "proc_map.h"
#include "slab.h"
#include "substitute/lib/substitute.h"

#define PATCHES_POOL_SIZE 0x10000

#define NUM_PROC_MAP_BUCKETS 16

#define FUNC_TO_UINTPTR_T(x) (((uintptr_t)(x))&0xFFFFFFFE)

#define MEM_SHARED_START ((void*)0xE0000000)

SceUID g_patch_pool;

static tai_proc_map_t *g_map;

static SceUID g_hooks_lock;

static SceClass g_taihen_class;

static int init_patch(void *dat) {
  tai_patch_t *patch;

  patch = (tai_patch_t *)dat;
  LOG("init of: %p", patch);
  return 0;
}

static int free_patch(void *dat) {
  tai_patch_t *patch;

  patch = (tai_patch_t *)dat;
  LOG("cleanup of: %p", patch);
  return 0;
}

int patches_init(void) {
  SceKernelMemPoolCreateOpt opt;
  int ret;

  memset(&opt, 0, sizeof(opt));
  opt.size = sizeof(opt);
  opt.uselock = 1;
  g_patch_pool = sceKernelMemPoolCreate("tai_patches", PATCHES_POOL_SIZE, &opt);
  LOG("sceKernelMemPoolCreate(tai_patches): 0x%08X", g_patch_pool);
  if (g_patch_pool < 0) {
    return g_patch_pool;
  }
  g_map = proc_map_alloc(NUM_PROC_MAP_BUCKETS);
  if (g_map == NULL) {
    LOG("Failed to create proc map.");
    return TAI_ERROR_SYSTEM;
  }
  g_hooks_lock = sceKernelCreateMutexForKernel("tai_hooks_lock", SCE_KERNEL_MUTEX_ATTR_RECURSIVE, 0, NULL);
  LOG("sceKernelCreateMutexForKernel(tai_hooks_lock): 0x%08X", g_hooks_lock);
  if (g_hooks_lock < 0) {
    return g_hooks_lock;
  }
  ret = sceKernelCreateClass(&g_taihen_class, "taiHENClass", sceKernelGetUidClass(), sizeof(tai_patch_t), init_patch, free_patch);
  LOG("sceKernelCreateClass(taiHENClass): 0x%08X", ret);
  if (ret < 0) {
    return ret;
  }
  return TAI_SUCCESS;
}

void patches_deinit(void) {
  LOG("Cleaning up patches subsystem.");
  // TODO: Find out how to clean up class
  sceKernelDeleteMutexForKernel(g_hooks_lock);
  sceKernelMemPoolDestroy(g_patch_pool);
  proc_map_free(g_map);
  g_map = NULL;
  g_patch_pool = 0;
  g_hooks_lock = 0;
}

static inline void hex_dump(uintptr_t paddr, const char *addr, unsigned int size)
{
    unsigned int i;
    for (i = 0; i < (size >> 4); i++)
    {
        LOG("0x%08X: %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X\n", 
            paddr, 
            addr[0], addr[1], addr[2], addr[3], addr[4], addr[5], addr[6], addr[7], addr[8], 
            addr[9], addr[10], addr[11], addr[12], addr[13], addr[14], addr[15]
        );
        paddr += 0x10;
        addr += 0x10;
    }
}

#ifdef __arm__

void cache_flush(SceUID pid, uintptr_t vma, size_t len) {
  uintptr_t vma_align;
  int flags;
  int my_context[3];
  int ret;
  int *other_context;
  int dacr;

  vma_align = vma & ~0x1F;
  len = ((vma + len + 0x1F) & ~0x1F) - vma_align;
  LOG("cache flush: vma %p, vma_align %p, len %x", vma, vma_align, len);

  if (pid == KERNEL_PID) {
    sceKernelCpuDcacheFlush((void *)vma_align, len);
    sceKernelCpuIcacheAndL2Flush((void *)vma_align, len);
  } else {
    // TODO: Take care of SHARED_PID
    flags = sceKernelCpuDisableInterrupts();
    sceKernelCpuSaveContext(my_context);
    ret = sceKernelGetPidContext(pid, &other_context);
    if (ret >= 0) {
      sceKernelCpuRestoreContext(other_context);
      asm volatile ("mrc p15, 0, %0, c3, c0, 0" : "=r" (dacr));
      asm volatile ("mcr p15, 0, %0, c3, c0, 0" :: "r" (0x15450FC3));
      sceKernelCpuDcacheFlush((void *)vma_align, len);
      sceKernelCpuIcacheAndL2Flush((void *)vma_align, len);
      hex_dump(vma_align, (char *)vma_align, len);
      asm volatile ("mcr p15, 0, %0, c3, c0, 0" :: "r" (dacr));
    }
    sceKernelCpuRestoreContext(my_context);
    sceKernelCpuEnableInterrupts(flags);
    LOG("sceKernelSwitchVmaForPid(%d): 0x%08X\n", pid, ret);
  }
  asm volatile ("isb" ::: "memory");
}
#endif

struct hook_args {
  SceUID pid;
  struct substitute_function_hook *hook;
  struct substitute_function_hook_record **saved;
};

static int do_hooking(void *args) {
  int flags;
  int ret;
  int my_context[3];
  int *other_context;
  int dacr;
  struct hook_args *uargs = (struct hook_args *)args;

  if (uargs->pid == KERNEL_PID) {
    ret = substitute_hook_functions(uargs->hook, 1, uargs->saved, SUBSTITUTE_RELAXED);
  } else {
    flags = sceKernelCpuDisableInterrupts();
    sceKernelCpuSaveContext(my_context);
    ret = sceKernelGetPidContext(uargs->pid, &other_context);
    if (ret >= 0) {
      sceKernelCpuRestoreContext(other_context);
      asm volatile ("mrc p15, 0, %0, c3, c0, 0" : "=r" (dacr));
      asm volatile ("mcr p15, 0, %0, c3, c0, 0" :: "r" (0x15450FC3));
      ret = substitute_hook_functions(uargs->hook, 1, uargs->saved, SUBSTITUTE_RELAXED);
      asm volatile ("mcr p15, 0, %0, c3, c0, 0" :: "r" (dacr));
    }
    sceKernelCpuRestoreContext(my_context);
    sceKernelCpuEnableInterrupts(flags);
  }
  return ret;
}

static int do_unhooking(void *saved) {
  return substitute_free_hooks((struct substitute_function_hook_record *)saved, 1);
}

static int tai_hook_function(struct slab_chain *slab, void *target_func, const void *src_func, void **old, void **saved) {
  struct hook_args uargs;
  struct substitute_function_hook hook;
  int ret;

  if (target_func == src_func) {
    LOG("no hook needed");
    return TAI_SUCCESS; // no need for hook
  }

  hook.function = target_func;
  hook.replacement = (void *)src_func;
  hook.old_ptr = old;
  hook.options = 0;
  hook.opt = slab;
  LOG("Calling substitute_hook_functions");
  // TODO: Take care of SHARED_PID
  uargs.pid = slab->pid;
  uargs.hook = &hook;
  uargs.saved = (struct substitute_function_hook_record **)saved;
  ret = sceKernelRunWithStack(0x4000, do_hooking, &uargs);
  LOG("Done hooking");
  if (ret != SUBSTITUTE_OK) {
    LOG("libsubstitute error: %s", substitute_strerror(ret));
    return TAI_ERROR_HOOK_ERROR;
  }
  return TAI_SUCCESS;
}

static int tai_unhook_function(void *saved) {
  int ret;
  LOG("Calling substitute_free_hooks");
  ret = sceKernelRunWithStack(0x4000, do_unhooking, saved);
  if (ret != SUBSTITUTE_OK) {
    LOG("libsubstitute error: %s", substitute_strerror(ret));
    return TAI_ERROR_HOOK_ERROR;
  }
  return TAI_SUCCESS;
}

static int tai_force_memcpy(SceUID dst_pid, void *dst, const void *src, size_t size) {
  int ret;
  if (dst_pid == KERNEL_PID) {
      ret = sceKernelCpuUnrestrictedMemcpy(dst, src, size);
      LOG("sceKernelCpuUnrestrictedMemcpy(%p, %p, 0x%08X): 0x%08X", dst, src, size, ret);
  } else {
      ret = sceKernelRxMemcpyKernelToUserForPid(dst_pid, (uintptr_t)dst, src, size);
      LOG("sceKernelRxMemcpyKernelToUserForPid(%x, %p, %p, 0x%08X): 0x%08X", dst_pid, dst, src, size, ret);
  }
  cache_flush(dst_pid, (uintptr_t)dst, size);
  return ret;
}

int tai_memcpy_to_kernel(SceUID src_pid, void *dst, const char *src, size_t size) {
  int ret;
  if (src_pid == KERNEL_PID) {
    memcpy(dst, src, size);
    LOG("memcpy(%p, %p, 0x%08X)", dst, src, size);
  } else {
    ret = sceKernelMemcpyUserToKernelForPid(src_pid, dst, (uintptr_t)src, size);
    LOG("sceKernelMemcpyUserToKernelForPid(%x, %p, %p, 0x%08X): 0x%08X", src_pid, dst, src, size, ret);
  }
  return 0;
}

static int hooks_add_hook(tai_hook_list_t *hooks, tai_hook_t *item) {
  tai_hook_t *head;
  int ret;

  LOG("Adding hook %p to chain %p", item, hooks);
  sceKernelLockMutexForKernel(g_hooks_lock, 1, NULL);
  if (hooks->head == NULL) { // first hook for this list
    ret = tai_hook_function(item->patch->slab, hooks->func, item->u.func, &hooks->old, &hooks->saved);
    if (ret >= 0) {
      hooks->head = item;
      item->next = NULL;
      item->u.next = (uintptr_t)NULL;
      item->u.old = hooks->old;
      cache_flush(item->patch->pid, slab_getmirror(item->patch->slab, item), sizeof(tai_hook_t));
    } else {
      LOG("Hook failed, do not add to chain");
    }
  } else {
    head = hooks->head;
    item->next = head->next;
    item->u.next = head->u.next;
    item->u.old = hooks->old;
    head->next = item;
    head->u.next = slab_getmirror(item->patch->slab, item);
    LOG("Added hook to existing chain %p", head);
    // flush cache for head + item, which were modified
    cache_flush(item->patch->pid, slab_getmirror(item->patch->slab, head), sizeof(tai_hook_t));
    cache_flush(item->patch->pid, head->u.next, sizeof(tai_hook_t));
    ret = 1;
  }
  sceKernelUnlockMutexForKernel(g_hooks_lock, 1);

  return ret;
}

static int hooks_remove_hook(tai_hook_list_t *hooks, tai_hook_t *item) {
  tai_hook_t **cur;
  uintptr_t tmp, *cur_user;
  int ret;

  LOG("Removing hook %p for %p", item, hooks);
  sceKernelLockMutexForKernel(g_hooks_lock, 1, NULL);
  if (hooks->head == item) { // first hook for this list
    // we must remove the patch
    tai_unhook_function(hooks->saved);
    hooks->saved = NULL;
    // set head to the next item
    hooks->head = item->next;
    if (hooks->head != NULL) {
      // add a patch to the new head
      ret = tai_hook_function(item->patch->slab, hooks->func, hooks->head->u.func, &hooks->old, &hooks->saved);
      // update the old pointers
      for (cur = &hooks->head; *cur != NULL; cur = &(*cur)->next) {
        (*cur)->u.old = hooks->old;
      }
      // clear cache of mirror for the last item since it uses the old pointer
      cache_flush(item->patch->pid, (uintptr_t)cur - offsetof(tai_hook_t, next), sizeof(tai_hook_t));
    } else {
      ret = 0;
    }
  } else {
    cur = &hooks->head;
    cur_user = &tmp;
    ret = -1;
    while (1) {
      if (*cur) {
        if (*cur == item) {
          *cur = item->next; // remove from list
          *cur_user = item->u.next;
          // clear cache since pointers were changed
          cache_flush(item->patch->pid, (uintptr_t)cur - offsetof(tai_hook_t, next), sizeof(tai_hook_t));
          ret = 0;
          break;
        } else {
          cur = &(*cur)->next;
          cur_user = &(*cur)->u.next;
        }
      } else {
        break;
      }
    }
  }
  sceKernelUnlockMutexForKernel(g_hooks_lock, 1);
  return ret;
}

SceUID tai_hook_func_abs(tai_hook_ref_t *p_hook, SceUID pid, void *dest_func, const void *hook_func) {
  SceCreateUidObjOpt opt;
  tai_patch_t *patch, *tmp;
  tai_hook_t *hook;
  int ret;
  struct slab_chain *slab;
  uintptr_t exe_addr;

  LOG("Hooking %p to %p for pid %x", hook_func, dest_func, pid);
  if (hook_func >= MEM_SHARED_START) {
    if (pid == KERNEL_PID) {
      return TAI_ERROR_INVALID_KERNEL_ADDR; // invalid hook address
    } else {
      return TAI_ERROR_NOT_IMPLEMENTED; // TODO: add support for this
    }
  }

  hook = NULL;
  if (pid == KERNEL_PID) {
    ret = sceKernelCreateUidObj(&g_taihen_class, "tai_patch_hook", NULL, (SceObjectBase **)&patch);
  } else {
    memset(&opt, 0, sizeof(opt));
    opt.flags = 8;
    opt.pid = pid;
    ret = sceKernelCreateUidObj(&g_taihen_class, "tai_patch_hook_user", &opt, (SceObjectBase **)&patch);
  }
  LOG("sceKernelCreateUidObj(tai_patch_hook): 0x%08X, %p", ret, patch);
  if (ret < 0) {
    return ret;
  }

  sceKernelLockMutexForKernel(g_hooks_lock, 1, NULL);
  patch->type = HOOKS;
  patch->uid = ret;
  patch->pid = pid;
  patch->addr = FUNC_TO_UINTPTR_T(dest_func);
  patch->size = FUNC_SAVE_SIZE;
  patch->next = NULL;
  patch->data.hooks.func = dest_func;
  patch->data.hooks.saved = NULL;
  patch->data.hooks.head = NULL;
  if (proc_map_try_insert(g_map, patch, &tmp) < 1) {
    ret = sceKernelDeleteUid(patch->uid);
    LOG("sceKernelDeleteUid(old): 0x%08X", ret);
    if (tmp == NULL || tmp->type != HOOKS) {
      // error
      LOG("this hook overlaps an existing hook");
      ret = TAI_ERROR_PATCH_EXISTS;
      goto err;
    } else {
      // we have an existing patch
      LOG("found existing patch %p, discarding %p", tmp, patch);
      patch = tmp;
    }
  }

  hook = slab_alloc(patch->slab, &exe_addr);
  if (hook == NULL) {
    ret = -1;
    goto err;
  }
  hook->u.func = (void *)hook_func;
  hook->patch = patch;

  ret = hooks_add_hook(&patch->data.hooks, hook);
  if (ret < 0 && patch->data.hooks.head == NULL) {
    LOG("failed to add hook and patch is now empty, freeing hook %p", hook);
    slab_free(patch->slab, hook);
    hook = NULL;
    proc_map_remove(g_map, patch);
    sceKernelDeleteUid(patch->uid);
    patch = NULL;
  } else if (ret >= 0) {
    ret = patch->uid;
    *p_hook = slab_getmirror(patch->slab, hook);
  }

err:
  // error and we have allocated a hook
  if (ret < 0 && patch && hook) {
    LOG("freeing hook %p", hook);
    slab_free(patch->slab, hook);
  }

  sceKernelUnlockMutexForKernel(g_hooks_lock, 1);

  return ret;
}

int tai_hook_release(SceUID uid, tai_hook_ref_t hook_ref) {
  tai_hook_t **cur, *hook;
  tai_patch_t *patch;
  struct slab_chain *slab;
  int ret;

  ret = sceKernelGetObjForUid(uid, &g_taihen_class, (SceObjectBase **)&patch);
  LOG("sceKernelGetObjForUid(%x): 0x%08X", uid, ret);
  if (ret < 0) {
    return ret;
  }
  sceKernelLockMutexForKernel(g_hooks_lock, 1, NULL);
  slab = patch->slab;
  for (cur = &patch->data.hooks.head; *cur != NULL; cur = &(*cur)->next) {
    if (slab_getmirror(slab, *cur) == hook_ref) {
      hook = *cur;
      LOG("Found hook %p for ref %p", hook, hook_ref);
      ret = hooks_remove_hook(&patch->data.hooks, hook);
      *cur = hook->next;
      LOG("freeing hook");
      slab_free(slab, hook);
      if (patch->data.hooks.head == NULL) {
        LOG("patch is now empty, freeing it");
        proc_map_remove(g_map, patch);
        sceKernelDeleteUid(patch->uid);
      }
      ret = TAI_SUCCESS;
      goto end;
    }
  }
  LOG("Cannot find hook for uid %x ref %p", uid, hook_ref);
  ret = TAI_ERROR_NOT_FOUND;
end:
  sceKernelUnlockMutexForKernel(g_hooks_lock, 1);

  return ret;
}

SceUID tai_inject_abs(SceUID pid, void *dest, const void *src, size_t size) {
  tai_patch_t *patch, *tmp;
  void *saved;
  int ret;

  // TODO: Check that dest is not inside our slab structure... that could corrupt kernel code

  LOG("Injecting %p with %p for size 0x%08X at pid %x", dest, src, size, pid);
  ret = sceKernelCreateUidObj(&g_taihen_class, "tai_patch_inject", NULL, (SceObjectBase **)&patch);
  LOG("sceKernelCreateUidObj(tai_patch_inject): 0x%08X, %p", ret, patch);
  if (ret < 0) {
    return ret;
  }

  saved = sceKernelMemPoolAlloc(g_patch_pool, size);
  LOG("sceKernelMemPoolAlloc(g_patch_pool, 0x%08X): %p", size, saved);
  if (saved == NULL) {
    sceKernelDeleteUid(ret);
    return TAI_ERROR_MEMORY;
  }

  // try to save old data
  if (tai_memcpy_to_kernel(pid, saved, dest, size) < 0) {
    LOG("Invalid address for memcpy");
    sceKernelDeleteUid(ret);
    sceKernelMemPoolFree(g_patch_pool, saved);
    return TAI_ERROR_INVALID_ARGS;
  }

  sceKernelLockMutexForKernel(g_hooks_lock, 1, NULL);
  patch->type = INJECTION;
  patch->uid = ret;
  patch->pid = pid;
  patch->addr = (uintptr_t)dest;
  patch->size = size;
  patch->next = NULL;
  patch->data.inject.saved = saved;
  patch->data.inject.size = size;
  patch->data.inject.patch = patch;
  if (proc_map_try_insert(g_map, patch, &tmp) < 1) {
    ret = TAI_ERROR_PATCH_EXISTS;
  } else {
    ret = tai_force_memcpy(pid, dest, src, size);
  }

  if (ret < 0) {
    sceKernelDeleteUid(patch->uid);
    sceKernelMemPoolFree(g_patch_pool, saved);
  } else {
    ret = patch->uid;
  }

  sceKernelUnlockMutexForKernel(g_hooks_lock, 1);

  return ret;
}

int tai_inject_release(SceUID uid) {
  tai_inject_t *inject;
  tai_patch_t *patch;
  void *saved;
  void *dest;
  size_t size;
  int ret;
  SceUID pid;

  ret = sceKernelGetObjForUid(uid, &g_taihen_class, (SceObjectBase **)&patch);
  LOG("sceKernelGetObjForUid(%x): 0x%08X", uid, ret);
  if (ret < 0) {
    return ret;
  }
  if (patch->type != INJECTION || patch->uid != uid) {
    LOG("internal error: trying to free an invalid injection");
    return TAI_ERROR_SYSTEM;
  }
  inject = &patch->data.inject;
  LOG("Releasing injection %p for patch %p", inject, patch);
  sceKernelLockMutexForKernel(g_hooks_lock, 1, NULL);
  pid = patch->pid;
  dest = (void *)patch->addr;
  saved = inject->saved;
  size = inject->size;
  if (!proc_map_remove(g_map, patch)) {
    LOG("internal error, cannot remove patch from proc_map");
    ret = TAI_ERROR_SYSTEM;
  } else {
    ret = tai_force_memcpy(pid, dest, saved, size);
    sceKernelMemPoolFree(g_patch_pool, saved);
    sceKernelDeleteUid(patch->uid);
  }
  sceKernelUnlockMutexForKernel(g_hooks_lock, 1);

  return ret;
}

int tai_try_cleanup_process(SceUID pid) {
  tai_patch_t *patch, *next;
  LOG("Calling patches cleanup for pid %x", pid);
  sceKernelLockMutexForKernel(g_hooks_lock, 1, NULL);
  if (proc_map_remove_all_pid(g_map, pid, &patch) > 0) {
    while (patch != NULL) {
      next = patch->next;
      if (patch->type == INJECTION) {
        LOG("freeing injection saved data");
        sceKernelMemPoolFree(g_patch_pool, patch->data.inject.saved);
      } else if (patch->type == HOOKS) {
        LOG("freeing hook saved data");
        free(patch->data.hooks.saved);
      }
      LOG("deleting patch: %x", patch->uid);
      sceKernelDeleteUid(patch->uid);
      patch = next;
    }
  }
  sceKernelUnlockMutexForKernel(g_hooks_lock, 1);
  return 0;
}