From 7d2f643692eef4f7c4b75c237db63717cc3800ba Mon Sep 17 00:00:00 2001 From: ifratric Date: Thu, 31 Oct 2024 11:02:52 +0100 Subject: [PATCH] Implement -instrument_transitive --- Linux/debugger.cpp | 17 +++++----- Linux/debugger.h | 3 +- README.md | 2 ++ Windows/debugger.cpp | 27 +++++++++------- Windows/debugger.h | 3 +- macOS/debugger.cpp | 75 ++++++++++++++++++++++++-------------------- macOS/debugger.h | 8 +++-- tinyinst.cpp | 34 ++++++++++++++++---- tinyinst.h | 3 ++ 9 files changed, 109 insertions(+), 63 deletions(-) diff --git a/Linux/debugger.cpp b/Linux/debugger.cpp index ec38f63..201dc09 100644 --- a/Linux/debugger.cpp +++ b/Linux/debugger.cpp @@ -1499,7 +1499,8 @@ void Debugger::ExtractCodeRanges(void *module_base, size_t min_address, size_t max_address, std::list *executable_ranges, - size_t *code_size) + size_t *code_size, + bool do_protect) { std::string elf_filename; if(module_base) { @@ -1531,12 +1532,14 @@ void Debugger::ExtractCodeRanges(void *module_base, iter->addr_from = min_address; } } - - int ret = RemoteMprotect((void *)iter->addr_from, - (iter->addr_to - iter->addr_from), - iter->permissions ^ PROT_EXEC); - if(ret) { - FATAL("Could not apply memory protection"); + + if(do_protect) { + int ret = RemoteMprotect((void *)iter->addr_from, + (iter->addr_to - iter->addr_from), + iter->permissions ^ PROT_EXEC); + if(ret) { + FATAL("Could not apply memory protection"); + } } AddressRange range; diff --git a/Linux/debugger.h b/Linux/debugger.h index d00957e..6216c81 100644 --- a/Linux/debugger.h +++ b/Linux/debugger.h @@ -162,7 +162,8 @@ class Debugger { size_t min_address, size_t max_address, std::list *executable_ranges, - size_t *code_size); + size_t *code_size, + bool do_protect = true); void ProtectCodeRanges(std::list *executable_ranges); diff --git a/README.md b/README.md index ece5e89..478df8f 100755 --- a/README.md +++ b/README.md @@ -184,6 +184,8 @@ In addition to the general-purpose API documented above, TinyInst also implement `-instrument_module [module name]` specifies which module to instrument, multiple `-instrument_module` options can be specified to instrument multiple modules. +`-instrument_transitive [module name]` similar to `-instrument_module` except only code entered from other instrumented modules will run instrumented. Primarily used as optimization for calls like module1->module2->module1 where it's not important to instrument the entire module2 module, but module2->module1 entries are causing slowdowns. + `-indirect_instrumentation [none|local|global|auto]` which instrumentation to use for indirect jump/calls `-patch_return_addresses` - replaces return address with the original value, causes returns to be instrumented using whatever `-indirect_instrumentation` method is specified diff --git a/Windows/debugger.cpp b/Windows/debugger.cpp index 56bbc14..801644a 100644 --- a/Windows/debugger.cpp +++ b/Windows/debugger.cpp @@ -553,7 +553,8 @@ void Debugger::ExtractCodeRanges(void *module_base, size_t min_address, size_t max_address, std::list *executable_ranges, - size_t *code_size) + size_t *code_size, + bool do_protect) { LPCVOID end_address = (void *)max_address; LPCVOID cur_address = (void *)min_address; @@ -593,17 +594,19 @@ void Debugger::ExtractCodeRanges(void *module_base, FATAL("Error in ReadProcessMemory"); } - uint8_t low = meminfobuf.Protect & 0xFF; - low = low >> 4; - DWORD newProtect = (meminfobuf.Protect & 0xFFFFFF00) + low; - DWORD oldProtect; - if (!VirtualProtectEx(child_handle, - meminfobuf.BaseAddress, - meminfobuf.RegionSize, - newProtect, - &oldProtect)) - { - FATAL("Error in VirtualProtectEx"); + if(do_protect) { + uint8_t low = meminfobuf.Protect & 0xFF; + low = low >> 4; + DWORD newProtect = (meminfobuf.Protect & 0xFFFFFF00) + low; + DWORD oldProtect; + if (!VirtualProtectEx(child_handle, + meminfobuf.BaseAddress, + meminfobuf.RegionSize, + newProtect, + &oldProtect)) + { + FATAL("Error in VirtualProtectEx"); + } } newRange.from = (size_t)meminfobuf.BaseAddress; diff --git a/Windows/debugger.h b/Windows/debugger.h index bccab5d..471fecf 100644 --- a/Windows/debugger.h +++ b/Windows/debugger.h @@ -124,7 +124,8 @@ class Debugger { size_t min_address, size_t max_address, std::list *executable_ranges, - size_t *code_size); + size_t *code_size, + bool do_protect = true); void ProtectCodeRanges(std::list *executable_ranges); diff --git a/macOS/debugger.cpp b/macOS/debugger.cpp index c9b9f95..2f4248d 100644 --- a/macOS/debugger.cpp +++ b/macOS/debugger.cpp @@ -815,8 +815,9 @@ void Debugger::ExtractCodeRanges(void *base_address, size_t min_address, size_t max_address, std::list *executable_ranges, - size_t *code_size) { - + size_t *code_size, + bool do_protect) +{ if(!base_address) { ExtractSegmentCodeRanges(min_address, max_address, executable_ranges, code_size); return; @@ -876,7 +877,7 @@ void Debugger::ExtractCodeRanges(void *base_address, mach_vm_address_t segment_start_addr = (mach_vm_address_t)segment_cmd->vmaddr + file_vm_slide; mach_vm_address_t segment_end_addr = (mach_vm_address_t)segment_cmd->vmaddr + file_vm_slide + segment_cmd->vmsize; - ExtractSegmentCodeRanges(segment_start_addr, segment_end_addr, executable_ranges, code_size); + ExtractSegmentCodeRanges(segment_start_addr, segment_end_addr, executable_ranges, code_size, do_protect); #endif } @@ -889,7 +890,9 @@ void Debugger::ExtractCodeRanges(void *base_address, void Debugger::ExtractSegmentCodeRanges(mach_vm_address_t segment_start_addr, mach_vm_address_t segment_end_addr, std::list *executable_ranges, - size_t *code_size) { + size_t *code_size, + bool do_protect) +{ mach_vm_address_t cur_address = segment_start_addr; while (cur_address < segment_end_addr) { mach_vm_size_t region_size = 0; @@ -916,38 +919,42 @@ void Debugger::ExtractSegmentCodeRanges(mach_vm_address_t segment_start_addr, new_range.data = (char *)malloc(range_size); RemoteRead((void*)new_range.from, new_range.data, range_size); - retry_label: - RemoteProtect((void*)new_range.from, range_size, info.protection ^ VM_PROT_EXECUTE); - mach_vm_address_t region_addr = new_range.from; - mach_vm_size_t region_sz = range_size; - vm_region_submap_info_data_64_t region_info; - mach_target->GetRegionSubmapInfo(®ion_addr, (mach_vm_size_t*)®ion_sz, ®ion_info); - if (region_info.protection & VM_PROT_EXECUTE) { - if (retried) { - FATAL("Failed to mark the original code NON-EXECUTABLE\n"); - } - - kern_return_t krt; - krt = mach_vm_deallocate(mach_target->Task(), - (mach_vm_address_t)new_range.from, - range_size); - - if (krt == KERN_SUCCESS) { - mach_vm_address_t alloc_address = new_range.from; - krt = mach_vm_allocate(mach_target->Task(), - (mach_vm_address_t*)&alloc_address, - range_size, - VM_FLAGS_FIXED); - - if (krt == KERN_SUCCESS && alloc_address && new_range.from) { - RemoteWrite((void*)new_range.from, new_range.data, range_size); - } else { - FATAL("Unable to re-allocate memory after deallocate in ExtractSegmentCodeRanges\n"); + if(do_protect) { + retry_label: + RemoteProtect((void*)new_range.from, range_size, info.protection ^ VM_PROT_EXECUTE); + mach_vm_address_t region_addr = new_range.from; + mach_vm_size_t region_sz = range_size; + vm_region_submap_info_data_64_t region_info; + mach_target->GetRegionSubmapInfo(®ion_addr, (mach_vm_size_t*)®ion_sz, ®ion_info); + if (region_info.protection & VM_PROT_EXECUTE) { + if (retried) { + FATAL("Failed to mark the original code NON-EXECUTABLE\n"); + } + + kern_return_t krt; + krt = mach_vm_deallocate(mach_target->Task(), + (mach_vm_address_t)new_range.from, + range_size); + + if (krt == KERN_SUCCESS) { + mach_vm_address_t alloc_address = new_range.from; + krt = mach_vm_allocate(mach_target->Task(), + (mach_vm_address_t*)&alloc_address, + range_size, + VM_FLAGS_FIXED); + + if (krt == KERN_SUCCESS && alloc_address && new_range.from) { + RemoteWrite((void*)new_range.from, new_range.data, range_size); + } else { + FATAL("Unable to re-allocate memory after deallocate in ExtractSegmentCodeRanges\n"); + } } + + retried = true; + goto retry_label; } - - retried = true; - goto retry_label; + } else { + //WARN("skipping memory protection"); } AddressRange *last_range = NULL; diff --git a/macOS/debugger.h b/macOS/debugger.h index bd39126..36b0133 100644 --- a/macOS/debugger.h +++ b/macOS/debugger.h @@ -156,6 +156,9 @@ friend kern_return_t catch_mach_exception_raise_state_identity( DebuggerStatus Kill(); DebuggerStatus Continue(uint32_t timeout); DebuggerStatus Attach(unsigned int pid, uint32_t timeout); + DebuggerStatus GetDebuggerStatus() { + return dbg_last_status; + } bool IsTargetAlive(); bool IsTargetFunctionDefined() { return target_function_defined; } @@ -254,7 +257,8 @@ friend kern_return_t catch_mach_exception_raise_state_identity( size_t min_address, size_t max_address, std::list *executable_ranges, - size_t *code_size); + size_t *code_size, + bool do_protect = true); void ProtectCodeRanges(std::list *executable_ranges); @@ -398,7 +402,7 @@ friend kern_return_t catch_mach_exception_raise_state_identity( void ExtractSegmentCodeRanges(mach_vm_address_t segment_start_addr, mach_vm_address_t segment_end_addr, std::list *executable_ranges, - size_t *code_size); + size_t *code_size, bool do_protect = true); void HandleDyld(void *module); diff --git a/tinyinst.cpp b/tinyinst.cpp index 99b59d7..2474e18 100644 --- a/tinyinst.cpp +++ b/tinyinst.cpp @@ -49,6 +49,7 @@ ModuleInfo::ModuleInfo() { instrumented_code_size = 0; unwind_data = NULL; client_data = NULL; + do_protect = true; } void ModuleInfo::ClearInstrumentation() { @@ -900,16 +901,19 @@ void TinyInst::InstrumentModule(ModuleInfo *module) { printf("Module %s already instrumented, " "reusing instrumentation data\n", module->module_name.c_str()); - ProtectCodeRanges(&module->executable_ranges); + if (module->do_protect) { + ProtectCodeRanges(&module->executable_ranges); + } FixCrossModuleLinks(module); return; } - + ExtractCodeRanges(module->module_header, module->min_address, module->max_address, &module->executable_ranges, - &module->code_size); + &module->code_size, + module->do_protect); // allocate buffer for instrumented code module->instrumented_code_size = module->code_size * CODE_SIZE_MULTIPLIER; @@ -1213,6 +1217,18 @@ InstructionResult TinyInst::InstrumentInstruction(ModuleInfo *module, return INST_NOTHANDLED; } +void TinyInst::AddInstrumentedModule(char* name, bool do_protect) { + std::string module_name = name; + for(auto iter=instrumented_modules.begin(); iter!=instrumented_modules.end(); iter++) { + if((*iter)->module_name == module_name) { + FATAL("Duplicate instrumented modules, module %s is already being instrumented", name); + } + } + ModuleInfo *new_module = new ModuleInfo(); + new_module->module_name = module_name; + new_module->do_protect = do_protect; + instrumented_modules.push_back(new_module); +} // initializes instrumentation from command line options void TinyInst::Init(int argc, char **argv) { @@ -1298,9 +1314,15 @@ void TinyInst::Init(int argc, char **argv) { #endif for (const auto module_name: module_names) { - ModuleInfo *new_module = new ModuleInfo(); - new_module->module_name = module_name; - instrumented_modules.push_back(new_module); + AddInstrumentedModule(module_name, true); + // SAY("--- %s\n", module_name); + } + + std::list module_names_transitive; + GetOptionAll("-instrument_transitive", argc, argv, &module_names_transitive); + + for (const auto module_name: module_names_transitive) { + AddInstrumentedModule(module_name, false); // SAY("--- %s\n", module_name); } diff --git a/tinyinst.h b/tinyinst.h index ee1a8c4..22dfff6 100644 --- a/tinyinst.h +++ b/tinyinst.h @@ -165,6 +165,7 @@ class TinyInst : public Debugger { void InstrumentAddressRange(const char *name, size_t min_address, size_t max_address); private: + void AddInstrumentedModule(char* name, bool do_protect); bool HandleBreakpoint(void *address); void OnInstrumentModuleLoaded(void *module, ModuleInfo *target_module); ModuleInfo *IsInstrumentModule(char *module_name); @@ -321,6 +322,8 @@ class ModuleInfo { std::unordered_set entry_offsets; UnwindData *unwind_data; + + bool do_protect; // clients can use this to store additional data // about the module