From 0c97dcf1fd56dc97dcb89fdbe10aa8d73e1578e0 Mon Sep 17 00:00:00 2001
From: Kyle Bayes <kylebayes@google.com>
Date: Tue, 22 Oct 2024 12:10:29 -0700
Subject: [PATCH] Avoid overflowing the stack when allocating memory for
 polytope in EPA.

PiperOrigin-RevId: 688641864
Change-Id: I3eef7e821ffbfa8d97a946ecf86d098cb4f0c233
---
 src/engine/engine_collision_gjk.c | 21 ++++++++++++++++-----
 src/engine/engine_io.c            | 14 ++++++++++++++
 src/engine/engine_io.h            |  7 +++++--
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/src/engine/engine_collision_gjk.c b/src/engine/engine_collision_gjk.c
index 44e5b5da..f513dc19 100644
--- a/src/engine/engine_collision_gjk.c
+++ b/src/engine/engine_collision_gjk.c
@@ -1327,16 +1327,27 @@ mjtNum mjc_ccd(const mjCCDConfig* config, mjCCDStatus* status, mjCCDObj* obj1, m
     Polytope pt;
     pt.nfaces = pt.nmap = pt.nverts = 0;
 
-    // allocate memory for faces
-    pt.maxfaces = (6*N > 1000) ? 6*N : 1000;  // use 1000 faces as lower bound
-    pt.faces = mj_stackAllocByte(d, sizeof(Face) * pt.maxfaces, _Alignof(Face));
-    pt.map = mj_stackAllocByte(d, sizeof(Face*) * pt.maxfaces, _Alignof(Face*));
-
     // allocate memory for vertices
     pt.verts  = mj_stackAllocNum(d, 3*(5 + N));
     pt.verts1 = mj_stackAllocNum(d, 3*(5 + N));
     pt.verts2 = mj_stackAllocNum(d, 3*(5 + N));
 
+    // allocate memory for faces
+    pt.maxfaces = (6*N > 1000) ? 6*N : 1000;  // use 1000 faces as lower bound
+    size_t size1 = sizeof(Face) * pt.maxfaces;
+    size_t size2 = sizeof(Face*) * pt.maxfaces;
+
+    // since a generous upper bound is used, we need to rescale stack use if not enough
+    // memory is available
+    size_t max_size = mj_stackBytesAvailable(d) - 12*(N * sizeof(int));
+    if (size1 + size2 > max_size) {
+      pt.maxfaces = max_size / (sizeof(Face) + sizeof(Face*));
+      size1 = sizeof(Face) * pt.maxfaces;
+      size2 = sizeof(Face*) * pt.maxfaces;
+    }
+    pt.faces = mj_stackAllocByte(d, size1, _Alignof(Face));
+    pt.map = mj_stackAllocByte(d, size2, _Alignof(Face*));
+
     int ret;
     if (status->nsimplex == 2) {
       ret = polytope2(&pt, status, obj1, obj2);
diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c
index c66eeee6..88ca9bd2 100644
--- a/src/engine/engine_io.c
+++ b/src/engine/engine_io.c
@@ -1715,6 +1715,20 @@ void mj__freeStack(mjData* d)
 
 
 
+// returns the number of bytes available on the stack
+size_t mj_stackBytesAvailable(mjData* d) {
+  if (!d->threadpool) {
+    mjStackInfo stack_info = get_stack_info_from_data(d);
+    return stack_info.top - stack_info.limit;
+  } else {
+    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+    mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
+    return stack_info->top - stack_info->limit;
+  }
+}
+
+
+
 // allocate bytes on the stack
 void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment) {
   return stackalloc(d, bytes, alignment);
diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h
index 44b0764b..4b34784f 100644
--- a/src/engine/engine_io.h
+++ b/src/engine/engine_io.h
@@ -75,7 +75,7 @@ MJAPI void mj_saveModel(const mjModel* m, const char* filename, void* buffer, in
 // load binary MJB
 mjModel* mj_loadModelBuffer(const void* buffer, int buffer_sz);
 
-// de-allocate model
+// deallocate model
 MJAPI void mj_deleteModel(mjModel* m);
 
 // size of buffer needed to hold model
@@ -128,6 +128,9 @@ void mj__freeStack(mjData* d) __attribute__((noinline));
 
 #endif  // ADDRESS_SANITIZER
 
+// returns the number of bytes available on the stack
+MJAPI size_t mj_stackBytesAvailable(mjData* d);
+
 // mjData stack allocate
 MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
 
@@ -137,7 +140,7 @@ MJAPI mjtNum* mj_stackAllocNum(mjData* d, size_t size);
 // mjData stack allocate for array of ints
 MJAPI int* mj_stackAllocInt(mjData* d, size_t size);
 
-// de-allocate data
+// deallocate data
 MJAPI void mj_deleteData(mjData* d);
 
 // clear arena pointers in mjData