From a07e172770ed104508aeb85e74112d9b52ec9f36 Mon Sep 17 00:00:00 2001
From: Saran Tunyasuvunakool <stunya@google.com>
Date: Fri, 8 Sep 2023 03:28:41 -0700
Subject: [PATCH] Make MuJoCo work with `-fsanitize=alignment`.

PiperOrigin-RevId: 563702045
Change-Id: Ieca69595b7e2fc1e35ca69754d7cc7971825f48b
---
 src/engine/engine_io.c        | 15 +++++++++++----
 src/user/user_mesh.cc         | 27 ++++++++++++++++++---------
 test/engine/engine_io_test.cc | 24 ++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c
index 4d1b3d22..a9eb5444 100644
--- a/src/engine/engine_io.c
+++ b/src/engine/engine_io.c
@@ -1289,16 +1289,23 @@ static inline void* stackalloc(mjData* d, size_t size, size_t alignment) {
 #ifdef ADDRESS_SANITIZER
   // actual stack usage (without red zone bytes) is stored in the red zone
   if (d->pstack) {
-    size_t* prev_usage_ptr = (size_t*)(end_of_arena_ptr - d->pstack);
+    char* prev_pstack_ptr = (char*)(end_of_arena_ptr - d->pstack);
+    size_t prev_misalign = (uintptr_t)prev_pstack_ptr % _Alignof(size_t);
+    size_t* prev_usage_ptr =
+        (size_t*)(prev_pstack_ptr +
+                  (prev_misalign ? _Alignof(size_t) - prev_misalign : 0));
     ASAN_UNPOISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
     usage = current_alloc_usage + *prev_usage_ptr;
     ASAN_POISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
   }
 
   // store new stack usage in the red zone
-  ASAN_UNPOISON_MEMORY_REGION((void*)new_pstack_ptr, sizeof(size_t));
-  *(size_t*)new_pstack_ptr = usage;
-  ASAN_POISON_MEMORY_REGION((void*)new_pstack_ptr, sizeof(size_t));
+  size_t misalign = new_pstack_ptr % _Alignof(size_t);
+  size_t* usage_ptr =
+      (size_t*)(new_pstack_ptr + (misalign ? _Alignof(size_t) - misalign : 0));
+  ASAN_UNPOISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
+  *usage_ptr = usage;
+  ASAN_POISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
 
   // unpoison the actual usable allocation
   ASAN_UNPOISON_MEMORY_REGION((void*)start_ptr, size);
diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc
index 57ba191f..606fc468 100644
--- a/src/user/user_mesh.cc
+++ b/src/user/user_mesh.cc
@@ -114,6 +114,12 @@ static T* VecToArray(std::vector<T>& vector,  bool clear = true){
   }
 }
 
+// Read data of type T from a potentially unaligned buffer pointer.
+template <typename T>
+static void ReadFromBuffer(T* dst, const char* src) {
+  std::memcpy(dst, src, sizeof(T));
+}
+
 //------------------ class mjCMesh implementation --------------------------------------------------
 
 mjCMesh::mjCMesh(mjCModel* _model, mjCDef* _def) {
@@ -961,7 +967,7 @@ void mjCMesh::LoadSTL(mjResource* resource) {
   }
 
   // get number of triangles, check bounds
-  nface_ = *(unsigned int*)(buffer+80);
+  ReadFromBuffer(&nface_, buffer + 80);
   if (nface_<1 || nface_>200000) {
     throw mjCError(this,
                    "number of faces should be between 1 and 200000 in STL file '%s';"
@@ -985,8 +991,10 @@ void mjCMesh::LoadSTL(mjResource* resource) {
   // add vertices and faces, including repeated for now
   for (int i=0; i<nface_; i++) {
     for (int j=0; j<3; j++) {
-      // get pointer to vertex coordiates
-      float* v = (float*)(stl+50*i+12*(j+1));
+      // read vertex coordinates
+      float v[3];
+      ReadFromBuffer(&v, stl+50*i+12*(j+1));
+
       for (int k=0; k < 3; k++) {
         if (std::isnan(v[k]) || std::isinf(v[k])) {
           throw mjCError(this, "STL file '%s' contains invalid vertices.",
@@ -1039,10 +1047,10 @@ void mjCMesh::LoadMSH(mjResource* resource) {
   }
 
   // get sizes from header
-  nvert_ = ((int*)buffer)[0];
-  nnormal_ = ((int*)buffer)[1];
-  ntexcoord_ = ((int*)buffer)[2];
-  nface_ = ((int*)buffer)[3];
+  ReadFromBuffer(&nvert_, buffer);
+  ReadFromBuffer(&nnormal_, buffer + sizeof(int));
+  ReadFromBuffer(&ntexcoord_, buffer + 2*sizeof(int));
+  ReadFromBuffer(&nface_, buffer + 3*sizeof(int));
 
   // check sizes
   if (nvert_<4 || nface_<0 || nnormal_<0 || ntexcoord_<0 ||
@@ -1058,7 +1066,8 @@ void mjCMesh::LoadMSH(mjResource* resource) {
   }
 
   // allocate and copy
-  float* fdata = (float*)(((int*)buffer) + 4);
+  using UnalignedFloat = char[sizeof(float)];
+  auto fdata = reinterpret_cast<UnalignedFloat*>(buffer + 4*sizeof(int));
   if (nvert_) {
     vert_ = (float*) mju_malloc(3*nvert_*sizeof(float));
     memcpy(vert_, fdata, 3*nvert_*sizeof(float));
@@ -1085,7 +1094,7 @@ void mjCMesh::LoadMSH(mjResource* resource) {
     memcpy(facetexcoord_, fdata, 3*nface_*sizeof(int));
   }
 
-  // rearange face data if left-handed scaling
+  // rearrange face data if left-handed scaling
   if (nface_ && !righthand) {
     for (int i=0; i<nface_; i++) {
       int tmp = face_[3*i+1];
diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc
index c73b28e3..b9f723cb 100644
--- a/test/engine/engine_io_test.cc
+++ b/test/engine/engine_io_test.cc
@@ -813,6 +813,30 @@ TEST_F(EngineIoTest, CanDetectStackFrameLeakage) {
   mj_deleteData(data);
   mj_deleteModel(model);
 }
+
+TEST_F(EngineIoTest, RedZoneAlignmentTest) {
+  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());
+
+  mj_markStack(data);
+  mj_stackAlloc(data, 1, 1);
+  mj_stackAlloc(data, 1, 1);
+  mj_freeStack(data);
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
 #endif
 
 }  // namespace