fix: ignore conflicts while bulk creating v2 packages

- This is a TOCTOU problem when multiple workers try to create the same PURL

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

vulnerabilities/models.py

@@ -3402,16 +3402,10 @@ def bulk_get_or_create_from_purls(self, purls: List[Union[PackageURL, str]]):
         """
         Return new or existing Packages given ``purls`` list of PackageURL object or PURL string.
         """
-        purl_strings = [str(p) for p in purls]
-        existing_packages = PackageV2.objects.filter(package_url__in=purl_strings)
-        existing_purls = set(existing_packages.values_list("package_url", flat=True))
-
-        all_packages = list(existing_packages)
         packages_to_create = []
-        for purl in purls:
-            if str(purl) in existing_purls:
-                continue
+        normalize_purls = []
 
+        for purl in purls:
             purl_dict = purl_to_dict(purl)
             purl = PackageURL(**purl_dict)
 
@@ -3422,16 +3416,16 @@ def bulk_get_or_create_from_purls(self, purls: List[Union[PackageURL, str]]):
             purl_dict["package_url"] = str(normalized)
             purl_dict["plain_package_url"] = str(utils.plain_purl(normalized))
 
+            normalize_purls.append(str(normalized))
             packages_to_create.append(PackageV2(**purl_dict))
 
         try:
-            new_packages = PackageV2.objects.bulk_create(packages_to_create)
+            PackageV2.objects.bulk_create(packages_to_create, ignore_conflicts=True)
         except Exception as e:
             logging.error(f"Error creating PackageV2: {e} \n {traceback_format_exc()}")
             return []
 
-        all_packages.extend(new_packages)
-        return all_packages
+        return PackageV2.objects.filter(package_url__in=normalize_purls)
 
     def only_vulnerable(self):
         return self._vulnerable(True)

vulnerabilities/tests/test_models.py

@@ -34,6 +34,7 @@
 from vulnerabilities.models import AdvisorySeverity
 from vulnerabilities.models import Alias
 from vulnerabilities.models import Package
+from vulnerabilities.models import PackageV2
 from vulnerabilities.models import Patch
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.severity_systems import CVSSV3
@@ -827,3 +828,38 @@ def test_advisoryv2_duplication_data(self):
         result = models.AdvisoryV2.objects.count()
 
         self.assertEqual(result, 2)
+
+
+class TestPackageV2BulkCreate(DjangoTestCase):
+    def setUp(self):
+        PackageV2.objects.get_or_create_from_purl(
+            "pkg:deb/ubuntu/linux@6.17.0-19.19?arch=source&distro=questing"
+        )
+        PackageV2.objects.get_or_create_from_purl("pkg:pypi/foo@1.2.3")
+        PackageV2.objects.get_or_create_from_purl("pkg:npm/foobar@3.2.3")
+        PackageV2.objects.get_or_create_from_purl("pkg:maven/foo@1.2.3")
+        PackageV2.objects.get_or_create_from_purl(
+            "pkg:deb/ubuntu/linux@6.17.0-4.4?arch=source&distro=questing"
+        )
+
+    def test_package_bulk_get_or_create_from_purls(self):
+        purls = [
+            "pkg:npm/foo@1.2.3",
+            "pkg:pypi/foo@1.2.3",
+            "pkg:npm/foobar@3.2.3",
+            "pkg:maven/foo@1.2.3",
+            "pkg:nuget/foo@1.2.3",
+            "pkg:deb/ubuntu/linux@6.17.0-22.22?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-20.20?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-14.14?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-12.12?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-8.8?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-7.7?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-6.6?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-5.5?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-4.4?arch=source&distro=questing",
+        ]
+        result_qs = PackageV2.objects.bulk_get_or_create_from_purls(purls)
+        result = [p.package_url for p in result_qs]
+
+        self.assertCountEqual(result, purls)