Free current module.

free_current_module.cc

#include "free_current_module.h"

#include "blook/blook.h"

#include <Windows.h>
#include <thread>

struct ReallocInfo {
  void *new_addr;
  void *old_addr;
  size_t old_size;
  void *old_thread;
  void *main;
};

DWORD WINAPI free_realloc_stage_3(ReallocInfo *info) {
  SetThreadDescription(GetCurrentThread(), L"free_realloc_stage_3");
  // free the new module
  auto proc = GetCurrentProcess();
  VirtualFreeEx(proc, info->new_addr, 0, MEM_RELEASE);
  CreateThread(nullptr, 0,
              (LPTHREAD_START_ROUTINE)info->main, nullptr, 0, nullptr);
  return 0;
}

DWORD WINAPI free_realloc_stage_2(ReallocInfo *info) {
  SetThreadDescription(GetCurrentThread(), L"free_realloc_stage_2");
  auto proc = GetCurrentProcess();
  FreeLibrary((HMODULE)info->old_addr);
  FreeLibrary((HMODULE)info->old_addr);
  while (!VirtualAllocEx(proc, info->old_addr, info->old_size,
                         MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE))
    ;

  // copy the module back to the old address
  std::memcpy(info->old_addr, info->new_addr, info->old_size);
  // open a new thread in the old module
  auto offset = (char *)&free_realloc_stage_3 - (char *)info->new_addr +
                (char *)info->old_addr;

  CreateThread(nullptr, 0, (LPTHREAD_START_ROUTINE)offset, info, 0, nullptr);
  return 0;
}
void mb_shell::free_and_reallocate_current_module(auto (*main)()->void) {
  auto current_proc = GetCurrentProcess();

  auto current_mod = blook::Process::self()->module().value();
  auto mod_base = current_mod->base().data();
  auto mod_size = current_mod->size();

  // copy self to a new location first
  auto temp_mod_addr =
      VirtualAllocEx(current_proc, nullptr, mod_size, MEM_COMMIT | MEM_RESERVE,
                     PAGE_EXECUTE_READWRITE);
  if (!temp_mod_addr) {
    std::cerr << "Failed to allocate memory for new module" << std::endl;
    return;
  }

  std::memcpy(temp_mod_addr, mod_base, mod_size);

  // start a new thread in the new module
  auto offset =
      (char *)&free_realloc_stage_2 - (char *)mod_base + (char *)temp_mod_addr;
  ReallocInfo *info = new ReallocInfo{temp_mod_addr, mod_base, mod_size,
                                      GetCurrentThread(), (void *)main};

  CreateThread(nullptr, 0, (LPTHREAD_START_ROUTINE)offset, info, 0, nullptr);
}