Просмотр исходного кода

Fixes #23445: Prevent concurrent REST API updates from reverting each other's changes (#23459)

Jeremy Stretch 2 дней назад
Родитель
Сommit
bb93fb913e

+ 1 - 1
docs/integrations/rest-api.md

@@ -701,7 +701,7 @@ $ curl -s -X PATCH \
     --data '{"status": "decommissioning"}'
     --data '{"status": "decommissioning"}'
 ```
 ```
 
 
-A literal `If-Match: *` value matches any current ETag and may be used to assert simply that the object exists. Submitting `If-Match` is optional; requests without the header retain prior (last-write-wins) behavior.
+A literal `If-Match: *` value matches any current ETag and may be used to assert simply that the object exists. Submitting `If-Match` is optional; requests without the header retain prior (last-write-wins) behavior for the fields they set.
 
 
 ### Adding and Removing Tags
 ### Adding and Removing Tags
 
 

+ 39 - 11
netbox/netbox/api/viewsets/__init__.py

@@ -1,8 +1,9 @@
 import logging
 import logging
 import warnings
 import warnings
+from contextlib import contextmanager
 from functools import cached_property
 from functools import cached_property
 
 
-from django.core.exceptions import ObjectDoesNotExist, PermissionDenied
+from django.core.exceptions import FieldError, ObjectDoesNotExist, PermissionDenied, ValidationError
 from django.db import router, transaction
 from django.db import router, transaction
 from django.db.models import ProtectedError, RestrictedError
 from django.db.models import ProtectedError, RestrictedError
 from rest_framework import mixins as drf_mixins
 from rest_framework import mixins as drf_mixins
