From 7076da690919559b0e1fe73dca9f7fee94fb5e24 Mon Sep 17 00:00:00 2001
From: Nimrod Gileadi <nimrod@google.com>
Date: Mon, 12 Dec 2022 17:48:23 -0800
Subject: [PATCH] Correctly validate the size of user sensors in
 mj_validateReference.

PiperOrigin-RevId: 494877162
Change-Id: I0a070d7813efea6216d15270f62b6009c55e1a54
---
 src/engine/engine_io.c        |  6 +++---
 test/engine/engine_io_test.cc | 19 +++++++++++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c
index 8d1ce4dd..cb7d61c6 100644
--- a/src/engine/engine_io.c
+++ b/src/engine/engine_io.c
@@ -1209,7 +1209,7 @@ void mj_deleteData(mjData* d) {
 const int nPOS[4] = {7, 4, 1, 1};
 const int nVEL[4] = {6, 3, 1, 1};
 
-static int sensorSize(mjtSensor sensor_type, int nuser_sensor) {
+static int sensorSize(mjtSensor sensor_type, int sensor_dim) {
   switch (sensor_type) {
   case mjSENS_TOUCH:
   case mjSENS_RANGEFINDER:
@@ -1254,7 +1254,7 @@ static int sensorSize(mjtSensor sensor_type, int nuser_sensor) {
     return 4;
 
   case mjSENS_USER:
-    return nuser_sensor;
+    return sensor_dim;
 
   case mjSENS_PLUGIN:
     return -1;
@@ -1593,7 +1593,7 @@ const char* mj_validateReferences(const mjModel* m) {
       }
       sensor_size = plugin->nsensordata(m, m->sensor_plugin[i], i);
     } else {
-      sensor_size = sensorSize(sensor_type, m->nuser_sensor);
+      sensor_size = sensorSize(sensor_type, m->sensor_dim[i]);
     }
     if (sensor_size < 0) {
         return "Invalid model: Bad sensor_type.";
diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc
index d555ef0f..e2a6c8da 100644
--- a/test/engine/engine_io_test.cc
+++ b/test/engine/engine_io_test.cc
@@ -434,6 +434,25 @@ TEST_F(ValidateReferencesTest, SensorsAddress) {
   }
 }
 
+TEST_F(ValidateReferencesTest, SensorsAddressUser) {
+  static const char xml[] = R"(
+  <mujoco>
+    <sensor>
+      <user dim="3" user="1 2 3 4 5" />
+      <user dim="2" />
+      <user dim="1" />
+    </sensor>
+  </mujoco>
+  )";
+
+  std::array<char, 1024> error;
+  mjModel* model = LoadModelFromString(xml, error.data(), error.size());
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+
+  EXPECT_THAT(mj_validateReferences(model), IsNull());
+  mj_deleteModel(model);
+}
+
 TEST_F(ValidateReferencesTest, SensorsObj) {
   static const char xml[] = R"(
   <mujoco>