debugger.cpp

#include "debugger.h"

namespace {
struct reg_descriptor
{
  std::string name;
  size_t offset;
};

const std::vector<reg_descriptor> g_registers = {
  {"rax", offsetof(user_regs_struct, rax)}, {"rbx", offsetof(user_regs_struct, rbx)},
  {"rcx", offsetof(user_regs_struct, rcx)}, {"rdx", offsetof(user_regs_struct, rdx)},
  {"rdi", offsetof(user_regs_struct, rdi)}, {"rsi", offsetof(user_regs_struct, rsi)},
  {"rbp", offsetof(user_regs_struct, rbp)}, {"rsp", offsetof(user_regs_struct, rsp)},
  {"r8", offsetof(user_regs_struct, r8)},   {"r9", offsetof(user_regs_struct, r9)},
  {"r10", offsetof(user_regs_struct, r10)}, {"r11", offsetof(user_regs_struct, r11)},
  {"r12", offsetof(user_regs_struct, r12)}, {"r13", offsetof(user_regs_struct, r13)},
  {"r14", offsetof(user_regs_struct, r14)}, {"r15", offsetof(user_regs_struct, r15)},
  {"rip", offsetof(user_regs_struct, rip)}, {"eflags", offsetof(user_regs_struct, eflags)}};

} // namespace

Debugger::Debugger()
  : state(NOT_LOADED)
  , pid(-1)
  , entry_point(0)
{
  if (cs_open(CS_ARCH_X86, CS_MODE_64, &capstone_handle) != CS_ERR_OK) {
    std::cerr << "** fatal: failed to initialize capstone.\n";
    exit(1);
  }
  cs_option(capstone_handle, CS_OPT_DETAIL, CS_OPT_ON);
}

Debugger::~Debugger() { cs_close(&capstone_handle); }

bool Debugger::parse_elf(const std::string &path)
{
  int fd = open(path.c_str(), O_RDONLY);
  if (fd < 0)
    return false;

  Elf64_Ehdr ehdr;
  if (read(fd, &ehdr, sizeof(ehdr)) != sizeof(ehdr)) {
    close(fd);
    return false;
  }

  entry_point = ehdr.e_entry;
  lseek(fd, ehdr.e_phoff, SEEK_SET);
  for (int i = 0; i < ehdr.e_phnum; i++) {
    Elf64_Phdr phdr;
    read(fd, &phdr, sizeof(phdr));
    if (phdr.p_type == PT_LOAD && (phdr.p_flags & PF_X)) {
      text_segment.start = phdr.p_vaddr;
      text_segment.end = phdr.p_vaddr + phdr.p_memsz;
      break;
    }
  }
  close(fd);
  return true;
}

void Debugger::load_program(const std::string &path)
{
  if (state != NOT_LOADED) {
    std::cerr << "** program already loaded.\n";
    return;
  }

  auto s = std::filesystem::status(path);
  auto p = s.permissions();
  bool is_exec = (p & (std::filesystem::perms::owner_exec | std::filesystem::perms::group_exec |
                       std::filesystem::perms::others_exec)) != std::filesystem::perms::none;

  if (!std::filesystem::exists(s) || !std::filesystem::is_regular_file(s) || !is_exec) {
    std::cerr << "** error: '" << path << "' is invalid or not executable.\n";
    return;
  }

  program_path = path;
  if (!parse_elf(path)) {
    std::cerr << "** failed to load program.\n";
    return;
  }
  std::cout << "** program '" << path << "' loaded. entry point 0x" << std::hex << entry_point << "\n";
  state = LOADED;
}

void Debugger::load_initial_program(const std::string &path)
{
  program_path = path;
  if (!parse_elf(path)) {
    std::cerr << "** failed to load program.\n";
    exit(1);
  }
  std::cout << "** program '" << path << "' loaded. entry point 0x" << std::hex << entry_point << "\n";
  state = LOADED;
}

