chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add mj_markStack and mj_freeStack as public API functions.

Browse Source 
Also add asan instrumentation to detect stack frame leakages (i.e. `mj_markStack` without a corresponding `mj_freeStack` in the same caller function).

PiperOrigin-RevId: 562625645
Change-Id: I4e3ff66ca0b9d08ed0a95cef45393db8e3053e22
This commit is contained in:
Saran Tunyasuvunakool
2023-09-04 17:03:23 -07:00
committed by Copybara-Service
parent 9308e1d383
commit 94a8705ad0
17 changed files with 311 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference/functions.rst
+19
View File
@@ -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
doc/changelog.rst
+2
View File
@@ -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
^^^^^^^^^^^^^^^
doc/includes/references.h
+3
View File
@@ -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);
include/mujoco/mjdata.h
+1
View File
@@ -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
include/mujoco/mjmacro.h
+2 -15
View File
@@ -15,8 +15,6 @@
#ifndef MUJOCO_MJMACRO_H_
#define MUJOCO_MJMACRO_H_
#include <stddef.h>
// include asan interface header, or provide stubs for poison/unpoison macros when not using asan
#ifdef ADDRESS_SANITIZER
#include <sanitizer/asan_interface.h>
@@ -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_
include/mujoco/mjxmacro.h
+1
View File
@@ -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 ) \
include/mujoco/mujoco.h
+7 -4
View File
@@ -30,10 +30,6 @@ extern "C" {
#include <stdlib.h>
#include <math.h>
#ifdef ADDRESS_SANITIZER
#include <sanitizer/asan_interface.h>
#endif
// type definitions
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
@@ -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);
introspect/functions.py
+28
View File
@@ -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',
introspect/structs.py
+5
View File
@@ -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'),
src/engine/engine_io.c
+91 -9
View File
@@ -35,6 +35,11 @@
#include "engine/engine_util_misc.h"
#include "engine/engine_vfs.h"
#ifdef ADDRESS_SANITIZER
#include <sanitizer/asan_interface.h>
#include <sanitizer/common_interface_defs.h>
#endif
#ifdef MEMORY_SANITIZER
#include <sanitizer/msan_interface.h>
#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]);
src/engine/engine_io.h
+6
View File
@@ -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);
src/engine/engine_print.c
+8 -6
View File
@@ -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, \
test/engine/engine_io_test.cc
+81 -3
View File
@@ -17,24 +17,26 @@
#include "src/engine/engine_io.h"
#include <array>
#include <climits>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <gtest/gtest-spi.h> // IWYU pragma: keep
#include <absl/strings/str_format.h>
#include <mujoco/mjxmacro.h>
#include <mujoco/mujoco.h>
#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"(
<mujoco>
<worldbody>
</worldbody>
</mujoco>
)";
std::array<char, 1024> 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"(
<mujoco>
<worldbody>
</worldbody>
</mujoco>
)";
std::array<char, 1024> 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
test/engine/engine_util_container_test.cc
+12 -6
View File
@@ -27,18 +27,21 @@ namespace {
using testing::NotNull;
template <typename T, int N, int Capacity>
template <int prev_size, typename T, int N, int capacity>
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<T, N - Capacity, 2 * Capacity>();
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<size_with_buffer, T, N - capacity,
2 * capacity>();
}
}
@@ -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<DataType, kNumElements, kInitialCapacity>();
GetExpectedStackUsageBytes<kFrameMarkerSize, DataType, kNumElements,
kInitialCapacity>();
EXPECT_EQ(d->maxuse_stack, kExpectedMaxUseStack);
for (int i = 0; i < kNumElements; ++i) {
test/fixture.h
+35
View File
@@ -15,12 +15,23 @@
#ifndef MUJOCO_TEST_FIXTURE_H_
#define MUJOCO_TEST_FIXTURE_H_
#include <csetjmp>
#include <cstring>
#include <optional>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include <absl/strings/string_view.h>
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mujoco.h>
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 <typename Return, typename... Args>
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);
test/xml/xml_native_reader_test.cc
+3 -2
View File
@@ -17,6 +17,7 @@
#include <array>
#include <limits>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
@@ -43,12 +44,12 @@ TEST_F(XMLReaderTest, MemorySize) {
{
static constexpr char xml[] = R"(
<mujoco>
<size memory="256"/>
<size memory="512"/>
</mujoco>
)";
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);
}
{
unity/Runtime/Bindings/MjBindings.cs
+7
View File
@@ -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);