From 4268d81b55b655d05de58ae0c18f10d6c27319f5 Mon Sep 17 00:00:00 2001
From: Nimrod Gileadi <nimrod@google.com>
Date: Thu, 4 Aug 2022 12:36:07 -0700
Subject: [PATCH] Null out mjData.buffer and mjData.stack before possibly
 calling mj_deleteData.

While creating an mjData, any error triggers mj_deleteData. If buffer or stack is not NULL at that point, an uninitialized pointer will be freed.

PiperOrigin-RevId: 465378329
Change-Id: I9be0eef0648e05e3e5f1361da346a5046485cde7
---
 src/engine/engine_io.c        |  5 ++++-
 test/engine/engine_io_test.cc | 24 ++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c
index 010ea6a7..255adfa2 100644
--- a/src/engine/engine_io.c
+++ b/src/engine/engine_io.c
@@ -838,6 +838,7 @@ static mjData* _makeData(const mjModel* m) {
 
   // compute buffer size
   d->nbuffer = 0;
+  d->buffer = d->stack = NULL;
 #define X(type, name, nr, nc)                                                \
   if (!safeAddToBufferSize(&offset, &d->nbuffer, sizeof(type), m->nr, nc)) { \
     mju_warning("Invalid data: " #name " too large.");                       \
@@ -871,7 +872,9 @@ static mjData* _makeData(const mjModel* m) {
 
 mjData* mj_makeData(const mjModel* m) {
   mjData* d = _makeData(m);
-  mj_resetData(m, d);
+  if (d) {
+    mj_resetData(m, d);
+  }
   return d;
 }
 
diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc
index 0dbc71ae..a4d2174d 100644
--- a/test/engine/engine_io_test.cc
+++ b/test/engine/engine_io_test.cc
@@ -17,6 +17,7 @@
 #include "src/engine/engine_io.h"
 
 #include <array>
+#include <climits>
 #include <cstring>
 #include <string>
 
@@ -24,6 +25,7 @@
 #include <gtest/gtest.h>
 #include <absl/strings/str_format.h>
 #include <mujoco/mjxmacro.h>
+#include "src/engine/engine_util_errmem.h"
 #include "test/fixture.h"
 
 namespace mujoco {
@@ -177,6 +179,28 @@ TEST_F(EngineIoTest, CopyDataWithPartialModel) {
   mj_deleteData(copy);
 }
 
+TEST_F(EngineIoTest, MakeDataReturnsNullOnFailure) {
+  constexpr char xml[] = "<mujoco/>";
+
+  std::array<char, 1024> error;
+  mjModel* model = LoadModelFromString(xml, error.data(), error.size());
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+
+  // trigger overflow intentionally
+  model->nbody = INT_MAX;
+  static bool warning = false;
+  warning = false;
+  mju_user_warning = [](const char* error) {
+    warning = true;
+  };
+  mjData* data = mj_makeData(model);
+  EXPECT_THAT(data, IsNull());
+  EXPECT_TRUE(warning) << "Expecting warning to be triggered.";
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 using ValidateReferencesTest = MujocoTest;
 
 TEST_F(ValidateReferencesTest, BodyReferences) {