Add mj_markStack and mj_freeStack as public API functions.

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

src/engine/engine_io.c

@@ -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

@@ -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

@@ -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,                                                     \