@@ -208,6 +209,7 @@ class NetBoxModelViewSet(
     mixins.BulkCreateModelMixin,
     mixins.BulkCreateModelMixin,
     mixins.BulkUpdateModelMixin,
     mixins.BulkUpdateModelMixin,
     mixins.BulkDestroyModelMixin,
     mixins.BulkDestroyModelMixin,
+    mixins.ObjectLockingMixin,
     mixins.ObjectValidationMixin,
     mixins.ObjectValidationMixin,
     mixins.CustomFieldsMixin,
     mixins.CustomFieldsMixin,
     mixins.ExportTemplatesMixin,
     mixins.ExportTemplatesMixin,
@@ -231,6 +233,22 @@ class NetBoxModelViewSet(
             obj.snapshot()
             obj.snapshot()
         return obj
         return obj
 
 
+    @contextmanager
+    def lock_object(self):
+        """
+        Lock the row of the requested object for the duration of the block (see lock_objects()). Enter this
+        immediately within a transaction, before retrieving the object.
+        """
+        manager = self.queryset.model._default_manager
+        lookup_url_kwarg = self.lookup_url_kwarg or self.lookup_field
+        try:
+            queryset = manager.filter(**{self.lookup_field: self.kwargs[lookup_url_kwarg]})
+        except (KeyError, FieldError, TypeError, ValueError, ValidationError):
+            # The URL has no lookup value, or the lookup is not valid for the model
+            queryset = manager.none()
+        with self.lock_objects(queryset):
+            yield
+
     def get_queryset(self):
     def get_queryset(self):
         qs = super().get_queryset()
         qs = super().get_queryset()
         return reapply_model_ordering(qs)
         return reapply_model_ordering(qs)
@@ -346,14 +364,20 @@ class NetBoxModelViewSet(
 
 
     def update(self, request, *args, **kwargs):
     def update(self, request, *args, **kwargs):
         partial = kwargs.pop('partial', False)
         partial = kwargs.pop('partial', False)
-        instance = self.get_object_with_snapshot()
 
 
-        # Enforce If-Match precondition (RFC 9110 §13.1.1)
-        self._validate_etag(self.request, instance)
+        # Lock the object before retrieving it, so that concurrent updates to the same object are applied one after
+        # the other. Otherwise, each would save its own copy of the object, reverting any fields which another
+        # update had changed since that copy was retrieved.
+        using = router.db_for_write(self.queryset.model)
+        with transaction.atomic(using=using), self.lock_object():
+            instance = self.get_object_with_snapshot()
 
 
-        serializer = self.get_serializer(instance, data=request.data, partial=partial)
-        serializer.is_valid(raise_exception=True)
-        self.perform_update(serializer)
+            # Enforce If-Match precondition (RFC 9110 §13.1.1)
+            self._validate_etag(self.request, instance)
+
+            serializer = self.get_serializer(instance, data=request.data, partial=partial)
+            serializer.is_valid(raise_exception=True)
+            self.perform_update(serializer)
 
 
         # After updating the instance, re-initialize the serializer with a queryset
         # After updating the instance, re-initialize the serializer with a queryset
         # to ensure related objects are prefetched.
         # to ensure related objects are prefetched.
@@ -377,11 +401,15 @@ class NetBoxModelViewSet(
         using = router.db_for_write(model)
         using = router.db_for_write(model)
         try:
         try:
             with transaction.atomic(using=using), mixins.discard_events_on_rollback(self, using=using):
             with transaction.atomic(using=using), mixins.discard_events_on_rollback(self, using=using):
-                # Re-check the If-Match ETag under a row-level lock to close the TOCTOU window
-                # between the initial check in update() and the actual write.
+                # Check the If-Match ETag under a row-level lock to close the TOCTOU window between any
+                # earlier check and the actual write. An object retrieved only once its row had been
+                # locked is already current; otherwise, lock its row and re-fetch it.
                 if self._get_if_match(self.request):
                 if self._get_if_match(self.request):
-                    locked = model.objects.select_for_update().get(pk=serializer.instance.pk)
-                    self._validate_etag(self.request, locked)
+                    if serializer.instance.pk in self.locked_pks:
+                        current = serializer.instance
+                    else:
+                        current = model._default_manager.select_for_update().get(pk=serializer.instance.pk)
+                    self._validate_etag(self.request, current)
                 instance = serializer.save()
                 instance = serializer.save()
                 self._validate_objects(instance)
                 self._validate_objects(instance)
         except ObjectDoesNotExist:
         except ObjectDoesNotExist:

+ 62 - 24
netbox/netbox/api/viewsets/mixins.py

@@ -30,10 +30,12 @@ __all__ = (
     'BulkUpdateModelMixin',
     'BulkUpdateModelMixin',
     'CustomFieldsMixin',
     'CustomFieldsMixin',
     'ExportTemplatesMixin',
     'ExportTemplatesMixin',
+    'ObjectLockingMixin',
     'ObjectValidationMixin',
     'ObjectValidationMixin',
     'discard_events_on_rollback',
     'discard_events_on_rollback',
     'get_duplicate_objects_response',
     'get_duplicate_objects_response',
     'get_invalid_entries_response',
     'get_invalid_entries_response',
+    'get_missing_objects_errors',
     'get_missing_objects_response',
     'get_missing_objects_response',
     'get_non_list_response',
     'get_non_list_response',
     'resolve_bulk_error_status',
     'resolve_bulk_error_status',
@@ -194,37 +196,28 @@ def get_duplicate_objects_response(object_ids):
     )
     )
 
 
 
 
-def get_missing_objects_response(object_ids, queryset):
+def get_missing_objects_errors(object_ids, found_pks):
     """
     """
-    Return a structured error Response naming each of the given object IDs which the queryset does
-    not match, or None if it matches them all.
-
-    An ID goes unmatched either because no such object exists or because the requesting user's
-    object-level permissions exclude it. The two are deliberately not distinguished, consistent with
-    the single-object endpoints, which return a 404 in both cases.
-
-    Bulk operations call this before performing any work: an unresolvable ID means the request names
-    an object the client cannot act on, so there is nothing to be gained by attempting the batch (it
-    would only be rolled back). Note that the status is 400 rather than the 409 a bulk delete
-    returns for a dependency conflict, as this is a problem with the request itself rather than with
-    the current state of the database.
+    Return a list of per-object errors naming each of the given object IDs which is not among the
+    PKs found.
     """
     """
-    found_pks = set(queryset.values_list('pk', flat=True))
-
-    errors = [
+    return [
         {
         {
             'id': object_id,
             'id': object_id,
             'errors': {
             'errors': {
                 'id': [_("Object with ID {id} does not exist").format(id=object_id)],
                 'id': [_("Object with ID {id} does not exist").format(id=object_id)],
             },
             },
         }
         }
-        # NetBox's bulk actions reject a repeated ID before reaching this point (see
-        # get_duplicate_objects_response), but de-duplicate anyway so that any other caller reports
-        # such an ID once rather than once per occurrence. dict.fromkeys() preserves the order of
-        # first appearance.
-        for object_id in dict.fromkeys(object_ids)
-        if object_id not in found_pks
+        for object_id in dict.fromkeys(object_ids) if object_id not in found_pks
     ]
     ]
+
+
+def get_missing_objects_response(object_ids, queryset):
+    """
+    Return a structured error Response naming each of the given object IDs which the queryset does
+    not match, or None if it matches them all.
+    """
+    errors = get_missing_objects_errors(object_ids, set(queryset.values_list('pk', flat=True)))
     if not errors:
     if not errors:
         return None
         return None
 
 
@@ -278,6 +271,37 @@ def discard_events_on_rollback(sender, using=None):
         clear_events.send(sender=sender)
         clear_events.send(sender=sender)
 
 
 
 
+class ObjectLockingMixin:
+    """
+    Lock the rows of objects to be modified before retrieving them.
+    """
+    # The primary keys of objects whose rows are locked by the current transaction
+    locked_pks = frozenset()
+
+    @contextmanager
+    def lock_objects(self, queryset):
+        """
+        Lock the rows of the objects matched by the given QuerySet, and record their primary keys as
+        `locked_pks` for the duration of the block. Enter this immediately within the transaction to
+        which the locks belong (they are held until it ends), before retrieving the objects to be
+        modified.
+        """
+        self.locked_pks = set(
+            queryset
+            # Lock rows in a consistent order, so that concurrent requests locking overlapping sets of
+            # objects cannot deadlock. This also clears any default ordering, which may require JOINs.
+            .order_by('pk')
+            # no_key: Avoid blocking foreign key inserts that reference these objects
+            .select_for_update(no_key=True)
+            .values_list('pk', flat=True)
+        )
+        try:
+            yield
+        finally:
+            # The locks are released when the transaction ends
+            self.locked_pks = frozenset()
+
+
 class BackgroundOperationMixin:
 class BackgroundOperationMixin:
     """
     """
     Enable optional background processing of REST API bulk write operations. When a write
     Enable optional background processing of REST API bulk write operations. When a write
@@ -512,7 +536,7 @@ class BulkCreateModelMixin:
         return created_pks, errors, resolve_bulk_error_status(error_statuses)
         return created_pks, errors, resolve_bulk_error_status(error_statuses)
 
 
 
 
-class BulkUpdateModelMixin:
+class BulkUpdateModelMixin(ObjectLockingMixin):
     """
     """
     Support bulk modification of objects using the list endpoint for a model. Accepts a PATCH action with a list of one
     Support bulk modification of objects using the list endpoint for a model. Accepts a PATCH action with a list of one
     or more JSON objects, each specifying the numeric ID of an object to be updated as well as the attributes to be set.
     or more JSON objects, each specifying the numeric ID of an object to be updated as well as the attributes to be set.
@@ -611,7 +635,21 @@ class BulkUpdateModelMixin:
         errors = []
         errors = []
         error_statuses = set()
         error_statuses = set()
         using = router.db_for_write(self.queryset.model)
         using = router.db_for_write(self.queryset.model)
-        with transaction.atomic(using=using), discard_events_on_rollback(self, using=using):
+        with (
+            transaction.atomic(using=using),
+            discard_events_on_rollback(self, using=using),
+            # Lock the objects before retrieving them, so that concurrent updates to the same objects
+            # are applied one after the other
+            self.lock_objects(self.queryset.model._default_manager.filter(pk__in=update_data.keys())),
+        ):
+            objects = list(objects)
+
+            # Report any object which can no longer be found (e.g. because it was deleted after
+            # get_missing_objects_response() checked for it), rather than silently omitting it
+            if missing := get_missing_objects_errors(update_data.keys(), {obj.pk for obj in objects}):
+                errors.extend(missing)
+                error_statuses.add(status.HTTP_400_BAD_REQUEST)
+
             # Validate and save each object in turn so subsequent validations see the DB
             # Validate and save each object in turn so subsequent validations see the DB
             # state left by prior saves (e.g. two items renamed to the same name: the second
             # state left by prior saves (e.g. two items renamed to the same name: the second
             # will fail validation rather than raising an integrity error on save).
             # will fail validation rather than raising an integrity error on save).

+ 135 - 1
netbox/netbox/tests/test_api.py

@@ -4,9 +4,10 @@ from unittest.mock import patch
 
 
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import NON_FIELD_ERRORS
 from django.core.exceptions import NON_FIELD_ERRORS
-from django.db import DEFAULT_DB_ALIAS, connections
+from django.db import DEFAULT_DB_ALIAS, connection, connections
 from django.db.backends.postgresql.psycopg_any import NumericRange
 from django.db.backends.postgresql.psycopg_any import NumericRange
 from django.test import RequestFactory, TestCase, TransactionTestCase, override_settings, tag
 from django.test import RequestFactory, TestCase, TransactionTestCase, override_settings, tag
+from django.test.utils import CaptureQueriesContext
 from django.urls import reverse
 from django.urls import reverse
 from django.utils.timezone import now
 from django.utils.timezone import now
 from rest_framework.exceptions import ValidationError
 from rest_framework.exceptions import ValidationError
@@ -16,6 +17,7 @@ from rest_framework.test import APIClient
 
 
 from core.models import DataFile, DataSource, ObjectType
 from core.models import DataFile, DataSource, ObjectType
 from dcim.api.serializers import InterfaceSerializer, RackSerializer
 from dcim.api.serializers import InterfaceSerializer, RackSerializer
+from dcim.api.views import SiteViewSet
 from dcim.models import Device, Rack, Site
 from dcim.models import Device, Rack, Site
 from extras.models import ExportTemplate
 from extras.models import ExportTemplate
 from ipam.api.serializers import RouteTargetSerializer, VRFSerializer
 from ipam.api.serializers import RouteTargetSerializer, VRFSerializer
@@ -24,6 +26,7 @@ from netbox.api.exceptions import QuerySetNotOrdered, SerializerNotFound
 from netbox.api.fields import ContentTypeField, IntegerRangeSerializer, RelatedObjectCountField
 from netbox.api.fields import ContentTypeField, IntegerRangeSerializer, RelatedObjectCountField
 from netbox.api.pagination import NetBoxPagination
 from netbox.api.pagination import NetBoxPagination
 from netbox.api.serializers import ValidatedModelSerializer, WritableNestedSerializer
 from netbox.api.serializers import ValidatedModelSerializer, WritableNestedSerializer
+from netbox.api.viewsets.mixins import BulkUpdateModelMixin
 from users.constants import TOKEN_PREFIX
 from users.constants import TOKEN_PREFIX
 from users.models import Token, User
 from users.models import Token, User
 from utilities.api import get_serializer_for_model
 from utilities.api import get_serializer_for_model
@@ -572,3 +575,134 @@ class BulkOperationRoutingTestCase(TransactionTestCase):
         self.assertEqual([e['id'] for e in response.data['errors']], [sites[0].pk])
         self.assertEqual([e['id'] for e in response.data['errors']], [sites[0].pk])
         self.assertEqual(Site.objects.count(), 2)
         self.assertEqual(Site.objects.count(), 2)
         flush_events.assert_not_called()
         flush_events.assert_not_called()
+
+
+class ObjectUpdateLockingTestCase(APITestCase):
+    """
+    An update must lock the rows of the objects being updated before retrieving them, and hold the locks until the
+    objects have been saved. A concurrent update to the same objects then waits for this one to complete, rather than
+    having its changes reverted when this one saves copies of the objects retrieved before they were committed.
+    """
+
+    @classmethod
+    def setUpTestData(cls):
+        cls.sites = (
+            Site(name='Site 1', slug='site-1'),
+            Site(name='Site 2', slug='site-2'),
+        )
+        Site.objects.bulk_create(cls.sites)
+
+    def setUp(self):
+        super().setUp()
+        self.add_permissions('dcim.change_site')
+
+    def assertLockedUntilSaved(self, queries):
+        """
+        Assert that the given queries lock the rows of the Sites being updated (once) before retrieving them, within a
+        transaction which does not end until they have been saved.
+        """
+        table = Site._meta.db_table
+        locks = [
+            i for i, sql in enumerate(queries)
+            if f'FROM "{table}"' in sql and sql.endswith(('FOR UPDATE', 'FOR NO KEY UPDATE'))
+        ]
+        self.assertEqual(len(locks), 1, "Expected exactly one query to lock the Sites' rows")
+        lock = locks[0]
+        # The first query to select the Sites' fields (rather than only their PKs)
+        retrieval = next(i for i, sql in enumerate(queries) if sql.startswith(f'SELECT "{table}"."id", '))
+        self.assertLess(lock, retrieval, "The Sites were retrieved before their rows were locked")
+        last_save = max(i for i, sql in enumerate(queries) if sql.startswith(f'UPDATE "{table}"'))
+
+        # Find the transaction which took the lock: the innermost savepoint open at that point (as each test runs
+        # within a transaction, a transaction opened by the view appears as a savepoint)
+        savepoints = []
+        for sql in queries[:lock]:
+            if sql.startswith('SAVEPOINT '):
+                savepoints.append(sql.removeprefix('SAVEPOINT '))
+            elif sql.startswith(('RELEASE SAVEPOINT ', 'ROLLBACK TO SAVEPOINT ')):
+                savepoints.remove(sql.split()[-1])
+        self.assertTrue(savepoints, "The Sites' rows were not locked within a transaction")
+        self.assertIn(
+            f'RELEASE SAVEPOINT {savepoints[-1]}',
+            queries[last_save + 1:],
+            "The transaction which locked the Sites' rows ended before they were saved",
+        )
+
+    @tag('regression')  # Ref: #23445
+    def test_update_locks_object_before_retrieving_it(self):
+        """
+        Lock the row of an object before retrieving it for an update. The If-Match ETag is then checked against the
+        object retrieved, rather than against a copy re-fetched under a second lock.
+        """
+        site = Site.objects.get(pk=self.sites[0].pk)
+        url = reverse('dcim-api:site-detail', kwargs={'pk': site.pk})
+        header = {**self.header, 'HTTP_IF_MATCH': SiteViewSet._get_etag(site)}
+
+        with CaptureQueriesContext(connection) as ctx:
+            response = self.client.patch(url, {'description': 'Description'}, format='json', **header)
+        self.assertHttpStatus(response, 200)
+        self.assertLockedUntilSaved([q['sql'] for q in ctx.captured_queries])
+
+    @tag('regression')  # Ref: #23445
+    def test_bulk_update_locks_objects_before_retrieving_them(self):
+        """
+        Lock the rows of all objects being updated in bulk before retrieving them. The If-Match ETag is then checked
+        against each object retrieved, rather than against a copy re-fetched under a second lock.
+        """
+        sites = Site.objects.filter(pk__in=[site.pk for site in self.sites])
+        data = [{'id': site.pk, 'description': 'Description'} for site in sites]
+        header = {**self.header, 'HTTP_IF_MATCH': ', '.join(SiteViewSet._get_etag(site) for site in sites)}
+
+        with CaptureQueriesContext(connection) as ctx:
+            response = self.client.patch(reverse('dcim-api:site-list'), data, format='json', **header)
+        self.assertHttpStatus(response, 200)
+        self.assertLockedUntilSaved([q['sql'] for q in ctx.captured_queries])
+
+    def test_bulk_update_with_stale_etag(self):
+        """Reject a bulk update if any object's current ETag does not match the If-Match header."""
+        data = [{'id': site.pk, 'description': 'Description'} for site in self.sites]
+        header = {**self.header, 'HTTP_IF_MATCH': 'W/"stale"'}
+
+        response = self.client.patch(reverse('dcim-api:site-list'), data, format='json', **header)
+        self.assertHttpStatus(response, 412)
+        self.assertFalse(Site.objects.filter(description='Description').exists())
+
+    @tag('regression')  # Ref: #23445
+    def test_bulk_update_with_object_deleted_concurrently(self):
+        """
+        Reject a bulk update if any object can no longer be found once the objects have been locked (e.g. because it
+        was deleted after the view checked that each object exists), rather than silently omitting it.
+        """
+        data = [{'id': site.pk, 'description': 'Description'} for site in self.sites]
+        perform_bulk_update = BulkUpdateModelMixin.perform_bulk_update
+
+        def delete_then_perform_bulk_update(viewset, *args, **kwargs):
+            # Delete a Site between the view's check for missing objects and the bulk update
+            Site.objects.filter(pk=self.sites[1].pk).delete()
+            return perform_bulk_update(viewset, *args, **kwargs)
+
+        with patch.object(BulkUpdateModelMixin, 'perform_bulk_update', delete_then_perform_bulk_update):
+            response = self.client.patch(reverse('dcim-api:site-list'), data, format='json', **self.header)
+        self.assertHttpStatus(response, 400)
+        self.assertEqual([e['id'] for e in response.data['errors']], [self.sites[1].pk])
+        self.assertFalse(Site.objects.filter(description='Description').exists())
+
+    def test_lock_object(self):
+        """
+        Lock the requested object for the duration of the block if it can be found by its URL lookup, and otherwise
+        take no lock.
+        """
+        for attrs, locked_pks in (
+            ({'kwargs': {'pk': self.sites[0].pk}}, {self.sites[0].pk}),
+            # Not a valid PK (e.g. for a viewset whose get_object() accepts a slug)
+            ({'kwargs': {'pk': 'site-1'}}, set()),
+            # No lookup value
+            ({'kwargs': {}}, set()),
+            # Not a model field (e.g. for a viewset whose get_object() filters on an annotation)
+            ({'kwargs': {'circuit_count': 0}, 'lookup_field': 'circuit_count'}, set()),
+        ):
+            with self.subTest(**attrs):
+                viewset = SiteViewSet(**attrs)
+                with viewset.lock_object():
+                    self.assertEqual(viewset.locked_pks, locked_pks)
+                self.assertEqual(viewset.locked_pks, set())