Add ASAN and MSAN instrumentation for arena and mjData stack.

PiperOrigin-RevId: 535635490
Change-Id: Ia1e58771bae25bf72c861f57452fb8faa2b918fb

src/engine/engine_collision_box.c

@@ -15,6 +15,7 @@
 #include <math.h>
 #include <string.h>
 
+#include <mujoco/mjmacro.h>
 #include "engine/engine_collision_primitive.h"
 #include "engine/engine_util_blas.h"
 

src/engine/engine_collision_convex.c

@@ -20,9 +20,9 @@
 #include <ccd/vec3.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_collision_primitive.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_collision_driver.c

@@ -18,6 +18,7 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_callback.h"
 #include "engine/engine_collision_convex.h"
@@ -25,7 +26,6 @@
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_collision_primitive.c

@@ -17,6 +17,7 @@
 #include <math.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_spatial.h"

src/engine/engine_core_constraint.c

@@ -19,12 +19,12 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_array_safety.h"
 #include "engine/engine_core_smooth.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
@@ -32,6 +32,10 @@
 #include "engine/engine_util_sparse.h"
 #include "engine/engine_util_spatial.h"
 
+#ifdef MEMORY_SANITIZER
+  #include <sanitizer/msan_interface.h>
+#endif
+
 #ifdef mjUSEPLATFORMSIMD
   #if defined(__AVX__) && defined(mjUSEDOUBLE)
     #define mjUSEAVX