void Debugger::start_program()
{
  if (state == NOT_LOADED) {
    std::cerr << "** no program loaded.\n";
    return;
  }
  if (state == RUNNING) {
    std::cerr << "** program already running.\n";
    return;
  }
  pid = fork();
  if (pid == 0) {
    ptrace(PTRACE_TRACEME, 0, nullptr, nullptr);
    execl(program_path.c_str(), program_path.c_str(), nullptr);
    perror("execl");
    exit(1);
  } else if (pid > 0) {
    int status;
    waitpid(pid, &status, 0);
    state = RUNNING;
    std::cout << "** pid " << pid << "\n";
    apply_all_breakpoints();
  } else {
    perror("fork");
  }
}

void Debugger::run_program()
{
  if (state == NOT_LOADED) {
    std::cerr << "** no program loaded.\n";
    return;
  }
  if (state == RUNNING) {
    std::cout << "** program " << program_path << " is already running\n";
    continue_execution();
    return;
  }
  start_program();
  if (state == RUNNING) {
    continue_execution();
  }
}

void Debugger::wait_for_signal()
{
  int status;
  waitpid(pid, &status, 0);
  if (WIFEXITED(status)) {
    std::cout << "** child process " << pid << " terminiated normally (code " << WEXITSTATUS(status) << ")\n";
    state = LOADED;
    pid = -1;
  } else if (WIFSTOPPED(status)) {
    int sig = WSTOPSIG(status);
    if (sig == SIGTRAP) {
      handle_breakpoint();
    } else {
      std::cout << "** child process " << pid << " stopped by signal " << sig << " (" << strsignal(sig) << ")\n";
    }
  } else if (WIFSIGNALED(status)) {
    std::cout << "** child process " << pid << " killed by signal " << WTERMSIG(status) << "\n";
    state = LOADED;
  }
}

void Debugger::handle_breakpoint()
{
  struct user_regs_struct regs;
  ptrace(PTRACE_GETREGS, pid, nullptr, &regs);
  uintptr_t pc = regs.rip - 1;
  auto it = addr_to_bp.find(pc);
  if (it != addr_to_bp.end()) {
    regs.rip = pc;
    ptrace(PTRACE_SETREGS, pid, nullptr, &regs);
    print_instruction_at(pc);
  }
}

void Debugger::apply_all_breakpoints()
{
  for (auto &[id, bp] : breakpoints) {
    uint64_t code = ptrace(PTRACE_PEEKTEXT, pid, (void *)bp.addr, nullptr);
    ptrace(PTRACE_POKETEXT, pid, (void *)bp.addr, (code & ~0xFFULL) | 0xCC);
  }
}

void Debugger::continue_execution()
{
  if (state != RUNNING) {
    std::cerr << "** program not running.\n";
    return;
  }
  struct user_regs_struct regs;
  ptrace(PTRACE_GETREGS, pid, nullptr, &regs);
  uintptr_t pc = regs.rip;
  auto it = addr_to_bp.find(pc);
  if (it != addr_to_bp.end()) {
    auto &bp = it->second;
    uint64_t code = ptrace(PTRACE_PEEKTEXT, pid, (void *)pc, nullptr);
    ptrace(PTRACE_POKETEXT, pid, (void *)pc, (code & ~0xFFULL) | (bp.original_data & 0xFF));
    ptrace(PTRACE_SINGLESTEP, pid, nullptr, nullptr);
    int status;
    waitpid(pid, &status, 0);
    ptrace(PTRACE_POKETEXT, pid, (void *)pc, (code & ~0xFFULL) | 0xCC);
  }
  ptrace(PTRACE_CONT, pid, nullptr, nullptr);
  wait_for_signal();
}

void Debugger::set_breakpoint(uintptr_t addr)
{
  if (state != RUNNING) {
    std::cerr << "** program not running.\n";
    return;
  }

  if (addr_to_bp.find(addr) != addr_to_bp.end()) {
    std::cout << "** breakpoint already exists at 0x" << std::hex << addr << "\n";
    return;
  }

  uint64_t original_data = ptrace(PTRACE_PEEKTEXT, pid, (void *)addr, nullptr);
  uint64_t trap_data = (original_data & ~0xFFULL) | 0xCC;
  ptrace(PTRACE_POKETEXT, pid, (void *)addr, trap_data);
  addr_to_bp[addr] = {addr, original_data};
  breakpoints[breakpoint_id_counter] = {addr, original_data};
  breakpoint_id_counter++;
}

