From 2a26bf1ba53c168614c66164b39f8797705e197c Mon Sep 17 00:00:00 2001
From: Nimrod Gileadi <nimrod@google.com>
Date: Mon, 20 Jun 2022 08:57:14 -0700
Subject: [PATCH] Disable MSAN poisoning when printing using engine_print.

When using MSAN, mjData->buffer is marked as poisoned, despite being initialized with zeroes.
This is useful for catching computations that accidentally use uninitialized values.

However, when using engine_print, it's OK to assume the buffer is initialized to zeros.

PiperOrigin-RevId: 456085776
Change-Id: I9d85e37fa82958eb54ac413140543a6581aa46ad
---
 src/engine/engine_print.c        | 20 ++++++++
 test/engine/CMakeLists.txt       |  3 ++
 test/engine/engine_print_test.cc | 87 ++++++++++++++++++++++++++++++++
 3 files changed, 110 insertions(+)
 create mode 100644 test/engine/engine_print_test.cc

diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c
index 92fb1bc2..accaf7b2 100644
--- a/src/engine/engine_print.c
+++ b/src/engine/engine_print.c
@@ -30,6 +30,10 @@
 #include "engine/engine_util_misc.h"
 #include "engine/engine_util_sparse.h"
 
+#ifdef MEMORY_SANITIZER
+  #include <sanitizer/msan_interface.h>
+#endif
+
 #define FLOAT_FORMAT "% -9.2g"
 #define FLOAT_FORMAT_MAX_LEN 20
 #define INT_FORMAT " %d"
@@ -717,6 +721,16 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
   // allocate full inertia
   M = mj_stackAlloc(d, m->nv*m->nv);
 
+#ifdef MEMORY_SANITIZER
+  // If memory sanitizer is active, d->buffer will be marked as poisoned, even
+  // though it's really initialized to 0. This catches unintentionally
+  // using uninitialized values, but in engine_print it's OK to output zeroes.
+
+  // save current poison status of buffer before marking unpoisoned
+  void* shadow = mju_malloc(d->nbuffer);
+  __msan_copy_shadow(shadow, d->buffer, d->nbuffer);
+  __msan_unpoison(d->buffer, d->nbuffer);
+#endif
   // ---------------------------------- print mjData fields
 
   fprintf(fp, "SIZES\n");
@@ -966,6 +980,12 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
   printArray("CFRC_INT", m->nbody, 6, d->cfrc_int, fp, float_format);
   printArray("CFRC_EXT", m->nbody, 6, d->cfrc_ext, fp, float_format);
 
+#ifdef MEMORY_SANITIZER
+  // restore poisoned status
+  __msan_copy_shadow(d->buffer, shadow, d->nbuffer);
+  mju_free(shadow);
+#endif
+
   if (filename) {
     fclose(fp);
   }
diff --git a/test/engine/CMakeLists.txt b/test/engine/CMakeLists.txt
index ce423dab..d8248b73 100644
--- a/test/engine/CMakeLists.txt
+++ b/test/engine/CMakeLists.txt
@@ -35,6 +35,9 @@ target_link_libraries(
   absl::str_format
 )
 
+mujoco_test(engine_print_test)
+target_link_libraries(engine_print_test fixture gmock)
+
 mujoco_test(engine_ray_test)
 target_link_libraries(engine_ray_test fixture gmock)
 
diff --git a/test/engine/engine_print_test.cc b/test/engine/engine_print_test.cc
new file mode 100644
index 00000000..2fd4dbcc
--- /dev/null
+++ b/test/engine/engine_print_test.cc
@@ -0,0 +1,87 @@
+// Copyright 2022 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/engine/engine_print.h"
+
+#include <array>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mujoco.h>
+#include "test/fixture.h"
+
+#ifdef MEMORY_SANITIZER
+  #include <sanitizer/msan_interface.h>
+#endif
+
+namespace mujoco {
+namespace {
+
+#ifdef MEMORY_SANITIZER
+using ::testing::Eq;
+using ::testing::Not;
+#endif
+using ::testing::NotNull;
+
+using EnginePrintTest = MujocoTest;
+
+
+constexpr const char* NullFile() {
+#ifdef _WIN32
+  return "NUL";
+#else
+  return "/dev/null";
+#endif
+}
+
+TEST_F(EnginePrintTest, PrintDataWorksWithMsan) {
+  constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body>
+        <joint/>
+        <geom size="1"/>
+      </body>
+    </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());
+
+  // Mark qacc_smooth[0] as initialized.
+  data->qacc_smooth[0] = 0;
+
+  // This will read "uninitialized" values from mjData, but shouldn't fail.
+  mj_printData(model, data, NullFile());
+
+  // After mj_printData, poisoned status should be restored correctly, so
+  // qacc_smooth[0] should be marked initialzed.
+  EXPECT_EQ(data->qacc_smooth[0], 0);
+
+#ifdef MEMORY_SANITIZER
+  EXPECT_THAT(__msan_test_shadow(data->buffer, data->nbuffer), Not(Eq(-1))) <<
+      "Expecting some of data->buffer to be marked uninitialized";
+#endif
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+}  // namespace
+}  // namespace mujoco