@@ -127,6 +131,10 @@ int mj_addContact(const mjModel* m, mjData* d, const mjContact* con) {
 
   // move arena pointer back to the end of the existing contact array and invalidate efc_ arrays
   d->parena = d->ncon * sizeof(mjContact);
+#ifdef ADDRESS_SANITIZER
+  ASAN_POISON_MEMORY_REGION(
+      (char*)d->arena + d->parena, (d->nstack - d->pstack) * sizeof(mjtNum) - d->parena);
+#endif
   clearEfc(d);
 
   // copy contact
@@ -1580,6 +1588,10 @@ void mj_makeConstraint(const mjModel* m, mjData* d) {
 
   // move arena pointer to end of contact array
   d->parena = d->ncon * sizeof(mjContact);
+#ifdef ADDRESS_SANITIZER
+  ASAN_POISON_MEMORY_REGION(
+      (char*)d->arena + d->parena, (d->nstack - d->pstack) * sizeof(mjtNum) - d->parena);
+#endif
 
 #define X(type, name, nr, nc)                                             \
   d->name = mj_arenaAlloc(d, sizeof(type) * (nr) * (nc), _Alignof(type)); \

src/engine/engine_core_smooth.c

@@ -18,11 +18,11 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"

src/engine/engine_derivative.c

@@ -15,11 +15,11 @@
 #include "engine/engine_derivative.h"
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_passive.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"

src/engine/engine_derivative_fd.c

@@ -18,6 +18,7 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_forward.h"
 #include "engine/engine_io.h"

src/engine/engine_forward.c

@@ -18,6 +18,7 @@
 #include <stdio.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
 #include "engine/engine_callback.h"

src/engine/engine_inverse.c

@@ -17,6 +17,7 @@
 #include <stddef.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_collision_driver.h"
 #include "engine/engine_core_constraint.h"

src/engine/engine_io.c

@@ -22,6 +22,7 @@
 #include <string.h>
 
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_array_safety.h"
@@ -33,6 +34,10 @@
 #include "engine/engine_util_misc.h"
 #include "engine/engine_vfs.h"
 
+#ifdef MEMORY_SANITIZER
+  #include <sanitizer/msan_interface.h>
+#endif
+
 #ifdef _MSC_VER
   #pragma warning (disable: 4305)  // disable MSVC warning: truncation from 'double' to 'float'
 #endif
@@ -1163,6 +1168,15 @@ void* mj_arenaAlloc(mjData* d, int bytes, int alignment) {
   void* result = (char*)d->arena + d->parena + padding;
   d->parena += padding + bytes;
   d->maxuse_arena = mjMAX(d->maxuse_arena, d->pstack*sizeof(mjtNum) + d->parena);
+
+#ifdef ADDRESS_SANITIZER
+  ASAN_UNPOISON_MEMORY_REGION(result, bytes);
+#endif
+
+#ifdef MEMORY_SANITIZER
+  __msan_allocated_memory(result, bytes);
+#endif
+
   return result;
 }
 
@@ -1172,12 +1186,19 @@ void* mj_arenaAlloc(mjData* d, int bytes, int alignment) {
 mjtNum* mj_stackAlloc(mjData* d, int size) {
   // return NULL if empty
   if (!size) {
-    return 0;
+    return NULL;
   }
 
+  // add red zone padding when built with asan, to detect out-of-bound accesses
+#ifdef ADDRESS_SANITIZER
+  #define mjREDZONE 4
+#else
+  #define mjREDZONE 0
+#endif
+
   // check size
   size_t stack_available_bytes = d->nstack * sizeof(mjtNum) - d->parena;
-  size_t stack_required_bytes = (d->pstack + size) * sizeof(mjtNum);
+  size_t stack_required_bytes = (d->pstack + size + 2*mjREDZONE) * sizeof(mjtNum);
   if (stack_required_bytes > stack_available_bytes) {
     mju_error("stack overflow: max = %zu, available = %zu, requested = %zu "
               "(ne = %d, nf = %d, nefc = %d, ncon = %d)",
@@ -1187,18 +1208,45 @@ mjtNum* mj_stackAlloc(mjData* d, int size) {
 
   // allocate at end of arena
   char* end_ptr = (char*)d->arena + d->nstack * sizeof(mjtNum);
-  char* result = end_ptr - (d->pstack + size + 1) * sizeof(mjtNum);
+  char* result = end_ptr - (d->pstack + size + mjREDZONE) * sizeof(mjtNum);
+  size_t new_pstack = d->pstack + size + 2*mjREDZONE;
+  #undef mjREDZONE
+
+  // new stack usage level
+  size_t usage;
 
 #ifdef ADDRESS_SANITIZER
   if ((uintptr_t)result % sizeof(mjtNum)) {
     mju_error("mj_stackAlloc fails to align to sizeof(mjtNum)");
   }
+
+  // actual stack usage (without red zone bytes) is stored in the red zone
+  if (d->pstack) {
+    size_t* prev_ptr = (size_t*)(end_ptr - d->pstack*sizeof(mjtNum));
+    ASAN_UNPOISON_MEMORY_REGION(prev_ptr, sizeof(size_t));
+    usage = *prev_ptr + size;
+    ASAN_POISON_MEMORY_REGION(prev_ptr, sizeof(size_t));
+  } else {
+    usage = size;
+  }
+
+  // store new stack usage in the red zone
+  size_t* cur_ptr = (size_t*)(end_ptr - new_pstack*sizeof(mjtNum));
+  ASAN_UNPOISON_MEMORY_REGION(cur_ptr, sizeof(size_t));
+  *cur_ptr = usage;
+  ASAN_POISON_MEMORY_REGION(cur_ptr, sizeof(size_t));
+
+  // unpoison the actual usable allocation
+  ASAN_UNPOISON_MEMORY_REGION(result, size*sizeof(mjtNum));
+#else
+  usage = d->pstack + size;
 #endif
 
-  // update max, return pointer to buffer
-  d->pstack += size;
-  d->maxuse_stack = mjMAX(d->maxuse_stack, d->pstack);
-  d->maxuse_arena = mjMAX(d->maxuse_arena, d->pstack*sizeof(mjtNum) + d->parena);
+  // update pstack and max usage statistics
+  d->pstack = new_pstack;
+  d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
+  d->maxuse_arena = mjMAX(d->maxuse_arena, usage*sizeof(mjtNum) + d->parena);
+
   return (mjtNum*)result;
 }
 
@@ -1232,6 +1280,16 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
 
   // clear arena pointers
   d->parena = 0;
+
+  // poison the entire arena+stack memory region when built with asan
+#ifdef ADDRESS_SANITIZER
+  ASAN_POISON_MEMORY_REGION(d->arena, d->nstack * sizeof(mjtNum));
+#endif
+
+#ifdef MEMORY_SANITIZER
+  __msan_allocated_memory(d->arena, d->nstack * sizeof(mjtNum));
+#endif
+
 #define X(type, name, nr, nc) d->name = NULL;
   MJDATA_ARENA_POINTERS
 #undef X

src/engine/engine_io.h

@@ -100,7 +100,7 @@ MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_va
 MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 
 // mjData arena allocate
-void* mj_arenaAlloc(mjData* d, int bytes, int alignment);
+MJAPI void* mj_arenaAlloc(mjData* d, int bytes, int alignment);
 
 // mjData stack allocate for array of mjtNums
 MJAPI mjtNum* mj_stackAlloc(mjData* d, int size);

src/engine/engine_macro.h

@@ -17,22 +17,11 @@
 
 #include <stdint.h>
 
+#include <mujoco/mjmacro.h>
 #include "engine/engine_callback.h"  // IWYU pragma: export
 
 //-------------------------------- utility macros --------------------------------------------------
 
-// mark and free stack
-#define mjMARKSTACK int _mark = d->pstack;
-#define mjFREESTACK d->pstack = _mark;
-
-// check bitflag
-#define mjDISABLED(x) (m->opt.disableflags & (x))
-#define mjENABLED(x)  (m->opt.enableflags & (x))
-
-// max and min macros
-#define mjMAX(a, b)    (((a) > (b)) ? (a) : (b))
-#define mjMIN(a, b)    (((a) < (b)) ? (a) : (b))
-
 // thread local macro
 #ifdef _MSC_VER
 #define mjTHREADLOCAL __declspec(thread)
@@ -49,21 +38,7 @@
 #define TM_START1 mjtNum _tm1 = (mjcb_time ? mjcb_time() : 0);
 #define TM_END1(i) {d->timer[i].duration += ((mjcb_time ? mjcb_time() : 0) - _tm1); d->timer[i].number++;}
 
-//-------------------------- sanitizer macros ------------------------------------------------------
-
-#ifdef ADDRESS_SANITIZER
-  #include <sanitizer/asan_interface.h>
-#elif defined(_MSC_VER)
-  #define ASAN_POISON_MEMORY_REGION(addr, size)
-  #define ASAN_UNPOISON_MEMORY_REGION(addr, size)
-#else
-  #define ASAN_POISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
-  #define ASAN_UNPOISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
-#endif
-
-#ifdef MEMORY_SANITIZER
-  #include <sanitizer/msan_interface.h>
-#endif
+//-------------------------- compiler builtin ------------------------------------------------------
 
 #ifndef __has_builtin
 #define __has_builtin(x) 0

src/engine/engine_passive.c

@@ -18,11 +18,11 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_callback.h"
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_crossplatform.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"

src/engine/engine_print.c

@@ -20,11 +20,11 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
@@ -770,7 +770,7 @@ 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) {         \
     const char* format = _Generic(                                            \
         d->name,                                                              \
         int : INT_FORMAT,                                                     \

src/engine/engine_ray.c

@@ -18,10 +18,10 @@
 #include <stddef.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_sensor.c

@@ -17,13 +17,13 @@
 #include <stddef.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
 #include "engine/engine_callback.h"
 #include "engine/engine_core_smooth.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_ray.h"
 #include "engine/engine_support.h"

src/engine/engine_setconst.c

@@ -17,12 +17,12 @@
 #include <stdio.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_core_smooth.h"
 #include "engine/engine_forward.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"

src/engine/engine_solver.c

@@ -18,11 +18,11 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_core_smooth.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"

src/engine/engine_support.c

@@ -17,12 +17,12 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_array_safety.h"
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_util_blas.c

@@ -16,6 +16,7 @@
 
 #include <string.h>
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjtnum.h>
 
 #ifdef mjUSEPLATFORMSIMD

src/engine/engine_util_errmem.c

@@ -27,6 +27,7 @@
 
 #include "engine/engine_array_safety.h"
 #include "engine/engine_macro.h"
+
 //------------------------- cross-platform aligned malloc/free -------------------------------------
 
 static inline void* mju_alignedMalloc(size_t size, size_t align) {

src/engine/engine_util_errmem.h

@@ -18,6 +18,7 @@
 #include <stddef.h>
 
 #include <mujoco/mjexport.h>
+#include <mujoco/mjmacro.h>
 
 #ifdef __cplusplus
 extern "C" {

src/engine/engine_util_misc.c

@@ -20,6 +20,7 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_array_safety.h"
 #include "engine/engine_macro.h"

src/engine/engine_util_solve.c

@@ -19,8 +19,8 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_util_sparse.c

@@ -18,9 +18,9 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjtnum.h>
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 
 

src/engine/engine_vis_init.c

@@ -17,10 +17,10 @@
 #include <math.h>
 #include <string.h>
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_array_safety.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
 

src/engine/engine_vis_interact.c

@@ -19,11 +19,11 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_core_smooth.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_ray.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"

src/engine/engine_vis_state.c

@@ -22,7 +22,6 @@
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_core_constraint.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_errmem.h"
@@ -30,6 +29,10 @@
 #include "engine/engine_vis_interact.h"
 #include "engine/engine_vis_visualize.h"
 
+#ifdef MEMORY_SANITIZER
+  #include <sanitizer/msan_interface.h>
+#endif
+
 // this source file needs to treat XMJV differently from other X macros
 #undef XMJV
 

src/engine/engine_vis_visualize.c

@@ -18,10 +18,10 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_array_safety.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"