#include <ksymtab.h>




#define KFUNC_TABLE_ADDRESS 0xC0101000

#define KSYMTAB_MAX 0x8086FF

#define IS_MODULE_FUNC(id)  ((id) & 0x80000000)
#define GET_MODULE_ID(id)   (((id) >> 20) & 0x7FF)
#define GET_FUNC_ID(id)     ((id) & 0xFFFFF)
#define EXISTS(id)          ((id) > 0x0)

#define MAKE_KERNEL_FUNC(id) ((id) & 0x7FFFFFFF)

#define MAKE_MODULE_FUNC(mid, fid) (0x80000000 | ((mid) << 20) | ((fid) & 0xFFFFF))

kfunc_t* kfunc_table = (kfunc_t*)KFUNC_TABLE_ADDRESS;

static uint32_t ktab_size = 0;


uint64_t kfunc_call(kfunc_t* func, uint32_t a, uint32_t b, uint32_t c, uint32_t d)
{
  uint32_t eax_ret, edx_ret;

  asm volatile (
    "push %[d]\n\t"
    "push %[c]\n\t"
    "push %[b]\n\t"
    "push %[a]\n\t"
    "call *%[fn]\n\t"
    "add $16, %%esp\n\t" /* clean up stack (4 args * 4 bytes) */
    : "=a"(eax_ret), "=d"(edx_ret)
    : [a]"r"(a), [b]"r"(b), [c]"r"(c), [d]"r"(d), [fn]"r"(func->addr)
    : "memory"
  );

  return ((uint64_t)edx_ret << 32) | eax_ret;
}



uint64_t call_kfunc_by_id(uint32_t id, uint32_t a, uint32_t b, uint32_t c, uint32_t d)
{
  for (int i = 0; i < (int)ktab_size; i++)
  {
    if (kfunc_table[i].id == id)
    {
      if (kfunc_table[i].addr == 0x0)
      {
        return -1;
      }
      return kfunc_call(&kfunc_table[i], a, b, c, d);
    }
  }
  return -1;
}


uint32_t add_kfunc(void* addr, bool module, uint16_t module_id, uint32_t function_id)
{
  if (ktab_size >= KSYMTAB_MAX)
  {
    return 0xFFFFFFFF;
  }

  kfunc_t kf = make_kfunc(addr, module, module_id, function_id);
  kfunc_table[ktab_size] = kf;

  ktab_size++;

  return kf.id;
}

/*
  Constructs and returns a kfunc_t:
  - 'addr' is the address of the function (can be NULL/0)
  - 'module' indicates if it's a module/driver function
  - 'module_id' is used only if 'module' is true (11 bits max)
  - 'function_id':
    - 20 bits if module is true
    - 31 bits if kernel function
*/
kfunc_t make_kfunc(void* addr, bool module, uint16_t module_id, uint32_t function_id)
{
  uint32_t id;

  if (module)
  {
    id = 0x80000000 | ((module_id & 0x7FF) << 20) | (function_id & 0xFFFFF);
  }
  else
  {
    if (function_id == UINT32_MAX)
    {
      function_id = ktab_size;
    }
    id = function_id & 0x7FFFFFFF;
  }

  kfunc_t result = {
    .id = id,
    .addr = (uint32_t)(uintptr_t)addr
  };

  return result;
}

kfunc_t* find_kfunc(uint32_t id)
{
  for (uint32_t i = 0; i < ktab_size; i++)
  {
    if (kfunc_table[i].id == id)
    {
      return &kfunc_table[i];
    }
  }
  return NULL;
}