Merge pull request #3120 from Ashutosh0x:fix/safe-overflow-msvc

PiperOrigin-RevId: 923296001
Change-Id: I037dcc55c91d6eb68b3e63b34a0e059c796b5116

src/engine/engine_io.c

@@ -185,9 +185,32 @@ static mjtSize safeAddToBufferSize(intptr_t* offset, mjtSize* nbuffer,
   if (__builtin_add_overflow(*nbuffer, to_add, nbuffer)) return 0;
   if (__builtin_add_overflow(*offset, to_add, offset)) return 0;
 #else
-  // TODO: offer a safe implementation for MSVC or other compilers that don't have the builtins
-  *nbuffer += SKIP(*offset) + type_size*nr*nc;
-  *offset += SKIP(*offset) + type_size*nr*nc;
+  // safe overflow checks for MSVC and other compilers without __builtin_*_overflow
+  {
+    size_t product;
+    size_t to_add;
+    size_t skip = SKIP(*offset);
+
+    // nc * nr
+    if (nr > 0 && (size_t)nc > SIZE_MAX / (size_t)nr) return 0;
+    product = (size_t)nc * (size_t)nr;
+
+    // product * type_size
+    if (type_size > 0 && product > SIZE_MAX / type_size) return 0;
+    product *= type_size;
+
+    // product + SKIP(*offset)
+    if (product > SIZE_MAX - skip) return 0;
+    to_add = product + skip;
+
+    // *nbuffer + to_add
+    if ((size_t)*nbuffer > SIZE_MAX - to_add) return 0;
+    *nbuffer += to_add;
+
+    // *offset + to_add
+    if (*offset > 0 && to_add > (size_t)(INTPTR_MAX - *offset)) return 0;
+    *offset += to_add;
+  }
 #endif
 
   return 1;

test/engine/engine_io_test.cc

@@ -921,5 +921,67 @@ TEST_F(EngineIoTest, RedZoneAlignmentTest) {
 }
 #endif
 
+// Regression test: crafted binary model with overflow-inducing sizes must be
+// safely rejected (not cause heap overflow). This exercises the overflow checks
+// in safeAddToBufferSize on all compilers including MSVC.
+TEST_F(EngineIoTest, LoadModelBufferRejectsOverflowingSizes) {
+  // construct a minimal valid-looking .mjb header
+  int header[5];
+  header[0] = 20;              // ID
+  header[1] = sizeof(mjtNum);  // floating point size
+  // We need the correct nsize and nptr from the current build.
+  // Rather than hardcoding, we create and save a trivial model, then mutate
+  // the sizes to trigger overflow.
+
+  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();
+
+  // save model to a buffer
+  int bufsize = mj_sizeModel(model);
+  ASSERT_GT(bufsize, 0);
+  std::vector<char> buffer(bufsize);
+  mj_saveModel(model, nullptr, buffer.data(), bufsize);
+  mj_deleteModel(model);
+
+  // locate the size fields in the buffer (after 5-int header)
+  const int header_bytes = 5 * sizeof(int);
+  ASSERT_GT(bufsize, header_bytes + 77 * (int)sizeof(mjtSize));
+
+  // mutate a size field to an extremely large value that would overflow
+  // when multiplied by sizeof(type). ntexdata is a good candidate since it
+  // is a byte count field and gets multiplied by sizeof(mjtByte)==1, but other
+  // fields multiply by sizeof(int) or sizeof(mjtNum), making overflow easier.
+  // Use memcpy to avoid undefined behavior from unaligned access.
+  const int size_offset = header_bytes + 49 * sizeof(mjtSize);
+
+  // set to overflow-inducing value
+  mjtSize overflow_val = static_cast<mjtSize>(SIZE_MAX / 2);
+  std::memcpy(buffer.data() + size_offset, &overflow_val, sizeof(mjtSize));
+
+  // also need to update the nbuffer field (last size) to avoid the early
+  // nbuffer mismatch check — but the overflow should be caught earlier
+  // in safeAddToBufferSize/mj_makeModel before we reach that check.
+
+  // intercept mju_warning because the test framework translates it to ADD_FAILURE
+  static bool warning_triggered = false;
+  warning_triggered = false;
+  mju_user_warning = [](const char* msg) {
+    warning_triggered = true;
+  };
+
+  // attempt to load — should return NULL, not crash
+  mjModel* bad_model = mj_loadModelBuffer(buffer.data(), bufsize);
+  EXPECT_THAT(bad_model, IsNull())
+      << "Expected mj_loadModelBuffer to reject overflow-inducing sizes";
+  EXPECT_TRUE(warning_triggered) << "Expected a warning about invalid sizes";
+
+  // clean up if somehow it succeeded
+  if (bad_model) {
+    mj_deleteModel(bad_model);
+  }
+}
+
 }  // namespace
 }  // namespace mujoco