From abf6d41b7f8cef59bc8d1f22321c412486551240 Mon Sep 17 00:00:00 2001 From: Saran Tunyasuvunakool Date: Tue, 20 Feb 2024 09:25:26 -0800 Subject: [PATCH] Fix incorrect data types in Python bindings of certain arrays. The mismatch between the `py::array_t` template argument and the type of the MuJoCo struct member being wrapped caused the `py::array_t` to be constructed via an unintended constructor overload, which in turn triggered `use-after-poison` under asan. This bug affected the `geom`, `flex`, `elem`, and `vert` arrays in `mjContact`, and all array members in `mjrContext`. PiperOrigin-RevId: 608632860 Change-Id: Ic227f691041b004a0fcf0cb6fa3d62c44861d9aa --- doc/changelog.rst | 6 ++++ python/mujoco/bindings_test.py | 8 +++++ python/mujoco/render.cc | 7 +++- python/mujoco/structs.h | 64 +++++++++++++++++++++++++++------- 4 files changed, 72 insertions(+), 13 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 2c6f6aaf..861cfc61 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -27,6 +27,12 @@ MJX 5. Fixed a bug in ``mjx.solve`` that was causing slow convergence when using ``mjSOL_NEWTON`` in :ref:`mjtSolver`. 6. Added support for :ref:`mjOption.impratio` to ``mjx.Model``. +Python bindings +^^^^^^^^^^^^^^^ +7. Fixed incorrect data types in the bindings for the ``geom``, ``vert``, ``elem``, and ``flex`` array members + of the ``mjContact`` struct, and all array members of the ``mjrContext`` struct. + + Version 3.1.2 (February 05, 2024) ----------------------------------- diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index 407557f0..d293448e 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -529,6 +529,7 @@ class MuJoCoBindingsTest(parameterized.TestCase): expected_ncon = 4 mujoco.mj_forward(self.model, self.data) self.assertLen(self.data.contact, expected_ncon) + np.testing.assert_array_equal(self.data.contact.geom, [[0, 1]] * 4) expected_pos = [] for contact in self.data.contact: @@ -551,6 +552,13 @@ class MuJoCoBindingsTest(parameterized.TestCase): self.assertLen(expected_H, expected_ncon) np.testing.assert_array_equal(self.data.contact.H, expected_H) + expected_geom = [] + for i, contact in enumerate(self.data.contact): + expected_geom.append([i, i + 1]) + contact.geom = expected_geom[-1] + self.assertLen(expected_geom, expected_ncon) + np.testing.assert_array_equal(self.data.contact.geom, expected_geom) + def test_realloc_con_efc(self): self.assertEmpty(self.data.contact) diff --git a/python/mujoco/render.cc b/python/mujoco/render.cc index b2922efc..8be1fc69 100644 --- a/python/mujoco/render.cc +++ b/python/mujoco/render.cc @@ -14,6 +14,8 @@ #include #include +#include +#include #include #include @@ -40,7 +42,10 @@ class MjWrapper : public WrapperBase { void Free(); -#define X(var) py_array_or_tuple_t var + #define X(var) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var X(fogRGBA); X(auxWidth); X(auxHeight); diff --git a/python/mujoco/structs.h b/python/mujoco/structs.h index 5aba80f5..f27e17ff 100644 --- a/python/mujoco/structs.h +++ b/python/mujoco/structs.h @@ -17,12 +17,17 @@ #include #include +#include #include #include #include #include +#include +#include #include #include +#include +#include #include #include @@ -36,6 +41,7 @@ namespace mujoco::python { namespace _impl { + template class WrapperBase { public: @@ -44,7 +50,7 @@ class WrapperBase { T* get() { return ptr_; } const T* get() const { return ptr_; } - const pybind11::handle owner() const { return owner_; } + pybind11::handle owner() const { return owner_; } protected: static void DefaultCapsuleDestructor(PyObject* pyobj) { @@ -167,7 +173,10 @@ class MjWrapper : public WrapperBase { MjWrapper(raw::MjOption* ptr, pybind11::handle owner); ~MjWrapper() = default; - #define X(var, dim) py_array_or_tuple_t var; + #define X(var, dim) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var; MJOPTION_VECTORS #undef X }; @@ -188,7 +197,10 @@ class MjWrapper MjWrapper(raw::MjVisualHeadlight* ptr, pybind11::handle owner); ~MjWrapper() = default; - #define X(var) py_array_or_tuple_t var + #define X(var) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var X(ambient); X(diffuse); X(specular); @@ -209,7 +221,10 @@ class MjWrapper : public WrapperBase { MjWrapper(raw::MjVisualRgba* ptr, pybind11::handle owner); ~MjWrapper() = default; - #define X(var) py_array_or_tuple_t var + #define X(var) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var X(fog); X(haze); X(force); @@ -271,7 +286,10 @@ class MjWrapper : public WrapperBase { MjWrapper(raw::MjStatistic* ptr, pybind11::handle owner); ~MjWrapper() = default; - #define X(var) py_array_or_tuple_t var + #define X(var) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var X(center); #undef X }; @@ -504,7 +522,10 @@ class MjWrapper : public WrapperBase { MjWrapper(raw::MjContact* ptr, pybind11::handle owner); ~MjWrapper() = default; - #define X(var) py_array_or_tuple_t var + #define X(var) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var X(pos); X(frame); X(friction); @@ -628,7 +649,10 @@ class MjWrapper : public WrapperBase { MjWrapper(MjWrapper&&) = default; ~MjWrapper() = default; - #define X(var) py_array_or_tuple_t var + #define X(var) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var X(refpos); X(refquat); X(refselpos); @@ -650,7 +674,10 @@ class MjWrapper : public WrapperBase { MjWrapper(MjWrapper&&) = default; ~MjWrapper() = default; - #define X(var) py_array_or_tuple_t var + #define X(var) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var X(lookat); #undef X }; @@ -671,7 +698,10 @@ class MjWrapper : public WrapperBase { explicit MjWrapper(raw::MjvGLCamera&& other); ~MjWrapper() = default; - #define X(var) py_array_or_tuple_t var + #define X(var) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var X(pos); X(forward); X(up); @@ -693,7 +723,10 @@ class MjWrapper : public WrapperBase { MjWrapper(raw::MjvGeom* ptr, pybind11::handle owner); ~MjWrapper() = default; - #define X(var) py_array_or_tuple_t var + #define X(var) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var X(texrepeat); X(size); X(pos); @@ -717,7 +750,10 @@ class MjWrapper : public WrapperBase { MjWrapper(raw::MjvLight* ptr, pybind11::handle owner); ~MjWrapper() = default; - #define X(var) py_array_or_tuple_t var + #define X(var) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var X(pos); X(dir); X(attenuation); @@ -741,7 +777,10 @@ class MjWrapper : public WrapperBase { MjWrapper(MjWrapper&&) = default; ~MjWrapper() = default; - #define X(var) py_array_or_tuple_t var + #define X(var) \ + py_array_or_tuple_t< \ + std::remove_all_extents_t> \ + var X(geomgroup); X(sitegroup); X(jointgroup); @@ -862,6 +901,7 @@ class ScopedMsanDisabler { void* shadow_; }; #endif + } // namespace _impl template