IA-5415: Fix org unit types infinite recursion

dynamic_fields/serializer.py

@@ -53,6 +53,8 @@ def to_representation(self, instance):
         return super().to_representation(instance)
 
     def get_fields_value_from_query(self):
+        if self.context.get("ignore_dynamic_fields", False):
+            return []
         if self.context.get("request", None):
             if self.string_field:
                 return [

iaso/api/org_unit_types/serializers.py

@@ -11,18 +11,36 @@
 from ..projects.serializers import ProjectSerializer
 
 
-def get_parents(type_id):
-    parents_ids = []
-    if type_id is not None:
-        queryset = OrgUnitType.objects.filter(sub_unit_types__id=type_id)
-        for parent in queryset.all():
-            if parent.id not in parents_ids:
-                parents_ids.append(parent.id)
-            great_parents = get_parents(parent.id)
-            for great_parent_id in great_parents:
-                if great_parent_id not in parents_ids:
-                    parents_ids.append(great_parent_id)
-    return parents_ids
+def has_cycle_after_adding(node_id, proposed_children, relation_name):
+    """
+    Returns True if adding proposed_children to node_id's relation_name
+    creates a cycle/loop.
+    """
+    if not node_id:
+        return False
+
+    proposed_children_ids = [c.id for c in proposed_children]
+    if node_id in proposed_children_ids:
+        return True
+
+    visited = set(proposed_children_ids)
+    queue = list(proposed_children_ids)
+
+    while queue:
+        curr_id = queue.pop(0)
+        if curr_id == node_id:
+            return True
+        try:
+            curr_node = OrgUnitType.objects.get(pk=curr_id)
+            children_field = getattr(curr_node, relation_name)
+            for child_id in children_field.values_list("id", flat=True):
+                if child_id not in visited:
+                    visited.add(child_id)
+                    queue.append(child_id)
+        except OrgUnitType.DoesNotExist:
+            continue
+
+    return False
 
 
 def validate_reference_forms(data):
@@ -105,25 +123,29 @@ def get_sub_unit_types(self, obj: OrgUnitType):
             unit_types,
             fields=["id", "name", "short_name", "depth", "created_at", "updated_at"],
             many=True,
-            context=self.context,
+            context={**self.context, "ignore_dynamic_fields": True},
         ).data
 
     def validate(self, data: typing.Mapping):
-        parents = get_parents(self.context["request"].data.get("id", None))
+        org_unit_type_id = self.instance.id if self.instance else self.context["request"].data.get("id")
         # validate sub org unit type
-        sub_types_errors = []
-        for sub_type in data.get("sub_unit_types", []):
-            if sub_type.id in parents:
-                sub_types_errors.append(sub_type.name)
-        if len(sub_types_errors) > 0:
-            raise serializers.ValidationError({"sub_unit_type_ids": sub_types_errors})
+        if "sub_unit_types" in data:
+            if has_cycle_after_adding(org_unit_type_id, data["sub_unit_types"], "sub_unit_types"):
+                raise serializers.ValidationError(
+                    {"sub_unit_type_ids": ["A loop was detected in the sub-unit types hierarchy."]}
+                )
         # validate sub org unit type allowed to be created
-        create_sub_types_errors = []
-        for sub_type in data.get("allow_creating_sub_unit_types", []):
-            if sub_type.id in parents:
-                create_sub_types_errors.append(sub_type.name)
-        if len(create_sub_types_errors) > 0:
-            raise serializers.ValidationError({"allow_creating_sub_unit_type_ids": create_sub_types_errors})
+        if "allow_creating_sub_unit_types" in data:
+            if has_cycle_after_adding(
+                org_unit_type_id, data["allow_creating_sub_unit_types"], "allow_creating_sub_unit_types"
+            ):
+                raise serializers.ValidationError(
+                    {
+                        "allow_creating_sub_unit_type_ids": [
+                            "A loop was detected in the allowed sub-unit types creation hierarchy."
+                        ]
+                    }
+                )
         # validate projects (access check)
         for project in data.get("projects", []):
             if self.context["request"].user.iaso_profile.account != project.account:
@@ -223,7 +245,7 @@ def get_sub_unit_types(self, obj: OrgUnitType):
             unit_types,
             fields=["id", "name", "short_name", "depth", "created_at", "updated_at"],
             many=True,
-            context=self.context,
+            context={**self.context, "ignore_dynamic_fields": True},
         ).data
 
     def get_allow_creating_sub_unit_types(self, obj: OrgUnitType):
@@ -237,25 +259,29 @@ def get_allow_creating_sub_unit_types(self, obj: OrgUnitType):
             unit_types,
             fields=["id", "name", "short_name", "depth", "created_at", "updated_at"],
             many=True,
-            context=self.context,
+            context={**self.context, "ignore_dynamic_fields": True},
         ).data
 
     def validate(self, data: typing.Mapping):
-        parents = get_parents(self.context["request"].data.get("id", None))
+        org_unit_type_id = self.instance.id if self.instance else self.context["request"].data.get("id")
         # validate sub org unit type
-        sub_types_errors = []
-        for sub_type in data.get("sub_unit_types", []):
-            if sub_type.id in parents:
-                sub_types_errors.append(sub_type.name)
-        if len(sub_types_errors) > 0:
-            raise serializers.ValidationError({"sub_unit_type_ids": sub_types_errors})
+        if "sub_unit_types" in data:
+            if has_cycle_after_adding(org_unit_type_id, data["sub_unit_types"], "sub_unit_types"):
+                raise serializers.ValidationError(
+                    {"sub_unit_type_ids": ["A loop was detected in the sub-unit types hierarchy."]}
+                )
         # validate sub org unit type allowed to be created
