From b04de39e4e07f7a988aabd390cad31806a5ec153 Mon Sep 17 00:00:00 2001 From: Saran Tunyasuvunakool Date: Thu, 22 Feb 2024 08:02:46 -0800 Subject: [PATCH] Fix false positives from mjData stack leakage detection. When built and run under address sanitizer (asan), `mj_markStack` and `mj_freeStack` are instrumented to detect leakage of mjData stack frames. When the compiler inlines several callees that call into mark/free in the same function, this instrumentation requires that the compiler retains separate mark/free calls for each original callee. This change adds necessary optimization barriers to ensure that this is the case. PiperOrigin-RevId: 609373993 Change-Id: I30dfd998fb66530735348e9715d92ee713f28210 --- include/mujoco/mujoco.h | 36 +++++++++++++++++++-- src/engine/engine_crossplatform.cc | 52 +++++++++++++++++++++++++++--- src/engine/engine_crossplatform.h | 3 +- src/engine/engine_io.c | 29 ++++++++--------- src/engine/engine_io.h | 9 ++++++ 5 files changed, 107 insertions(+), 22 deletions(-) diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index f32d6b5e..2e11f2ee 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -19,7 +19,7 @@ // this is a C-API -#if defined(__cplusplus) +#ifdef __cplusplus extern "C" { #endif @@ -187,6 +187,8 @@ MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_va // Reset data. If 0 <= key < nkey, set fields from specified keyframe. MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key); +#ifndef ADDRESS_SANITIZER + // Mark a new frame on the mjData stack. MJAPI void mj_markStack(mjData* d); @@ -194,6 +196,8 @@ MJAPI void mj_markStack(mjData* d); // to mj_markStack must no longer be used afterwards. MJAPI void mj_freeStack(mjData* d); +#endif // ADDRESS_SANITIZER + // Allocate a number of bytes on mjData stack at a specific alignment. // Call mju_error on stack overflow. MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment); @@ -1340,8 +1344,36 @@ MJAPI void mju_defaultTask(mjTask* task); // Wait for a task to complete. MJAPI void mju_taskJoin(mjTask* task); +//---------------------- Sanitizer instrumentation helpers ----------------------------------------- +// +// Most MuJoCo users can ignore these functions, the following comments are aimed primarily at +// MuJoCo developers. +// +// When built and run under address sanitizer (asan), mj_markStack and mj_freeStack are instrumented +// to detect leakage of mjData stack frames. When the compiler inlines several callees that call +// into mark/free into the same function, this instrumentation requires that the compiler retains +// separate mark/free calls for each original callee. The memory-clobbered asm blocks act as a +// barrier to prevent mark/free calls from being combined under optimization. -#if defined(__cplusplus) +#ifdef ADDRESS_SANITIZER + +void mj__markStack(mjData*) __attribute__((noinline)); +static inline void mj_markStack(mjData* d) __attribute__((always_inline)) { + asm volatile("" ::: "memory"); + mj__markStack(d); + asm volatile("" ::: "memory"); +} + +void mj__freeStack(mjData*) __attribute__((noinline)); +static inline void mj_freeStack(mjData* d) __attribute__((always_inline)) { + asm volatile("" ::: "memory"); + mj__freeStack(d); + asm volatile("" ::: "memory"); +} + +#endif // ADDRESS_SANITIZER + +#ifdef __cplusplus } #endif diff --git a/src/engine/engine_crossplatform.cc b/src/engine/engine_crossplatform.cc index 9ddcca6d..c30b3c0b 100644 --- a/src/engine/engine_crossplatform.cc +++ b/src/engine/engine_crossplatform.cc @@ -53,12 +53,23 @@ void CheckRosetta() { #include #include #include +#include #include namespace { -std::string_view SymbolizeCached(void* pc) { +const std::pair& +FuncNameAndDebugInfoCached(void* pc) { + static const std::unordered_set* const kIgnoredInlinedFunctions = + []() { + return new std::unordered_set{ + "mj_freeStack", + "mj_markStack", + }; + }(); + static auto* mu = new std::shared_mutex; - static auto* pc_to_func_name_map = new std::unordered_map; + static auto* pc_to_func_name_map = + new std::unordered_map>; { std::shared_lock lock(*mu); @@ -70,14 +81,43 @@ std::string_view SymbolizeCached(void* pc) { std::array buf; __sanitizer_symbolize_pc(pc, "%f", buf.data(), buf.size()); + + // buf contains sequence of null-terminated strings of inlined function names + // so we walk through the sequence until we find the first unignored function + std::string_view func_name(buf.data()); + int idx = 0; + while (kIgnoredInlinedFunctions->find(func_name.data()) != + kIgnoredInlinedFunctions->end()) { + func_name = func_name.data() + func_name.size() + 1; + ++idx; + } + + std::array buf2; + __sanitizer_symbolize_pc(pc, "%F at %S", buf2.data(), buf2.size()); + std::string_view debug_info(buf2.data()); + for (int i = 0; i < idx; ++i) { + debug_info = debug_info.data() + debug_info.size() + 1; + } + { std::unique_lock lock(*mu); - return pc_to_func_name_map->emplace(pc, buf.data()).first->second; + return pc_to_func_name_map + ->emplace( + pc, std::make_pair(std::string(func_name), std::string(debug_info))) + .first->second; } } + +const std::string& SymbolizeCached(void* pc) { + return FuncNameAndDebugInfoCached(pc).first; +} + +const std::string& DebugInfoCached(void* pc) { + return FuncNameAndDebugInfoCached(pc).second; +} } // namespace -int _mj_comparePcFuncName(void* pc1, void* pc2) { +int mj__comparePcFuncName(void* pc1, void* pc2) { static auto* mu = new std::shared_mutex; static auto* same_func_map = new std::map, bool>; @@ -96,4 +136,8 @@ int _mj_comparePcFuncName(void* pc1, void* pc2) { return same_func_map->emplace(pc_pair, is_same).first->second; } } + +const char* mj__getPcDebugInfo(void* pc) { + return DebugInfoCached(pc).c_str(); +} #endif // ADDRESS_SANITIZER diff --git a/src/engine/engine_crossplatform.h b/src/engine/engine_crossplatform.h index d52e12e9..5d506430 100644 --- a/src/engine/engine_crossplatform.h +++ b/src/engine/engine_crossplatform.h @@ -80,7 +80,8 @@ extern "C" { #endif #ifdef ADDRESS_SANITIZER -int _mj_comparePcFuncName(void* pc1, void* pc2); +int mj__comparePcFuncName(void* pc1, void* pc2); +const char* mj__getPcDebugInfo(void* pc); #endif #ifdef __cplusplus diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 7704e363..dbba8957 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -1439,10 +1439,12 @@ static inline void markstackinternal(mjData* d, mjStackInfo* stack_info) { // mjData mark stack frame -#ifdef ADDRESS_SANITIZER -__attribute__((noinline)) +#ifndef ADDRESS_SANITIZER +void mj_markStack(mjData* d) +#else +void mj__markStack(mjData* d) #endif -void mj_markStack(mjData* d) { +{ if (!d->threadpool) { mjStackInfo stack_info = get_stack_info_from_data(d); markstackinternal(d, &stack_info); @@ -1469,15 +1471,10 @@ static inline void freestackinternal(mjStackInfo* stack_info) { mjStackFrame* s = (mjStackFrame*) stack_info->stack_base; #ifdef ADDRESS_SANITIZER // raise an error if caller function name doesn't match the most recent caller of mj_markStack - if (!_mj_comparePcFuncName(s->pc, __sanitizer_return_address())) { - #define mjSYMBOLIZELEN 256 - char dbginfo[mjSYMBOLIZELEN]; - __sanitizer_symbolize_pc( - s->pc, "mj_markStack %F at %S has no corresponding mj_freeStack", - dbginfo, sizeof(dbginfo)); - dbginfo[mjSYMBOLIZELEN - 1] = '\0'; - mjERROR("%s", dbginfo); - #undef mjSYMBOLIZELEN + if (!mj__comparePcFuncName(s->pc, __sanitizer_return_address())) { + mjERROR("mj_markStack %s has no corresponding mj_freeStack (detected %s)", + mj__getPcDebugInfo(s->pc), + mj__getPcDebugInfo(__sanitizer_return_address())); } #endif @@ -1494,10 +1491,12 @@ static inline void freestackinternal(mjStackInfo* stack_info) { // mjData free stack frame -#ifdef ADDRESS_SANITIZER -__attribute__((noinline)) +#ifndef ADDRESS_SANITIZER +void mj_freeStack(mjData* d) +#else +void mj__freeStack(mjData* d) #endif -void mj_freeStack(mjData* d) { +{ if (!d->threadpool) { mjStackInfo stack_info = get_stack_info_from_data(d); freestackinternal(&stack_info); diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h index 0bc8aaed..be9f5e15 100644 --- a/src/engine/engine_io.h +++ b/src/engine/engine_io.h @@ -110,12 +110,21 @@ MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key); // mjData arena allocate MJAPI void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment); +#ifndef ADDRESS_SANITIZER + // mjData mark stack frame MJAPI void mj_markStack(mjData* d); // mjData free stack frame MJAPI void mj_freeStack(mjData* d); +#else + +void mj__markStack(mjData* d) __attribute__((noinline)); +void mj__freeStack(mjData* d) __attribute__((noinline)); + +#endif // ADDRESS_SANITIZER + // mjData stack allocate MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);