void Debugger::delete_breakpoint(int id)
{
  if (state != RUNNING) {
    std::cerr << "** program not running.\n";
    return;
  }
  auto it = breakpoints.find(id);
  if (it == breakpoints.end()) {
    std::cerr << "** breakpoint " << id << " not found.\n";
    return;
  }
  uintptr_t addr = it->second.addr;
  uint64_t original_byte = it->second.original_data & 0xFF;
  uint64_t current_code = ptrace(PTRACE_PEEKTEXT, pid, (void *)addr, nullptr);
  if ((current_code & 0xFF) == 0xCC) {
    uint64_t restored_code = (current_code & ~0xFFULL) | original_byte;
    ptrace(PTRACE_POKETEXT, pid, (void *)addr, (void *)restored_code);
  }
  addr_to_bp.erase(addr);
  breakpoints.erase(it);
  std::cout << "** breakpoint " << id << " deleted at 0x" << std::hex << addr << "\n";
}

void Debugger::list_breakpoints()
{
  for (const auto &[id, bp] : breakpoints) {
    std::cout << "  " << std::dec << id << ": " << std::hex << bp.addr << "\n";
  }
}

void Debugger::print_instruction_at(uintptr_t pc)
{
  cs_insn *insn;
  uint8_t code_buf[15];
  for (int i = 0; i < 15; i += 8) {
    uint64_t word = ptrace(PTRACE_PEEKTEXT, pid, (void *)(pc + i), nullptr);
    int copy_len = (15 - i > 8) ? 8 : 15 - i;
    memcpy(code_buf + i, &word, copy_len);
  }
  for (int i = 0; i < 15; i++) {
    auto it = addr_to_bp.find(pc + i);
    if (it != addr_to_bp.end())
      code_buf[i] = it->second.original_data & 0xFF;
  }
  size_t count = cs_disasm(capstone_handle, code_buf, 15, pc, 1, &insn);
  if (count > 0) {
    std::cout << "** breakpoint @ " << std::hex << std::setw(11) << std::right << pc << ": ";
    std::stringstream ss;
    for (int i = 0; i < insn[0].size; i++)
      ss << std::hex << std::setw(2) << std::setfill('0') << (int)insn[0].bytes[i] << " ";
    std::cout << std::left << std::setw(31) << std::setfill(' ') << ss.str();
    std::cout << std::left << std::setw(10) << insn[0].mnemonic << insn[0].op_str << "\n";
    cs_free(insn, count);
  }
}

void Debugger::show_breakpoint_info(uintptr_t pc)
{
  if (addr_to_bp.count(pc))
    print_instruction_at(pc);
}

void Debugger::single_step()
{
  if (state != RUNNING) {
    std::cerr << "** program not running.\n";
    return;
  }
  struct user_regs_struct regs;
  ptrace(PTRACE_GETREGS, pid, nullptr, &regs);
  uintptr_t pc = regs.rip;
  auto it = addr_to_bp.find(pc);
  if (it != addr_to_bp.end()) {
    auto &bp = it->second;
    uint64_t code = ptrace(PTRACE_PEEKTEXT, pid, (void *)pc, nullptr);
    ptrace(PTRACE_POKETEXT, pid, (void *)pc, (code & ~0xFFULL) | (bp.original_data & 0xFF));
    ptrace(PTRACE_SINGLESTEP, pid, nullptr, nullptr);
    int status;
    waitpid(pid, &status, 0);
    ptrace(PTRACE_POKETEXT, pid, (void *)pc, (code & ~0xFFULL) | 0xCC);
  } else {
    ptrace(PTRACE_SINGLESTEP, pid, nullptr, nullptr);
    int status;
    waitpid(pid, &status, 0);
  }
  ptrace(PTRACE_GETREGS, pid, nullptr, &regs);
  uintptr_t new_pc = regs.rip;
  if (addr_to_bp.count(new_pc))
    print_instruction_at(new_pc);
}

uint64_t Debugger::get_register_value(const std::string &reg_name, struct user_regs_struct &regs)
{
  for (const auto &rd : g_registers) {
    if (rd.name == reg_name) {
      return *(reinterpret_cast<uint64_t *>((uint8_t *)&regs + rd.offset));
    }
  }
  return 0;
}

