Fix for recursive custom units (#12815)

src/backend/InvenTree/InvenTree/conversion.py

@@ -94,19 +94,8 @@ def get_unit_registry():
     return _unit_registry
 
 
-def reload_unit_registry():
+def new_base_registry() -> pint.UnitRegistry:
-    """Reload the unit registry from the database.
+    """Construct a new pint UnitRegistry, with InvenTree's default (non-custom) unit definitions."""
-
-    This function is called at startup, and whenever the database is updated.
-    """
-    import time
-
-    t_start = time.time()
-
-    global _unit_registry
-
-    _unit_registry = None
-
     reg = pint.UnitRegistry(autoconvert_offset_to_baseunit=True)
 
     # Aliases for temperature units
@@ -124,6 +113,24 @@ def reload_unit_registry():
     reg.define('hundred = 100')
     reg.define('thousand = 1000')
 
+    return reg
+
+
+def reload_unit_registry():
+    """Reload the unit registry from the database.
+
+    This function is called at startup, and whenever the database is updated.
+    """
+    import time
+
+    t_start = time.time()
+
+    global _unit_registry
+
+    _unit_registry = None
+
+    reg = new_base_registry()
+
     # Allow for custom units to be defined in the database
     # Calculate a hash of all custom units
     hash_md5 = md5()
@@ -158,6 +165,41 @@ def reload_unit_registry():
     return reg
 
 
+def build_candidate_unit_registry(
+    pending_fmt_string: str, exclude_pk: Optional[int] = None
+) -> pint.UnitRegistry:
+    """Build a throwaway unit registry, to validate a pending (not yet saved) custom unit definition.
+
+    This constructs the registry that *would* result from saving the pending custom unit,
+    without touching the shared, cached unit registry. This allows us to detect issues
+    (such as a circular reference between two custom units) which only appear once every
+    custom unit definition is loaded together.
+
+    Arguments:
+        pending_fmt_string: The pint format string for the (not yet saved) custom unit
+        exclude_pk: If provided, exclude the CustomUnit with this primary key from the
+            existing database records (used when validating an update to an existing unit)
+
+    Returns:
+        A new pint.UnitRegistry instance, with all custom units (including the pending one) loaded
+    """
+    from common.models import CustomUnit
+
+    reg = new_base_registry()
+
+    custom_units = CustomUnit.objects.all()
+
+    if exclude_pk is not None:
+        custom_units = custom_units.exclude(pk=exclude_pk)
+
+    for cu in custom_units:
+        reg.define(cu.fmt_string())
+
+    reg.define(pending_fmt_string)
+
+    return reg
+
+
 def from_engineering_notation(value):
     """Convert a provided value to 'natural' representation from 'engineering' notation.
 

src/backend/InvenTree/common/api.py

@@ -18,6 +18,7 @@ from django.views.decorators.csrf import csrf_exempt
 import django_filters.rest_framework.filters as rest_filters
 import django_q.models
 import django_q.tasks
+import structlog
 from allauth.socialaccount import providers
 from allauth.socialaccount.models import SocialApp
 from django_filters.rest_framework.filterset import FilterSet
@@ -93,6 +94,8 @@ from InvenTree.permissions import (
 from InvenTree.serializers import EmptySerializer
 from scim.admin_api import ScimConfigViewSet
 
+logger = structlog.get_logger('inventree')
+
 admin_router = InvenTreeApiRouter()
 common_router = InvenTreeApiRouter()
 
@@ -619,11 +622,22 @@ class CustomUnitViewset(DataExportViewMixin, viewsets.ModelViewSet):
     def all(self, request, *args, **kwargs):
         """Return a list of all available units."""
         reg = InvenTree.conversion.get_unit_registry()
-        all_units = {k: self.get_unit(reg, k) for k in reg}
+
+        all_units = {}
+
+        for k in reg:
+            try:
+                if unit := self.get_unit(reg, k):
+                    all_units[k] = unit
+            except Exception:
+                # A single bad unit definition (e.g. a circular reference between
+                # two custom units) should not take down the entire endpoint
+                logger.exception("Failed to process unit '%s' in unit registry", k)
+
         data = {
             'default_system': reg.default_system,
             'available_systems': dir(reg.sys),
-            'available_units': {k: v for k, v in all_units.items() if v},
+            'available_units': all_units,
         }
         return Response(data)
 
@@ -631,11 +645,24 @@ class CustomUnitViewset(DataExportViewMixin, viewsets.ModelViewSet):
         """Parse a unit from the registry."""
         if not hasattr(reg, k):
             return None
+
         unit: type[UnitLike] = getattr(reg, k)
+
+        try:
+            compatible_units = [
+                str(a)
+                for a in unit.compatible_units()  # ty:ignore[missing-argument]
+            ]
+        except Exception:
+            # Guard against e.g. a circular / recursive custom unit definition,
+            # which would otherwise raise an uncaught RecursionError here
+            logger.exception("Failed to determine compatible units for '%s'", k)
+            return None
+
         return {
             'name': k,
             'is_alias': reg.get_name(k) == k,
-            'compatible_units': [str(a) for a in unit.compatible_units()],  # ty:ignore[missing-argument]
+            'compatible_units': compatible_units,
             'isdimensionless': unit.dimensionless,
         }
 

src/backend/InvenTree/common/models.py

@@ -1832,7 +1832,10 @@ class CustomUnit(models.Model):
         """Validate that the provided custom unit is indeed valid."""
         super().clean()
 
-        from InvenTree.conversion import get_unit_registry
+        from InvenTree.conversion import (
+            build_candidate_unit_registry,
+            get_unit_registry,
+        )
 
         registry = get_unit_registry()
 
@@ -1851,12 +1854,28 @@ class CustomUnit(models.Model):
         except Exception as exc:
             raise ValidationError({'definition': str(exc)})
 
-        # Finally, test that the entire custom unit definition is valid
+        # Test that the entire custom unit definition is valid
         try:
             registry.define(self.fmt_string())
         except Exception as exc:
             raise ValidationError(str(exc))
 
+        # Build a registry containing *every* custom unit (including this
+        # pending one), and try to resolve this unit's dimensionality.
+        # Useful for catching recursion errors.
+        try:
+            candidate_registry = build_candidate_unit_registry(
+                self.fmt_string(), exclude_pk=self.pk
+            )
+            getattr(candidate_registry, self.name).compatible_units()
+        except Exception as exc:
+            raise ValidationError(
+                _(
+                    'Unit definition results in a circular or invalid reference: %(error)s'
+                )
+                % {'error': str(exc)}
+            )
+
     name = models.CharField(
         max_length=50,
         verbose_name=_('Name'),

src/backend/InvenTree/common/tests.py

@@ -2137,6 +2137,32 @@ class CustomUnitAPITest(InvenTreeAPITestCase):
         for name in invalid_name_values:
             self.patch(url, {'name': name}, expected_code=400)
 
+    def test_validation_circular(self):
+        """Test that circular / recursive unit definitions are rejected.
+
+        Ref: https://github.com/inventree/InvenTree/issues/12813
+        """
+        self.user.is_staff = True
+        self.user.save()
+
+        a = CustomUnit.objects.create(name='circular_a', definition='meter')
+        b = CustomUnit.objects.create(name='circular_b', definition='3 * circular_a')
+
+        # Editing 'a' to reference 'b' introduces a circular reference
+        response = self.patch(
+            reverse('api-custom-unit-detail', kwargs={'pk': a.pk}),
+            {'definition': '2 * circular_b'},
+            expected_code=400,
+        )
+
+        self.assertIn('non_field_errors', response.data)
+
+        # Ensure the original (non-circular) definition was not overwritten
+        a.refresh_from_db()
+        b.refresh_from_db()
+        self.assertEqual(a.definition, 'meter')
+        self.assertEqual(b.definition, '3 * circular_a')
+
     def test_api(self):
         """Test the CustomUnit API."""
         response = self.get(reverse('api-custom-unit-all'))
@@ -2145,6 +2171,36 @@ class CustomUnitAPITest(InvenTreeAPITestCase):
         self.assertIn('available_units', response.data)
         self.assertEqual(len(response.data['available_units']) > 100, True)
 
+    def test_api_circular_unit(self):
+        """Test that a pre-existing circular unit definition does not break the 'all units' endpoint.
+
+        It is not possible to *create* a circular definition via the API (refer to
+        test_validation_circular), but this test guards against any other way such
+        a definition could end up in the database (e.g. a direct DB edit, or a bug
+        in some other validation path).
+
+        Ref: https://github.com/inventree/InvenTree/issues/12813
+        """
+        import InvenTree.conversion as conversion
+
+        a = CustomUnit.objects.create(name='circular_c', definition='meter')
+        b = CustomUnit.objects.create(name='circular_d', definition='3 * circular_c')
+
+        # Bypass model validation entirely, to simulate a pre-existing bad definition
+        CustomUnit.objects.filter(pk=a.pk).update(definition='2 * circular_d')
+        conversion.reload_unit_registry()
+
+        try:
+            response = self.get(reverse('api-custom-unit-all'), expected_code=200)
+
+            # The broken units are excluded, but the endpoint does not crash
+            self.assertNotIn('circular_c', response.data['available_units'])
+            self.assertNotIn('circular_d', response.data['available_units'])
+            self.assertGreater(len(response.data['available_units']), 100)
+        finally:
+            CustomUnit.objects.filter(pk__in=[a.pk, b.pk]).delete()
+            conversion.reload_unit_registry()
+
 
 class ContentTypeAPITest(InvenTreeAPITestCase):
     """Unit tests for the ContentType API."""