Index: include/xray/xray_interface.h =================================================================== --- include/xray/xray_interface.h +++ include/xray/xray_interface.h @@ -67,6 +67,14 @@ // result values. extern XRayPatchingStatus __xray_unpatch(); +// This patches a specific function id. See XRayPatchingStatus for possible +// result values. +extern XRayPatchingStatus __xray_patch_function(int32_t FuncId); + +// This unpatches a specific function id. See XRayPatchingStatus for possible +// result values. +extern XRayPatchingStatus __xray_unpatch_function(int32_t FuncId); + // Use XRay to log the first argument of each (instrumented) function call. // When this function exits, all threads will have observed the effect and // start logging their subsequent affected function calls (if patched). Index: lib/xray/xray_init.cc =================================================================== --- lib/xray/xray_init.cc +++ lib/xray/xray_init.cc @@ -25,6 +25,8 @@ void __xray_init(); extern const XRaySledEntry __start_xray_instr_map[] __attribute__((weak)); extern const XRaySledEntry __stop_xray_instr_map[] __attribute__((weak)); +extern const XRayFunctionSledIndex __start_xray_fn_idx[] __attribute__((weak)); +extern const XRayFunctionSledIndex __stop_xray_fn_idx[] __attribute__((weak)); } using namespace __xray; @@ -55,6 +57,8 @@ __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex); XRayInstrMap.Sleds = __start_xray_instr_map; XRayInstrMap.Entries = __stop_xray_instr_map - __start_xray_instr_map; + XRayInstrMap.SledsIndex = __start_xray_fn_idx; + XRayInstrMap.Functions = __stop_xray_fn_idx - __start_xray_fn_idx; } __sanitizer::atomic_store(&XRayInitialized, true, __sanitizer::memory_order_release); Index: lib/xray/xray_interface.cc =================================================================== --- lib/xray/xray_interface.cc +++ lib/xray/xray_interface.cc @@ -132,12 +132,48 @@ return CleanupInvoker{Fn}; } +inline bool patchSled(const XRaySledEntry &Sled, bool Enable, + int32_t FuncId) XRAY_NEVER_INSTRUMENT { + // While we're here, we should patch the nop sled. To do that we mprotect + // the page containing the function to be writeable. + const uint64_t PageSize = GetPageSizeCached(); + void *PageAlignedAddr = + reinterpret_cast(Sled.Address & ~(PageSize - 1)); + std::size_t MProtectLen = (Sled.Address + cSledLength) - + reinterpret_cast(PageAlignedAddr); + MProtectHelper Protector(PageAlignedAddr, MProtectLen); + if (Protector.MakeWriteable() == -1) { + printf("Failed mprotect: %d\n", errno); + return XRayPatchingStatus::FAILED; + } + + bool Success = false; + switch (Sled.Kind) { + case XRayEntryType::ENTRY: + Success = patchFunctionEntry(Enable, FuncId, Sled, __xray_FunctionEntry); + break; + case XRayEntryType::EXIT: + Success = patchFunctionExit(Enable, FuncId, Sled); + break; + case XRayEntryType::TAIL: + Success = patchFunctionTailExit(Enable, FuncId, Sled); + break; + case XRayEntryType::LOG_ARGS_ENTRY: + Success = patchFunctionEntry(Enable, FuncId, Sled, __xray_ArgLoggerEntry); + break; + default: + Report("Unsupported sled kind: %d\n", int(Sled.Kind)); + return false; + } + return Success; +} + // controlPatching implements the common internals of the patching/unpatching // implementation. |Enable| defines whether we're enabling or disabling the // runtime XRay instrumentation. XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT { if (!__sanitizer::atomic_load(&XRayInitialized, - __sanitizer::memory_order_acquire)) + __sanitizer::memory_order_acquire)) return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized. uint8_t NotPatching = false; @@ -179,38 +215,7 @@ ++FuncId; CurFun = F; } - - // While we're here, we should patch the nop sled. To do that we mprotect - // the page containing the function to be writeable. - void *PageAlignedAddr = - reinterpret_cast(Sled.Address & ~(PageSize - 1)); - std::size_t MProtectLen = (Sled.Address + cSledLength) - - reinterpret_cast(PageAlignedAddr); - MProtectHelper Protector(PageAlignedAddr, MProtectLen); - if (Protector.MakeWriteable() == -1) { - printf("Failed mprotect: %d\n", errno); - return XRayPatchingStatus::FAILED; - } - - bool Success = false; - switch (Sled.Kind) { - case XRayEntryType::ENTRY: - Success = patchFunctionEntry(Enable, FuncId, Sled, __xray_FunctionEntry); - break; - case XRayEntryType::EXIT: - Success = patchFunctionExit(Enable, FuncId, Sled); - break; - case XRayEntryType::TAIL: - Success = patchFunctionTailExit(Enable, FuncId, Sled); - break; - case XRayEntryType::LOG_ARGS_ENTRY: - Success = patchFunctionEntry(Enable, FuncId, Sled, __xray_ArgLoggerEntry); - break; - default: - Report("Unsupported sled kind: %d\n", int(Sled.Kind)); - continue; - } - (void)Success; + patchSled(Sled, Enable, FuncId); } __sanitizer::atomic_store(&XRayPatching, false, __sanitizer::memory_order_release); @@ -226,6 +231,61 @@ return controlPatching(false); } +XRayPatchingStatus patchSpecificFunction(int32_t FuncId, bool Enable) XRAY_NEVER_INSTRUMENT { + if (!__sanitizer::atomic_load(&XRayInitialized, + __sanitizer::memory_order_acquire)) + return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized. + + uint8_t NotPatching = false; + if (!__sanitizer::atomic_compare_exchange_strong( + &XRayPatching, &NotPatching, true, __sanitizer::memory_order_acq_rel)) + return XRayPatchingStatus::ONGOING; // Already patching. + + // Next, we look for the function index. + XRaySledMap InstrMap; + { + __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex); + InstrMap = XRayInstrMap; + } + + // If we don't have an index, we can't unpatch individual functions. + if (InstrMap.Functions == 0) + return XRayPatchingStatus::NOT_INITIALIZED; + + // FuncId may not be zero or negative. + if (FuncId <= 0) + return XRayPatchingStatus::FAILED; + + // Check whether the function id is out of bounds. + if (static_cast(FuncId) >= InstrMap.Functions) + return XRayPatchingStatus::FAILED; + + // Now we patch ths sleds for this specific function. + auto SledRange = InstrMap.SledsIndex[FuncId - 1]; + auto *f = reinterpret_cast(SledRange.Begin); + auto *e = reinterpret_cast(SledRange.End); + + // TODO: Figure out how we can undo partial patching, or signal partial + // success. + while (f != e) { + if (!patchSled(*f, Enable, FuncId)) + return XRayPatchingStatus::FAILED; + ++f; + } + + __sanitizer::atomic_store(&XRayPatching, false, + __sanitizer::memory_order_release); + return XRayPatchingStatus::SUCCESS; +} + +XRayPatchingStatus __xray_patch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT { + return patchSpecificFunction(FuncId, true); +} + +XRayPatchingStatus __xray_unpatch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT { + return patchSpecificFunction(FuncId, false); +} + int __xray_set_handler_arg1(void (*Handler)(int32_t, XRayEntryType, uint64_t)) { if (!__sanitizer::atomic_load(&XRayInitialized, __sanitizer::memory_order_acquire)) Index: lib/xray/xray_interface_internal.h =================================================================== --- lib/xray/xray_interface_internal.h +++ lib/xray/xray_interface_internal.h @@ -39,6 +39,18 @@ #error "Unsupported word size." #endif }; + +struct XRayFunctionSledIndex { +#if SANITIZER_WORDSIZE == 64 + uint64_t Begin; + uint64_t End; +#elif SANITIZER_WORDSIZE == 32 + uint32_t Begin; + uint32_t End; +#else +#error "Unsupported word size." +#endif +}; } namespace __xray { @@ -46,6 +58,8 @@ struct XRaySledMap { const XRaySledEntry *Sleds; size_t Entries; + const XRayFunctionSledIndex *SledsIndex; + size_t Functions; }; bool patchFunctionEntry(bool Enable, uint32_t FuncId, Index: test/xray/TestCases/Linux/coverage-sample.cc =================================================================== --- /dev/null +++ test/xray/TestCases/Linux/coverage-sample.cc @@ -0,0 +1,71 @@ +// Check that we can patch and unpatch specific function ids. +// +// RUN: %clangxx_xray -std=c++11 %s -o %t +// RUN: XRAY_OPTIONS="patch_premain=false xray_naive_log=false" %run %t | FileCheck %s + +#include "xray/xray_interface.h" + +#include +#include + +std::set function_ids; + +[[clang::xray_never_instrument]] void coverage_handler(int32_t fid, + XRayEntryType) { + static bool patching = false; + if (patching) return; + patching = true; + function_ids.insert(fid); + __xray_unpatch_function(fid); + patching = false; +} + +[[clang::xray_always_instrument]] void baz() { + // do nothing! +} + +[[clang::xray_always_instrument]] void bar() { + baz(); +} + +[[clang::xray_always_instrument]] void foo() { + bar(); +} + +[[clang::xray_always_instrument]] int main(int argc, char *argv[]) { + __xray_set_handler(coverage_handler); + __xray_patch(); + foo(); + __xray_unpatch(); + + // print out the function_ids. + for (const auto id : function_ids) + printf("patched: %d\n", id); + + // CHECK-DAG: patched: [[F1:.*]] + // CHECK-DAG: patched: [[F2:.*]] + // CHECK-DAG: patched: [[F3:.*]] + + // make a copy of the function_ids, then patch them later. + auto called_fns = function_ids; + + // clear the function_ids. + function_ids.clear(); + + // patch the functions we've called before. + for (const auto id : called_fns) + __xray_patch_function(id); + + // then call them again. + foo(); + __xray_unpatch(); + + // confirm that we've seen the same functions again. + for (const auto id : function_ids) + printf("patched: %d\n", id); + // CHECK-DAG: patched: [[F1]] + // CHECK-DAG: patched: [[F2]] + // CHECK-DAG: patched: [[F3]] + + return function_ids == called_fns ? 0 : 1; +}