-        create_sub_types_errors = []
-        for sub_type in data.get("allow_creating_sub_unit_types", []):
-            if sub_type.id in parents:
-                create_sub_types_errors.append(sub_type.name)
-        if len(create_sub_types_errors) > 0:
-            raise serializers.ValidationError({"allow_creating_sub_unit_type_ids": create_sub_types_errors})
+        if "allow_creating_sub_unit_types" in data:
+            if has_cycle_after_adding(
+                org_unit_type_id, data["allow_creating_sub_unit_types"], "allow_creating_sub_unit_types"
+            ):
+                raise serializers.ValidationError(
+                    {
+                        "allow_creating_sub_unit_type_ids": [
+                            "A loop was detected in the allowed sub-unit types creation hierarchy."
+                        ]
+                    }
+                )
         # validate projects (access check)
         for project in data.get("projects", []):
             if self.context["request"].user.iaso_profile.account != project.account:
@@ -284,8 +310,17 @@ class Meta:
 
     def get_sub_unit_types(self, obj):
         """Recursively serialize sub_unit_types to build complete hierarchy"""
+        visited = self.context.setdefault("visited_org_unit_types", set())
+        if obj.id in visited:
+            return []
+
+        child_visited = set(visited)
+        child_visited.add(obj.id)
+        child_context = {**self.context, "visited_org_unit_types": child_visited}
+
         sub_types = obj.sub_unit_types.all()
-        return OrgUnitTypeHierarchySerializer(sub_types, many=True, context=self.context).data
+        sub_types = [t for t in sub_types if t.id not in child_visited]
+        return OrgUnitTypeHierarchySerializer(sub_types, many=True, context=child_context).data
 
 
 class OrgUnitTypesDropdownSerializer(serializers.ModelSerializer):

iaso/tests/api/test_org_unit_types_v2.py

@@ -776,3 +776,68 @@ def test_num_queries_retrieve(self):
             response = self.client.get(f"{self.BASE_URL}{self.org_unit_type_1.id}/")
 
         self.assertJSONResponse(response, status.HTTP_200_OK)
+
+    def test_prevent_loop_sub_unit_types(self):
+        """Test that `sub_unit_types` is validated to prevent infinite recursion loops."""
+        self.client.force_authenticate(self.jane)
+        self.org_unit_type_1.sub_unit_types.set([self.org_unit_type_2])
+
+        response = self.client.put(
+            f"{self.BASE_URL}{self.org_unit_type_2.id}/",
+            data={
+                "name": self.org_unit_type_2.name,
+                "short_name": self.org_unit_type_2.short_name,
+                "project_ids": [self.ead.id],
+                "sub_unit_type_ids": [self.org_unit_type_1.id],
+                "allow_creating_sub_unit_type_ids": [],
+            },
+            format="json",
+        )
+        data = self.assertJSONResponse(response, status.HTTP_400_BAD_REQUEST)
+        self.assertIn("sub_unit_type_ids", data)
+
+    def test_prevent_loop_allow_creating_sub_unit_types(self):
+        """Test that `allow_creating_sub_unit_types` is validated to prevent infinite recursion loops."""
+        self.client.force_authenticate(self.jane)
+        self.org_unit_type_1.allow_creating_sub_unit_types.set([self.org_unit_type_2])
+
+        response = self.client.put(
+            f"{self.BASE_URL}{self.org_unit_type_2.id}/",
+            data={
+                "name": self.org_unit_type_2.name,
+                "short_name": self.org_unit_type_2.short_name,
+                "project_ids": [self.ead.id],
+                "sub_unit_type_ids": [],
+                "allow_creating_sub_unit_type_ids": [self.org_unit_type_1.id],
+            },
+            format="json",
+        )
+        data = self.assertJSONResponse(response, status.HTTP_400_BAD_REQUEST)
+        self.assertIn("allow_creating_sub_unit_type_ids", data)
+
+    def test_serialization_resilience_to_existing_loops(self):
+        """Test that if a cycle exists, querying the api doesn't throw a recursion error."""
+        self.client.force_authenticate(self.jane)
+        self.org_unit_type_1.sub_unit_types.set([self.org_unit_type_2])
+        self.org_unit_type_2.sub_unit_types.set([self.org_unit_type_1])
+
+        response = self.client.get(f"{self.BASE_URL}?fields=id,name,sub_unit_types")
+        self.assertJSONResponse(response, status.HTTP_200_OK)
+
+    def test_serialization_resilience_to_existing_loops_allow_creating_sub_unit_types(self):
+        """Test that if a cycle exists on allow_creating_sub_unit_types, querying the api doesn't throw a recursion error."""
+        self.client.force_authenticate(self.jane)
+        self.org_unit_type_1.allow_creating_sub_unit_types.set([self.org_unit_type_2])
+        self.org_unit_type_2.allow_creating_sub_unit_types.set([self.org_unit_type_1])
+
+        response = self.client.get(f"{self.BASE_URL}?fields=id,name,allow_creating_sub_unit_types")
+        self.assertJSONResponse(response, status.HTTP_200_OK)
+
+    def test_hierarchy_resilience_to_existing_loops(self):
+        """Test that if a cycle exists, querying the hierarchy endpoint doesn't throw a recursion error."""
+        self.client.force_authenticate(self.jane)
+        self.org_unit_type_1.sub_unit_types.set([self.org_unit_type_2])
+        self.org_unit_type_2.sub_unit_types.set([self.org_unit_type_1])
+
+        response = self.client.get(f"{self.BASE_URL}{self.org_unit_type_1.id}/hierarchy/")
+        self.assertJSONResponse(response, status.HTTP_200_OK)