diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index f87ebf2c..3a6eb564 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -1232,6 +1232,25 @@ mj_resetDataKeyframe Reset data, set fields from specified keyframe. +.. _mj_markStack: + +mj_markStack +~~~~~~~~~~~~ + +.. mujoco-include:: mj_markStack + +Mark a new frame on the :ref:`mjData` stack. + +.. _mj_freeStack: + +mj_freeStack +~~~~~~~~~~~~ + +.. mujoco-include:: mj_freeStack + +Free the current :ref:`mjData` stack frame. All pointers returned by mj_stackAlloc since the last call +to mj_markStack must no longer be used afterwards. + .. _mj_stackAlloc: mj_stackAlloc diff --git a/doc/changelog.rst b/doc/changelog.rst index cebd8ce5..db3b4ab9 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -43,6 +43,8 @@ General functionality of allocating ``mjtNum`` arrays is still available through the new function :ref:`mj_stackAllocNum`. #. Renamed the ``nstack`` field in :ref:`mjModel` and :ref:`mjData` to ``narena``. Changed ``narena``, ``pstack``, and ``maxuse_stack`` to count number of bytes rather than number of :ref:`mjtNum` |-| s. +#. Added new functions ``mj_markStack`` and ``mj_freeStack`` that manages ``mjData`` stack frames in a fully + encapsulated way (i.e. without having to introduce a local variable at the call site). Python bindings ^^^^^^^^^^^^^^^ diff --git a/doc/includes/references.h b/doc/includes/references.h index e6f81619..792c5dcf 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -130,6 +130,7 @@ struct mjData_ { // stack pointer size_t pstack; // first available mjtNum address in stack + size_t pbase; // value of pstack when mj_markStack was last called // arena pointer size_t parena; // first available byte in arena @@ -2205,6 +2206,8 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src); void mj_resetData(const mjModel* m, mjData* d); void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_value); void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key); +void mj_markStack(mjData* d); +void mj_freeStack(mjData* d); void* mj_stackAlloc(mjData* d, size_t bytes, size_t alignment); mjtNum* mj_stackAllocNum(mjData* d, int size); int* mj_stackAllocInt(mjData* d, int size); diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 756b7675..51522ca3 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -157,6 +157,7 @@ struct mjData_ { // stack pointer size_t pstack; // first available mjtNum address in stack + size_t pbase; // value of pstack when mj_markStack was last called // arena pointer size_t parena; // first available byte in arena diff --git a/include/mujoco/mjmacro.h b/include/mujoco/mjmacro.h index 713d5c52..e3972015 100644 --- a/include/mujoco/mjmacro.h +++ b/include/mujoco/mjmacro.h @@ -15,8 +15,6 @@ #ifndef MUJOCO_MJMACRO_H_ #define MUJOCO_MJMACRO_H_ -#include - // include asan interface header, or provide stubs for poison/unpoison macros when not using asan #ifdef ADDRESS_SANITIZER #include @@ -33,8 +31,8 @@ #define mjMIN(a, b) (((a) < (b)) ? (a) : (b)) // mjData stack frame management -#define mjMARKSTACK size_t _mark = d->pstack; -#define mjFREESTACK d->pstack = _mark; +#define mjMARKSTACK mj_markStack(d); +#define mjFREESTACK mj_freeStack(d); // return current value of mjOption enable/disable flags #define mjDISABLED(x) (m->opt.disableflags & (x)) @@ -49,15 +47,4 @@ #endif #endif -// implementation of mjFREESTACK when using the address sanitizer -#ifdef ADDRESS_SANITIZER - #undef mjFREESTACK - #define mjFREESTACK { \ - d->pstack = _mark; \ - ASAN_POISON_MEMORY_REGION( \ - (char*)d->arena + d->parena, \ - d->narena - d->pstack - d->parena); \ - } -#endif - #endif // MUJOCO_MJMACRO_H_ diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 7d5c5724..bccfcde2 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -625,6 +625,7 @@ X( size_t, nbuffer ) \ X( int, nplugin ) \ X( size_t, pstack ) \ + X( size_t, pbase ) \ X( size_t, parena ) \ X( size_t, maxuse_stack ) \ X( size_t, maxuse_arena ) \ diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index d3c1841a..ce813376 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -30,10 +30,6 @@ extern "C" { #include #include -#ifdef ADDRESS_SANITIZER -#include -#endif - // type definitions #include #include @@ -188,6 +184,13 @@ MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_va // Reset data, set fields from specified keyframe. MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key); +// Mark a new frame on the mjData stack. +MJAPI void mj_markStack(mjData* d); + +// Free the current mjData stack frame. All pointers returned by mj_stackAlloc since the last call +// to mj_markStack must no longer be used afterwards. +MJAPI void mj_freeStack(mjData* d); + // Allocate a number of bytes on mjData stack at a specific alignment. // Call mju_error on stack overflow. MJAPI void* mj_stackAlloc(mjData* d, size_t bytes, size_t alignment); diff --git a/introspect/functions.py b/introspect/functions.py index 4ee0d35c..984f4eda 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -677,6 +677,34 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Reset data, set fields from specified keyframe.', )), + ('mj_markStack', + FunctionDecl( + name='mj_markStack', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='d', + type=PointerType( + inner_type=ValueType(name='mjData'), + ), + ), + ), + doc='Mark a new frame on the mjData stack.', + )), + ('mj_freeStack', + FunctionDecl( + name='mj_freeStack', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='d', + type=PointerType( + inner_type=ValueType(name='mjData'), + ), + ), + ), + doc='Free the current mjData stack frame. All pointers returned by mj_stackAlloc since the last call to mj_markStack must no longer be used afterwards.', # pylint: disable=line-too-long + )), ('mj_stackAlloc', FunctionDecl( name='mj_stackAlloc', diff --git a/introspect/structs.py b/introspect/structs.py index 5ac4330a..1573fb0c 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -3502,6 +3502,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='size_t'), doc='first available mjtNum address in stack', ), + StructFieldDecl( + name='pbase', + type=ValueType(name='size_t'), + doc='value of pstack when mj_markStack was last called', + ), StructFieldDecl( name='parena', type=ValueType(name='size_t'), diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 668891a3..f713c718 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -35,6 +35,11 @@ #include "engine/engine_util_misc.h" #include "engine/engine_vfs.h" +#ifdef ADDRESS_SANITIZER + #include + #include +#endif + #ifdef MEMORY_SANITIZER #include #endif @@ -43,6 +48,13 @@ #pragma warning (disable: 4305) // disable MSVC warning: truncation from 'double' to 'float' #endif +// add red zone padding when built with asan, to detect out-of-bound accesses +#ifdef ADDRESS_SANITIZER + #define mjREDZONE 32 +#else + #define mjREDZONE 0 +#endif + static const int MAX_ARRAY_SIZE = INT_MAX / 4; // compute a % b with a fast code path if the second argument is a power of 2 @@ -54,6 +66,12 @@ static inline size_t fastmod(size_t a, size_t b) { return a % b; } +typedef struct { + size_t pbase; // value of d->pbase immediately before mj_markStack + size_t pstack; // value of d->pstack immediately before mj_markStack + void* pc; // program counter of the call site of mj_markStack (only set when under asan) +} mjStackFrame; + //------------------------------ mjLROpt ----------------------------------------------------------- // set default options for length range computation @@ -1226,13 +1244,6 @@ static inline void* stackalloc(mjData* d, size_t size, size_t alignment) { return NULL; } - // add red zone padding when built with asan, to detect out-of-bound accesses -#ifdef ADDRESS_SANITIZER - #define mjREDZONE 32 -#else - #define mjREDZONE 0 -#endif - // size of entire arena/stack in bytes size_t stack_size_bytes = d->narena; @@ -1284,8 +1295,6 @@ static inline void* stackalloc(mjData* d, size_t size, size_t alignment) { ASAN_UNPOISON_MEMORY_REGION((void*)start_ptr, size); #endif -#undef mjREDZONE - // update pstack and max usage statistics d->pstack = new_pstack; d->maxuse_stack = mjMAX(d->maxuse_stack, usage); @@ -1294,6 +1303,73 @@ static inline void* stackalloc(mjData* d, size_t size, size_t alignment) { return (void*)start_ptr; } + + +// mjData mark stack frame +#ifdef ADDRESS_SANITIZER +__attribute__((noinline)) +#endif +void mj_markStack(mjData* d) { + size_t pstack_old = d->pstack; + mjStackFrame* s = + (mjStackFrame*) stackalloc(d, sizeof(mjStackFrame), _Alignof(mjStackFrame)); + s->pbase = d->pbase; + s->pstack = pstack_old; +#ifdef ADDRESS_SANITIZER + // store the program counter to the caller so that we can compare against mj_freeStack later + s->pc = __sanitizer_return_address(); +#endif + d->pbase = d->pstack - mjREDZONE; +} + + + +// mjData free stack frame +#ifdef ADDRESS_SANITIZER +__attribute__((noinline)) +#endif +void mj_freeStack(mjData* d) { + if (mjUNLIKELY(!d->pbase)) { + return; + } + + mjStackFrame* s = (mjStackFrame*) ((char*)d->arena + d->narena - d->pbase); +#ifdef ADDRESS_SANITIZER + #define mjSYMBOLIZELEN 256 + + // symbolize s->pc to get the function name of most recent caller to mj_markStack + char markstack_func[mjSYMBOLIZELEN]; + __sanitizer_symbolize_pc(s->pc, "%f", markstack_func, mjSYMBOLIZELEN); + markstack_func[mjSYMBOLIZELEN - 1] = '\0'; + + // symbolize current program counter to get the function name of caller to this function + char freestack_func[mjSYMBOLIZELEN]; + __sanitizer_symbolize_pc(__sanitizer_return_address(), "%f", freestack_func, mjSYMBOLIZELEN); + freestack_func[mjSYMBOLIZELEN - 1] = '\0'; + + // raise an error if caller function name doesn't match the most recent caller of mj_markStack + if (strncmp(markstack_func, freestack_func, mjSYMBOLIZELEN)) { + 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 +#endif + + // restore pbase and pstack + d->pbase = s->pbase; + d->pstack = s->pstack; + + // if running under asan, poison the newly freed memory region +#ifdef ADDRESS_SANITIZER + ASAN_POISON_MEMORY_REGION((char*)d->arena + d->parena, d->narena - d->pstack - d->parena); +#endif +} + void* mj_stackAlloc(mjData* d, size_t bytes, size_t alignment) { return stackalloc(d, bytes, alignment); } @@ -1320,6 +1396,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { // clear stack pointer d->pstack = 0; + d->pbase = 0; // clear arena pointers d->parena = 0; @@ -1487,6 +1564,11 @@ void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key) { // de-allocate mjData void mj_deleteData(mjData* d) { if (d) { +#ifdef ADDRESS_SANITIZER + // raise an error if there's a dangling stack frame + mj_freeStack(d); +#endif + // destroy plugin instances for (int i = 0; i < d->nplugin; ++i) { const mjpPlugin* plugin = mjp_getPluginAtSlot(d->plugin[i]); diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h index 7ad94015..b0595747 100644 --- a/src/engine/engine_io.h +++ b/src/engine/engine_io.h @@ -106,6 +106,12 @@ MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key); // mjData arena allocate MJAPI void* mj_arenaAlloc(mjData* d, size_t bytes, size_t alignment); +// mjData mark stack frame +MJAPI void mj_markStack(mjData* d); + +// mjData free stack frame +MJAPI void mj_freeStack(mjData* d); + // mjData stack allocate MJAPI void* mj_stackAlloc(mjData* d, size_t bytes, size_t alignment); diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 8cd95131..d9455c4d 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -730,6 +730,11 @@ void mj_printModel(const mjModel* m, const char* filename) { // valid printf-style format string for a single float value void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, const char* float_format) { + // stack in use, SHOULD NOT OCCUR + if (d->pstack) { + mjERROR("attempting to print mjData when stack is in use"); + } + mjtNum *M; mjMARKSTACK; @@ -739,11 +744,6 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, float_format = FLOAT_FORMAT; } - // stack in use, SHOULD NOT OCCUR - if (d->pstack) { - mjERROR("attempting to print mjData when stack is in use"); - } - // get file FILE* fp; if (filename) { @@ -776,7 +776,9 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, fprintf(fp, "SIZES\n"); #define X(type, name) \ - if (strcmp(#name, "pstack") != 0 && strcmp(#name, "parena") != 0) { \ + if (strcmp(#name, "pstack") != 0 && \ + strcmp(#name, "pbase") != 0 && \ + strcmp(#name, "parena") != 0) { \ const char* format = _Generic( \ d->name, \ int : INT_FORMAT, \ diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc index 4ccebcce..c73b28e3 100644 --- a/test/engine/engine_io_test.cc +++ b/test/engine/engine_io_test.cc @@ -17,24 +17,26 @@ #include "src/engine/engine_io.h" #include -#include +#include #include #include #include -#include -#include #include +#include #include #include +#include // IWYU pragma: keep #include #include +#include #include "src/engine/engine_util_errmem.h" #include "test/fixture.h" namespace mujoco { namespace { +using ::testing::ContainsRegex; // NOLINT(misc-unused-using-decls) asan only using ::testing::HasSubstr; using ::testing::IsNull; using ::testing::NotNull; @@ -737,5 +739,81 @@ TEST_F(ValidateReferencesTest, Tuples) { mj_deleteModel(model); } +TEST_F(EngineIoTest, CanMarkAndFreeStack) { + constexpr char xml[] = R"( + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); + + mjData* data = mj_makeData(model); + ASSERT_THAT(data, NotNull()); + + auto pstack_before = data->pstack; + mj_markStack(data); + EXPECT_GT(data->pstack, pstack_before); + mj_freeStack(data); + EXPECT_EQ(data->pstack, pstack_before); + + mj_deleteData(data); + mj_deleteModel(model); +} + +#ifdef ADDRESS_SANITIZER +void MarkFreeStack(mjData* d, bool free) { + mj_markStack(d); + if (free) { + mj_freeStack(d); + } +} + +TEST_F(EngineIoTest, CanDetectStackFrameLeakage) { + static constexpr char xml[] = R"( + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); + + mjData* data = mj_makeData(model); + ASSERT_THAT(data, NotNull()); + + // MarkFreeStack correctly calls mj_freeStack, should not error. + MarkFreeStack(data, /* free= */ true); + + // MarkFreeStack calls mj_markStack without mj_freeStack, the next call to + // mj_freeStack should detect the stack frame leakage. + mj_markStack(data); + MarkFreeStack(data, /* free= */ false); + EXPECT_THAT( + MjuErrorMessageFrom(mj_freeStack)(data), + ContainsRegex( + "mj_markStack in MarkFreeStack at .*engine_io_test\\.cc.* has no " + "corresponding mj_freeStack")); + + // Dangling stack frames should be detected in mj_deleteData. + mj_resetData(model, data); + mj_markStack(data); + EXPECT_THAT( + MjuErrorMessageFrom(mj_deleteData)(data), + ContainsRegex( + "mj_markStack in .+EngineIoTest_CanDetectStackFrameLeakage.+ has no " + "corresponding mj_freeStack")); + + mj_resetData(model, data); + mj_deleteData(data); + mj_deleteModel(model); +} +#endif + } // namespace } // namespace mujoco diff --git a/test/engine/engine_util_container_test.cc b/test/engine/engine_util_container_test.cc index bd9766e2..42408739 100644 --- a/test/engine/engine_util_container_test.cc +++ b/test/engine/engine_util_container_test.cc @@ -27,18 +27,21 @@ namespace { using testing::NotNull; -template +template constexpr int GetExpectedStackUsageBytes() { if constexpr (N <= 0) { - return 0; + return prev_size; } else { constexpr auto RoundUpToAlignment = [](int x, int alignment) { return alignment * (x / alignment + ((x % alignment) ? 1 : 0)); }; - return RoundUpToAlignment(sizeof(mjArrayList), alignof(mjArrayList)) + - RoundUpToAlignment(Capacity * sizeof(T), alignof(std::max_align_t)) + - GetExpectedStackUsageBytes(); + constexpr int size_with_arraylist = RoundUpToAlignment( + prev_size + sizeof(mjArrayList), alignof(mjArrayList)); + constexpr int size_with_buffer = RoundUpToAlignment( + size_with_arraylist + capacity * sizeof(T), alignof(std::max_align_t)); + return GetExpectedStackUsageBytes(); } } @@ -48,6 +51,7 @@ TEST(TestMjArrayList, TestMjArrayListSingleThreaded) { ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error.data(); mjData* d = mj_makeData(m); mjMARKSTACK; + using DataType = int; constexpr int kInitialCapacity = 10; mjArrayList* array_list = @@ -59,8 +63,10 @@ TEST(TestMjArrayList, TestMjArrayListSingleThreaded) { } EXPECT_EQ(mju_arrayListSize(array_list), kNumElements); + constexpr int kFrameMarkerSize = 2 * sizeof(size_t) + sizeof(void*); constexpr int kExpectedMaxUseStack = - GetExpectedStackUsageBytes(); + GetExpectedStackUsageBytes(); EXPECT_EQ(d->maxuse_stack, kExpectedMaxUseStack); for (int i = 0; i < kNumElements; ++i) { diff --git a/test/fixture.h b/test/fixture.h index 83111f1c..cdf097ec 100644 --- a/test/fixture.h +++ b/test/fixture.h @@ -15,12 +15,23 @@ #ifndef MUJOCO_TEST_FIXTURE_H_ #define MUJOCO_TEST_FIXTURE_H_ +#include +#include +#include +#include +#include + #include #include #include #include #include +extern "C" { +MJAPI void _mjPRIVATE__set_tls_error_fn(decltype(mju_user_error)); +MJAPI decltype(mju_user_error) _mjPRIVATE__get_tls_error_fn(); +} + namespace mujoco { // Installs and uninstalls error callbacks on MuJoCo that fail the currently @@ -44,6 +55,30 @@ class MujocoTest : public ::testing::Test { MujocoErrorTestGuard error_guard; }; +template +auto MjuErrorMessageFrom(Return (*func)(Args...)) { + thread_local std::jmp_buf current_jmp_buf; + thread_local char err_msg[1000]; + + auto* old_error_handler = _mjPRIVATE__get_tls_error_fn(); + auto* new_error_handler = +[](const char* msg) -> void { + std::strncpy(err_msg, msg, sizeof(err_msg)); + std::longjmp(current_jmp_buf, 1); + }; + + return [func, old_error_handler, + new_error_handler](Args... args) -> std::string { + if (setjmp(current_jmp_buf) == 0) { + err_msg[0] = '\0'; + _mjPRIVATE__set_tls_error_fn(new_error_handler); + func(args...); + } + + _mjPRIVATE__set_tls_error_fn(old_error_handler); + return err_msg; + }; +} + // Returns a path to a data file, under the mujoco/test directory. const std::string GetTestDataFilePath(absl::string_view path); diff --git a/test/xml/xml_native_reader_test.cc b/test/xml/xml_native_reader_test.cc index 09d5c519..dad4144e 100644 --- a/test/xml/xml_native_reader_test.cc +++ b/test/xml/xml_native_reader_test.cc @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -43,12 +44,12 @@ TEST_F(XMLReaderTest, MemorySize) { { static constexpr char xml[] = R"( - + )"; mjModel* model = LoadModelFromString(xml, error.data(), error.size()); ASSERT_THAT(model, NotNull()) << error.data(); - EXPECT_EQ(model->narena, 256); + EXPECT_EQ(model->narena, 512); mj_deleteModel(model); } { diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 7b318b97..e886c3c5 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -574,6 +574,7 @@ public unsafe struct mjData_ { public UIntPtr nbuffer; public int nplugin; public UIntPtr pstack; + public UIntPtr pbase; public UIntPtr parena; public UIntPtr maxuse_stack; public UIntPtr maxuse_arena; @@ -3087,6 +3088,12 @@ public static unsafe extern void mj_resetDataDebug(mjModel_* m, mjData_* d, byte [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mj_resetDataKeyframe(mjModel_* m, mjData_* d, int key); +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mj_markStack(mjData_* d); + +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mj_freeStack(mjData_* d); + [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void* mj_stackAlloc(mjData_* d, UIntPtr bytes, UIntPtr alignment);