diff --git a/include/environment.hpp b/include/environment.hpp new file mode 100644 index 0000000..0639138 --- /dev/null +++ b/include/environment.hpp @@ -0,0 +1,88 @@ +#pragma once +#include +#include +#include +#include + +#include "wirekit.hpp" + +struct TraceContext { + enum class EXECUTION_MODE { STARTING, USER, KERNEL /* Currently executing a system call */ }; + + pid_t pid; + EXECUTION_MODE mode; + user_regs_struct regs; + bool regs_dirty; +}; + +class TraceContextCollection { + std::unordered_map> m_contexts; + pid_t m_active; +public: + TraceContextCollection(const TraceContextCollection& other) = delete; + TraceContextCollection(TraceContextCollection&& other); + TraceContextCollection& operator=(const TraceContextCollection& other) = delete; + TraceContextCollection& operator=(TraceContextCollection&& other); + + TraceContextCollection(); + ~TraceContextCollection(); + + bool add_context(pid_t pid); + void delete_context(pid_t pid); + void clear(); + + TraceContext* context(pid_t pid); + const TraceContext* context(pid_t pid) const; + + void context_set_active(pid_t pid); + void context_clear_active(); + TraceContext* context_get_active(); + const TraceContext* context_get_active() const; +}; + +class Environment { + struct HookPair { + hook_t entry = nullptr; + hook_t exit = nullptr; + }; + + std::unordered_map m_hooks; + HookPair m_default_hooks; + TraceContextCollection m_trace_contexts; +public: + Environment(const Environment& other) = delete; + Environment(Environment&& other) = delete; + Environment& operator=(const Environment& other) = delete; + Environment& operator=(Environment&& other) = delete; + + Environment(); + ~Environment(); + + void register_hooks(reg_t syscall, hook_t entry, hook_t exit); + void register_default_hooks(hook_t entry, hook_t exit); + + bool entry_hook(reg_t syscall) const; + bool exit_hook(reg_t syscall) const; + bool default_entry_hook() const; + bool default_exit_hook() const; + + TraceContextCollection& contexts(); + const TraceContextCollection& contexts() const; + + void clear(); +}; + +// An EnvironmentScope calls Environment::clear() when destructed +class EnvironmentScope { + Environment& m_environment; +public: + EnvironmentScope(const EnvironmentScope& other) = delete; + EnvironmentScope(EnvironmentScope&& other) = delete; + EnvironmentScope& operator=(const EnvironmentScope& other) = delete; + EnvironmentScope& operator=(EnvironmentScope&& other) = delete; + + EnvironmentScope(Environment& environment); + ~EnvironmentScope(); +}; + +extern Environment ENV; diff --git a/include/interface.hpp b/include/interface.hpp new file mode 100644 index 0000000..06d3577 --- /dev/null +++ b/include/interface.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include "wirekit.hpp" + +extern err_t(*wirekit_prepare_handle)(); +extern err_t(*wirekit_command_start_handle)(int, const char* const*); +extern void(*wirekit_command_exit_handle)(); +extern void(*wirekit_subject_start_handle)(); +extern void(*wirekit_subject_exit_handle)(); + +// Exposure to wirekit is achieved by adding these functions to rewire's symbol tables +#define EXPOSED __attribute__((visibility("default"))) + +EXPOSED err_t rewire_syscall_hook(reg_t syscall, hook_t entry, hook_t exit); +EXPOSED err_t rewire_syscall_unhook(reg_t syscall); +EXPOSED err_t rewire_syscall_hook_default(hook_t entry, hook_t exit); +EXPOSED err_t rewire_syscall_unhook_default(); +EXPOSED err_t rewire_subject_id(pid_t* pid); +EXPOSED err_t rewire_subject_get_regs(user_regs_struct* regs); +EXPOSED err_t rewire_subject_set_regs(const user_regs_struct* regs); +EXPOSED err_t rewire_subject_load_cstr(const char* subject_addr, char* buf, uint32_t* read, uint32_t n); diff --git a/include/wirekit.hpp b/include/wirekit.hpp new file mode 100644 index 0000000..17a18a7 --- /dev/null +++ b/include/wirekit.hpp @@ -0,0 +1,33 @@ +#pragma once +#include +#include +#include + +using reg_t = decltype(user_regs_struct::orig_rax); +using hook_t = void(*)(); +using err_t = uint8_t; + +constexpr err_t REWIRE_SUCCESS = 0; +constexpr err_t REWIRE_FAILURE = 1; +constexpr err_t REWIRE_FAILURE_NOCTX = 2; + +extern "C" { + +// Exposed rewire functions - to be used by wirekit +err_t rewire_syscall_hook(reg_t syscall, hook_t entry, hook_t exit); +err_t rewire_syscall_unhook(reg_t syscall); // Alias for rewire_syscall_hook(syscall, nullptr, nullptr) +err_t rewire_syscall_hook_default(hook_t entry, hook_t exit); +err_t rewire_syscall_unhook_default(); // Alias for rewire_syscall_hook_default(nullptr, nullptr); +err_t rewire_subject_id(pid_t* pid); +err_t rewire_subject_get_regs(user_regs_struct* regs); +err_t rewire_subject_set_regs(const user_regs_struct* regs); +err_t rewire_subject_load_cstr(const char* subject_addr, char* buf, uint32_t* read, uint32_t n); + +// Wirekit control functions - called by rewire for setup +err_t wirekit_prepare(); // NOCTX - Called when wirekit is loaded +err_t wirekit_command_start(int argc, const char* const* argv); // NOCTX - Called before command is executed +void wirekit_command_exit(); // Optional, NOCTX - Called after command finished executing +void wirekit_subject_start(); // Optional - Called when a new thread starts working on the current command +void wirekit_subject_exit(); // Optional - Called when a thread stops working on the current command + +} diff --git a/src/environment.cpp b/src/environment.cpp new file mode 100644 index 0000000..f9c646e --- /dev/null +++ b/src/environment.cpp @@ -0,0 +1,182 @@ +#include "environment.hpp" + +#include + +Environment ENV; + +// Trace context collection + +TraceContextCollection::TraceContextCollection(TraceContextCollection&& other) { + m_contexts = std::move(other.m_contexts); + m_active = other.m_active; + other.context_clear_active(); +} + +TraceContextCollection& TraceContextCollection::operator=(TraceContextCollection&& other) { + if (this == &other) + return *this; + + m_contexts = std::move(other.m_contexts); + m_active = other.m_active; + other.context_clear_active(); + return *this; +} + +TraceContextCollection::TraceContextCollection() {} + +TraceContextCollection::~TraceContextCollection() {} + +bool TraceContextCollection::add_context(pid_t pid) { + if (m_contexts.contains(pid)) + return false; + + std::unique_ptr new_context = std::make_unique(); + if (!new_context) + return false; + + new_context->pid = pid; + new_context->mode = TraceContext::EXECUTION_MODE::STARTING; + std::memset(&new_context->regs, 0, sizeof(new_context->regs)); + new_context->regs_dirty = false; + + m_contexts.emplace(pid, std::move(new_context)); + return true; +} + +void TraceContextCollection::delete_context(pid_t pid) { + if (m_contexts.contains(pid)) + m_contexts.erase(pid); + if (m_active == pid) + context_clear_active(); +} + +void TraceContextCollection::clear() { + m_contexts.clear(); + context_clear_active(); +} + +TraceContext* TraceContextCollection::context(pid_t pid) { + if (!m_contexts.contains(pid)) + return nullptr; + + return m_contexts.at(pid).get(); +} + +const TraceContext* TraceContextCollection::context(pid_t pid) const { + if (!m_contexts.contains(pid)) + return nullptr; + + return m_contexts.at(pid).get(); +} + +void TraceContextCollection::context_set_active(pid_t pid) { + if (m_contexts.contains(pid)) + m_active = pid; + else + context_clear_active(); +} + +void TraceContextCollection::context_clear_active() { + m_active = -1; +} + +TraceContext* TraceContextCollection::context_get_active() { + return context(m_active); +} + +const TraceContext* TraceContextCollection::context_get_active() const { + return context(m_active); +} + +// Environment + +Environment::Environment() { + clear(); +} + +Environment::~Environment() {} + +void Environment::register_hooks(reg_t syscall, hook_t entry, hook_t exit) { + if (!entry && !exit) { + if (m_hooks.contains(syscall)) + m_hooks.erase(syscall); + + return; + } + + if (!m_hooks.contains(syscall)) + m_hooks.emplace(syscall, HookPair{}); + + HookPair& hooks = m_hooks.at(syscall); + hooks.entry = entry; + hooks.exit = exit; +} + +void Environment::register_default_hooks(hook_t entry, hook_t exit) { + m_default_hooks.entry = entry; + m_default_hooks.exit = exit; +} + +bool Environment::entry_hook(reg_t syscall) const { + if (!m_hooks.contains(syscall)) + return false; + + const HookPair& hooks = m_hooks.at(syscall); + if (!hooks.entry) + return false; + + hooks.entry(); + return true; +} + +bool Environment::exit_hook(reg_t syscall) const { + if (!m_hooks.contains(syscall)) + return false; + + const HookPair& hooks = m_hooks.at(syscall); + if (!hooks.exit) + return false; + + hooks.exit(); + return true; +} + +bool Environment::default_entry_hook() const { + if (!m_default_hooks.entry) + return false; + + m_default_hooks.entry(); + return true; +} + +bool Environment::default_exit_hook() const { + if (!m_default_hooks.exit) + return false; + + m_default_hooks.exit(); + return true; +} + +TraceContextCollection& Environment::contexts() { + return m_trace_contexts; +} + +const TraceContextCollection& Environment::contexts() const { + return m_trace_contexts; +} + +void Environment::clear() { + m_hooks.clear(); + m_default_hooks.entry = nullptr; + m_default_hooks.exit = nullptr; + m_trace_contexts.clear(); +} + +// EnvironmentScope to clear Environment on scope exit + +EnvironmentScope::EnvironmentScope(Environment& environment) : m_environment{ environment } {} + +EnvironmentScope::~EnvironmentScope() { + m_environment.clear(); +} + diff --git a/src/interface.cpp b/src/interface.cpp new file mode 100644 index 0000000..85386bd --- /dev/null +++ b/src/interface.cpp @@ -0,0 +1,88 @@ +#include "interface.hpp" +#include "environment.hpp" + +#include +#include +#include + +err_t(*wirekit_prepare_handle)() = nullptr; +err_t(*wirekit_command_start_handle)(int, const char* const*) = nullptr; +void(*wirekit_command_exit_handle)() = nullptr; +void(*wirekit_subject_start_handle)() = nullptr; +void(*wirekit_subject_exit_handle)() = nullptr; + +err_t rewire_syscall_hook(reg_t syscall, hook_t entry, hook_t exit) { + ENV.register_hooks(syscall, entry, exit); + return REWIRE_SUCCESS; +} + +err_t rewire_syscall_unhook(reg_t syscall) { + rewire_syscall_hook(syscall, nullptr, nullptr); + return REWIRE_SUCCESS; +} + +err_t rewire_syscall_hook_default(hook_t entry, hook_t exit) { + ENV.register_default_hooks(entry, exit); + return REWIRE_SUCCESS; +} + +err_t rewire_syscall_unhook_default() { + rewire_syscall_hook_default(nullptr, nullptr); + return REWIRE_SUCCESS; +} + +err_t rewire_subject_id(pid_t* pid) { + const TraceContext* ctx = ENV.contexts().context_get_active(); + if (!ctx) + return REWIRE_FAILURE_NOCTX; + + *pid = ctx->pid; + return REWIRE_SUCCESS; +} + +err_t rewire_subject_get_regs(user_regs_struct* regs) { + const TraceContext* ctx = ENV.contexts().context_get_active(); + if (!ctx) + return REWIRE_FAILURE_NOCTX; + + std::memcpy(regs, &ctx->regs, sizeof(ctx->regs)); + return REWIRE_SUCCESS; +} + +err_t rewire_subject_set_regs(const user_regs_struct* regs) { + TraceContext* ctx = ENV.contexts().context_get_active(); + if (!ctx) + return REWIRE_FAILURE_NOCTX; + + std::memcpy(&ctx->regs, regs, sizeof(ctx->regs)); + ctx->regs_dirty = true; + return REWIRE_SUCCESS; +} + +err_t rewire_subject_load_cstr(const char* subject_addr, char* buf, uint32_t* read, uint32_t n) { + const TraceContext* ctx = ENV.contexts().context_get_active(); + if (!ctx) + return REWIRE_FAILURE_NOCTX; + + *read = 0; + while (*read < n) { + errno = 0; + long word = ptrace(PTRACE_PEEKDATA, ctx->pid, (char*)subject_addr, NULL); // TODO + if (errno == -1l && errno != 0) // errno differentiates successful -1 word from error + return REWIRE_FAILURE; + + for (unsigned i = 0; i < sizeof(long) && *read < n; i++) { + char c = *(reinterpret_cast(&word) + i); // Unlike bit operations, this is independent of endianness + *(buf++) = c; + (*read)++; + + if (c == 0) + return REWIRE_SUCCESS; // Encountered terminating null byte + } + + subject_addr += sizeof(long); + } + + return REWIRE_SUCCESS; +} + diff --git a/src/main.cpp b/src/main.cpp index 473b7e2..ece3f75 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -1,9 +1,20 @@ #include #include +#include +#include +#include +#include +#include +#include +#include #include "cli.hpp" #include "dl.hpp" #include "install.hpp" +#include "interface.hpp" +#include "environment.hpp" + +#define FORCEQUIT { std::cout << "Forcefully quitting!" << std::endl; while (true) exit(1); } int launch_help(int argc, char** argv) { std::cout << "Use rewire via one of the following commands:" << std::endl; @@ -95,6 +106,200 @@ int launch_list(int argc, char** argv) { return 0; } +pid_t dispatch_tracee(int argc, const char* const* argv) { + pid_t child = fork(); + + if (child != 0) + return child; // Either -1 on error or child pid + + ptrace(PTRACE_TRACEME, 0, 0, 0); + execvp(argv[0], (char* const*)argv); // Removing const here is in child process + // Note -> Tracee stops with SIGTRAP immediately after execvp! + + // Child should never return from execvp + while (true) exit(1); +} + +class KillWrapper { + pid_t m_pid; +public: + KillWrapper(const KillWrapper& other) = delete; + KillWrapper(KillWrapper&& other) = delete; + KillWrapper& operator=(const KillWrapper& other) = delete; + KillWrapper& operator=(KillWrapper&& other) = delete; + + KillWrapper(pid_t pid) : m_pid{ pid } {} + + ~KillWrapper() { + if (m_pid != -1) + kill(m_pid, SIGKILL); + } + + void release() { + m_pid = -1; + } +}; + +void run_command(int argc, const char* const* argv) { + EnvironmentScope es { ENV }; // Make sure environment is cleared after each command + + // Notify wirekit about command start + if (wirekit_command_start_handle(argc, argv) != REWIRE_SUCCESS) { + std::cout << "Wirekit prohibited command execution!" << std::endl; + return; + } + + // Dispatch main command process + pid_t main_pid = dispatch_tracee(argc, argv); + if (main_pid == -1) { + std::cout << "Failed to start command main process!" << std::endl; + return; + } + + KillWrapper kw { main_pid }; // Make sure it gets whiped until in main trace loop + + int exec_status; + if (waitpid(main_pid, &exec_status, 0) != main_pid || !WIFSTOPPED(exec_status) || WSTOPSIG(exec_status) != SIGTRAP) { + std::cout << "Execvp failed to run the command!" << std::endl; + return; + } + + /* + PTRACE_O_EXITKILL -> Kill tracees when tracer exits + PTRACE_O_TRACESYSGOOD -> SIGTRAP signals caused by system call entry/exit stops have 0x80 bit set + PTRACE_O_TRACEFORK, PTRACE_O_TRACEVFORK, PTRACE_O_TRACECLONE -> Automatically trace new children (inherit options and mode) + */ + if (ptrace(PTRACE_SETOPTIONS, main_pid, PTRACE_O_EXITKILL, PTRACE_O_TRACESYSGOOD, + PTRACE_O_TRACEFORK, PTRACE_O_TRACEVFORK, PTRACE_O_TRACECLONE) != 0) { + std::cout << "Tracer failed to set trace options of command main process!" << std::endl; + return; + } + + // Let first thread actually start running (until syscall entry) + if (ptrace(PTRACE_SYSCALL, main_pid, 0, 0) != 0) { + std::cout << "Failed to start command main process execution!" << std::endl; + return; + } + + // Inform wirekit about first thread and initialize its trace context + ENV.contexts().add_context(main_pid); + ENV.contexts().context(main_pid)->mode = TraceContext::EXECUTION_MODE::USER; + + ENV.contexts().context_set_active(main_pid); + + if (wirekit_subject_start_handle) + wirekit_subject_start_handle(); + + ENV.contexts().context_clear_active(); + + // Start tracer loop + kw.release(); + + int status; + pid_t pid; + + while ((pid = waitpid(-1, &status, __WALL)) != -1 || errno != ECHILD /* No (grand)children left */) { + if (pid == -1) { + std::cout << "Tracer interrupted by signal (errno=" << errno << ")!" << std::endl; + FORCEQUIT; + } + + /* + New processes and threads inherit PTRACE_SYSCALL execution mode from parent and thus end up + here when they perform their first syscall. Start logging context at this point. + */ + if (!ENV.contexts().context(pid)) { + ENV.contexts().add_context(pid); + ENV.contexts().context(pid)->mode = TraceContext::EXECUTION_MODE::USER; + + ENV.contexts().context_set_active(pid); + + if (wirekit_subject_start_handle) + wirekit_subject_start_handle(); + + ENV.contexts().context_clear_active(); + } + + // Activate context (definitely exists now) + ENV.contexts().context_set_active(pid); + + if (WIFEXITED(status)) { + // A child exited -> notify wirekit and stop tracing it + if (wirekit_subject_exit_handle) + wirekit_subject_exit_handle(); + + ENV.contexts().delete_context(pid); + } else if (WIFSTOPPED(status)) { + int sig = WSTOPSIG(status); + + if (sig == SIGTRAP) { + int event = (unsigned)status >> 16; + + if (event == PTRACE_EVENT_FORK || event == PTRACE_EVENT_VFORK || event == PTRACE_EVENT_CLONE) { + ; // New thread but nothing to do -> tracer attaches on first syscall + } else { + std::cout << "Received non-syscall SIGTRAP without thread creation event!" << std::endl; + FORCEQUIT; + } + } else if (sig == (SIGTRAP | 0x80) /* 0x80 thanks to PTRACE_O_TRACESYSGOOD */) { + // Tracee was trapped on either syscall entry or exit + TraceContext* ctx = ENV.contexts().context_get_active(); + + if (ptrace(PTRACE_GETREGS, pid, 0, &ctx->regs) != 0) { + std::cout << "Failed to read tracee registers!" << std::endl; + FORCEQUIT; + } + + ctx->regs_dirty = false; + + reg_t syscall = ctx->regs.orig_rax; + + if (ctx->mode == TraceContext::EXECUTION_MODE::USER) + ENV.entry_hook(syscall) || ENV.default_entry_hook(); // Short-circuit guaranteed by C++ standard ;) + else if (ctx->mode == TraceContext::EXECUTION_MODE::KERNEL) + ENV.exit_hook(syscall) || ENV.default_exit_hook(); + + if (ctx->regs_dirty) { + if (ptrace(PTRACE_SETREGS, pid, NULL, &ctx->regs) != 0) { + std::cout << "Failed to write tracee registers!" << std::endl; + FORCEQUIT; + } + } + + ctx->regs_dirty = false; + + auto progress_mode = [](TraceContext::EXECUTION_MODE m) { + switch (m) { + case TraceContext::EXECUTION_MODE::STARTING: + return TraceContext::EXECUTION_MODE::USER; + case TraceContext::EXECUTION_MODE::USER: + return TraceContext::EXECUTION_MODE::KERNEL; + case TraceContext::EXECUTION_MODE::KERNEL: + return TraceContext::EXECUTION_MODE::USER; + } + + return TraceContext::EXECUTION_MODE::STARTING; + }; + + ctx->mode = progress_mode(ctx->mode); + } + + // Let tracee continue until next system call event (entry or exit) + if (ptrace(PTRACE_SYSCALL, pid, 0, 0) != 0) { + std::cout << "Failed to continue command execution after SIGTRAP!" << std::endl; + FORCEQUIT; + } + } // else { Weird signal is happening, should be irrelevant to wirekit though! } + + // Deactivate current context + ENV.contexts().context_clear_active(); + } + + // Notify wirekit that command has finished executing (no threads are left) + if (wirekit_command_exit_handle) + wirekit_command_exit_handle(); +} + int launch_run(int argc, char** argv) { if (argc < 2) { std::cout << "Invalid usage, try \"rewire help\" for more information!" << std::endl; @@ -119,16 +324,64 @@ int launch_run(int argc, char** argv) { return 1; } - /* - TODO - environment.hpp/cpp - class Environment -> kit handles, hooks, trace contexts - class CommandScope -> hooks and trace contexts can only be added while scope holds environment, all cleared on scope exit - class TraceContext -> execution data on individual tracee threads - class TraceContextCollection -> management of all trace contexts and selection of active context - */ + // Collect wirekit control function handles + if (!(wirekit_prepare_handle = (err_t(*)())wirekit.resolve("wirekit_prepare"))) { + std::cout << "Wirekit is missing \"wirekit_prepare\" implementation!" << std::endl; + return 1; + } - // TODO -> Execute single command (argc > 2) or shell + if (!(wirekit_command_start_handle = (err_t(*)(int, const char* const*))wirekit.resolve("wirekit_command_start"))) { + std::cout << "Wirekit is missing \"wirekit_command_start\" implementation!" << std::endl; + return 1; + } + + wirekit_command_exit_handle = (void(*)())wirekit.resolve("wirekit_command_exit"); + wirekit_subject_start_handle = (void(*)())wirekit.resolve("wirekit_subject_start"); + wirekit_subject_exit_handle = (void(*)())wirekit.resolve("wirekit_subject_exit"); + + // Prepare wirekit + err_t prepare_err; + if ((prepare_err = wirekit_prepare_handle()) != REWIRE_SUCCESS) { + std::cout << "Wirekit preparation failed (wirekit_prepare returned " << prepare_err << ")!" << std::endl; + return 1; + } + + // Execute commands in rewired syscall environment + if (argc > 2) { + run_command(argc - 2, (const char**)(argv + 2)); // Execute a single command + } else { + // Shell mode + while (true) { + std::cout << "> "; + + std::string command; + std::getline(std::cin, command); + + // Split command at whitespaces + std::stringstream command_stream { command }; + std::vector command_split; // No need to reserve, presumably not many elements + + std::string part; + while (command_stream >> std::quoted(part)) + command_split.emplace_back(std::move(part)); + + if (command_split.empty()) + continue; + + if (command_split[0] == "exit") + break; + + // Compile argv (cstr pointer array) + std::vector command_split_cstr; + command_split_cstr.resize(command_split.size() + 1, nullptr); // +1 for nullptr termination + + for (std::size_t i = 0; i < command_split.size(); i++) + command_split_cstr[i] = command_split[i].c_str(); + + // Execute command + run_command((int)command_split_cstr.size() - 1, command_split_cstr.data()); + } + } return 0; }