From 7f11a0df1240a58cb43ea8efb445e8fd05e1a1fc Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Mon, 17 Feb 2025 04:00:35 -0800
Subject: [PATCH] Add check for height fields size during compilation.

Fixes #2427.

PiperOrigin-RevId: 727805729
Change-Id: Idf3192e5fb238249faaeae6b4fb9de4c55755d96
---
 python/mujoco/specs_test.py | 24 ++++++++++++++++++++++++
 src/user/user_objects.cc    |  3 +++
 2 files changed, 27 insertions(+)

diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py
index 27e46327..09b6e413 100644
--- a/python/mujoco/specs_test.py
+++ b/python/mujoco/specs_test.py
@@ -1135,5 +1135,29 @@ class SpecsTest(absltest.TestCase):
     ):
       print(mj_model.bind(joints).invalid)
 
+  def test_incorrect_hfield_size(self):
+    nrow = 300
+    ncol = 400
+    hdata = np.random.uniform(size=(1, 1))
+    model_spec = mujoco.MjSpec()
+    model_spec.add_hfield(
+        name='hfield',
+        size=[1, 1, 1, 1e-3],
+        ncol=ncol,
+        nrow=nrow,
+        userdata=hdata.flatten(),
+    )
+    model_spec.worldbody.add_geom(
+        name='hfield',
+        type=mujoco.mjtGeom.mjGEOM_HFIELD,
+        pos=np.array([0, 0, 1]),
+        hfieldname='hfield',
+    )
+    with self.assertRaisesRegex(
+        ValueError, r"Error: elevation data length must match nrow\*ncol\n"
+        "Element name 'hfield', id 0",
+    ):
+      model_spec.compile()
+
 if __name__ == '__main__':
   absltest.main()
diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc
index 74062dff..0e8c0044 100644
--- a/src/user/user_objects.cc
+++ b/src/user/user_objects.cc
@@ -3696,6 +3696,9 @@ void mjCHField::Compile(const mjVFS* vfs) {
 
   // copy userdata into data
   if (!userdata_.empty()) {
+    if (nrow*ncol != userdata_.size()) {
+      throw mjCError(this, "elevation data length must match nrow*ncol");
+    }
     data.assign(nrow*ncol, 0);
     if (data.empty()) {
       throw mjCError(this, "could not allocate buffers in hfield");