void Debugger::get_register(const std::string &reg)
{
  if (state != RUNNING) {
    std::cerr << "** program not running.\n";
    return;
  }
  struct user_regs_struct regs;
  ptrace(PTRACE_GETREGS, pid, nullptr, &regs);
  uint64_t value = get_register_value(reg, regs);
  std::cout << reg << " = " << std::dec << value << " (0x" << std::hex << value << ")\n";
}

void Debugger::set_register_value(const std::string &reg_name, uint64_t value, struct user_regs_struct &regs)
{
  for (const auto &rd : g_registers) {
    if (rd.name == reg_name) {
      *(reinterpret_cast<uint64_t *>((uint8_t *)&regs + rd.offset)) = value;
      return;
    }
  }
}

void print_reg(const std::string &name, unsigned long long val, bool force_fill = false)
{
  std::cout << std::left << std::setw(4) << name << " ";
  if (val == 0 && !force_fill)
    std::cout << std::left << std::setw(18) << "0";
  else {
    if (force_fill)
      std::cout << std::hex << std::setfill('0') << std::right << std::setw(16) << val << std::left << std::setfill(' ')
                << "  ";
    else
      std::cout << std::hex << std::left << std::setw(18) << val;
  }
}

void Debugger::get_registers()
{
  if (state != RUNNING) {
    std::cerr << "** program not running.\n";
    return;
  }
  struct user_regs_struct regs;
  ptrace(PTRACE_GETREGS, pid, nullptr, &regs);
  print_reg("RAX", regs.rax);
  print_reg("RBX", regs.rbx);
  print_reg("RCX", regs.rcx);
  print_reg("RDX", regs.rdx);
  std::cout << "\n";
  print_reg("R8", regs.r8);
  print_reg("R9", regs.r9);
  print_reg("R10", regs.r10);
  print_reg("R11", regs.r11);
  std::cout << "\n";
  print_reg("R12", regs.r12);
  print_reg("R13", regs.r13);
  print_reg("R14", regs.r14);
  print_reg("R15", regs.r15);
  std::cout << "\n";
  print_reg("RDI", regs.rdi);
  print_reg("RSI", regs.rsi);
  print_reg("RBP", regs.rbp);
  print_reg("RSP", regs.rsp);
  std::cout << "\n";
  print_reg("RIP", regs.rip);
  print_reg("FLAGS", regs.eflags, true);
  std::cout << "\n";
  std::cout << std::dec << std::setfill(' ');
}

void Debugger::set_register(const std::string &reg, uint64_t value)
{
  if (state != RUNNING) {
    std::cerr << "** program not running.\n";
    return;
  }
  struct user_regs_struct regs;
  ptrace(PTRACE_GETREGS, pid, nullptr, &regs);
  set_register_value(reg, value, regs);
  ptrace(PTRACE_SETREGS, pid, nullptr, &regs);
}

void Debugger::dump_memory(uintptr_t addr)
{
  if (state != RUNNING) {
    std::cerr << "** program not running.\n";
    return;
  }
  if (addr == 0) {
    std::cout << "** no addr is given.\n";
    return;
  }
  for (int line = 0; line < 5; line++) {
    std::cout << " " << std::hex << std::setw(7) << std::setfill('0') << (addr + line * 16) << ": ";
    uint8_t bytes[16];
    for (int i = 0; i < 16; i++) {
      bytes[i] = ptrace(PTRACE_PEEKTEXT, pid, (void *)(addr + line * 16 + i), nullptr) & 0xFF;
      std::cout << std::setw(2) << std::setfill('0') << std::hex << (int)bytes[i] << " ";
    }
    std::cout << "|";
    for (int i = 0; i < 16; i++)
      std::cout << (char)((bytes[i] >= 32 && bytes[i] <= 126) ? bytes[i] : '.');
    std::cout << "|\n";
  }
}

void Debugger::disassemble(uintptr_t addr)
{
  if (state != RUNNING) {
    std::cerr << "** program not running.\n";
    return;
  }
  if (addr == 0) {
    std::cout << "** no addr is given.\n";
    return;
  }
  cs_insn *insn;
  uint8_t code[150];
  for (int i = 0; i < 150; i++) {
    uintptr_t current_ptr = addr + i;
    code[i] = ptrace(PTRACE_PEEKTEXT, pid, (void *)(current_ptr), nullptr) & 0xFF;
    auto it = addr_to_bp.find(current_ptr);
    if (it != addr_to_bp.end())
      code[i] = it->second.original_data & 0xFF;
  }
  size_t count = cs_disasm(capstone_handle, code, 150, addr, 10, &insn);
  for (size_t i = 0; i < count; i++) {
    if (insn[i].address >= text_segment.end) {
      std::cout << "** the address is out of range\n";
      break;
    }
    std::cout << "      " << std::hex << std::setw(6) << std::setfill(' ') << std::right << insn[i].address << ": ";
    std::string bytes_str = "";
    char buf[4];
    for (int j = 0; j < insn[i].size; j++) {
      sprintf(buf, "%02x ", (int)insn[i].bytes[j]);
      bytes_str += buf;
    }
    std::cout << std::setw(32) << std::setfill(' ') << std::left << bytes_str << std::setw(10) << insn[i].mnemonic
              << insn[i].op_str << "\n";
  }
  cs_free(insn, count);
}

void Debugger::show_vmmap()
{
  if (state != RUNNING) {
    std::cerr << "** program not running.\n";
    return;
  }
  std::stringstream path;
  path << "/proc/" << pid << "/maps";
  std::ifstream maps(path.str());
  std::string line;
  while (getline(maps, line)) {
    std::stringstream iss(line);
    std::string addr_range, perms, offset;
    iss >> addr_range >> perms >> offset;
    size_t dash = addr_range.find('-');
    std::cout << std::setfill('0') << std::setw(16) << std::right << addr_range.substr(0, dash) << "-" << std::setw(16)
              << addr_range.substr(dash + 1) << " " << perms.substr(0, 3) << " " << std::hex
              << stoull(offset, nullptr, 16) << "\n";
  }
}

void Debugger::print_help()
{
  std::cout << "Available commands:\nbreak {addr}, cont, delete {id}, disasm, dump, exit, get, getregs, help, list, "
               "load, run, vmmap, set, si, start\n";
}

void Debugger::exit_debugger()
{
  if (pid > 0) {
    kill(pid, SIGKILL);
    waitpid(pid, nullptr, 0);
  }
  exit(0);
}

void Debugger::handle_command(const std::string &line)
{
  if (line.empty())
    return;
  std::stringstream iss(line);
  std::string cmd;
  iss >> cmd;
  if (cmd == "load") {
    std::string path;
    iss >> path;
    load_program(path);
  } else if (cmd == "start")
    start_program();
  else if (cmd == "run" || cmd == "r")
    run_program();
  else if (cmd == "cont" || cmd == "c")
    continue_execution();
  else if (cmd == "break" || cmd == "b") {
    std::string a;
    iss >> a;
    set_breakpoint(stoull(a, 0, 16));
  } else if (cmd == "delete") {
    int id;
    iss >> id;
    delete_breakpoint(id);
  } else if (cmd == "list" || cmd == "l")
    list_breakpoints();
  else if (cmd == "si")
    single_step();
  else if (cmd == "get" || cmd == "g") {
    std::string r;
    iss >> r;
    get_register(r);
  } else if (cmd == "getregs")
    get_registers();
  else if (cmd == "set" || cmd == "s") {
    std::string r;
    uint64_t v;
    iss >> r >> std::hex >> v;
    set_register(r, v);
  } else if (cmd == "dump" || cmd == "x") {
    std::string a;
    if (iss >> a)
      dump_memory(stoull(a, 0, 16));
    else
      dump_memory(0);
  } else if (cmd == "disasm" || cmd == "d") {
    std::string a;
    if (iss >> a)
      disassemble(stoull(a, 0, 16));
    else
      disassemble(0);
  } else if (cmd == "vmmap" || cmd == "m")
    show_vmmap();
  else if (cmd == "help" || cmd == "h")
    print_help();
  else if (cmd == "exit" || cmd == "q")
    exit_debugger();
}

void Debugger::run(bool use_script, const std::string &script_file)
{
  if (use_script) {
    std::ifstream script(script_file);
    std::string line;
    while (getline(script, line))
      handle_command(line);
  } else {
    std::string line;
    while (std::cout << "sdb> " && getline(std::cin, line))
      handle_command(line);
  }
}