소스 검색

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

Jeremy Stretch 2 일 전
부모
커밋
bb93fb913e
4개의 변경된 파일과 237개의 추가작업 그리고 37개의 파일을 삭제
  1. 1 1
      docs/integrations/rest-api.md
  2. 39 11
      netbox/netbox/api/viewsets/__init__.py
  3. 62 24
      netbox/netbox/api/viewsets/mixins.py
  4. 135 1
      netbox/netbox/tests/test_api.py

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

@@ -701,7 +701,7 @@ $ curl -s -X PATCH \
     --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
 

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

@@ -1,8 +1,9 @@
 import logging
 import warnings
+from contextlib import contextmanager
 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.models import ProtectedError, RestrictedError
 from rest_framework import mixins as drf_mixins
@@ -208,6 +209,7 @@ class NetBoxModelViewSet(
     mixins.BulkCreateModelMixin,
     mixins.BulkUpdateModelMixin,
     mixins.BulkDestroyModelMixin,
+    mixins.ObjectLockingMixin,
     mixins.ObjectValidationMixin,
     mixins.CustomFieldsMixin,
     mixins.ExportTemplatesMixin,
@@ -231,6 +233,22 @@ class NetBoxModelViewSet(
             obj.snapshot()
         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):
         qs = super().get_queryset()
         return reapply_model_ordering(qs)
@@ -346,14 +364,20 @@ class NetBoxModelViewSet(
 
     def update(self, request, *args, **kwargs):
         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
         # to ensure related objects are prefetched.
@@ -377,11 +401,15 @@ class NetBoxModelViewSet(
         using = router.db_for_write(model)
         try:
             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):
-                    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()
                 self._validate_objects(instance)
         except ObjectDoesNotExist:

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

@@ -30,10 +30,12 @@ __all__ = (
     'BulkUpdateModelMixin',
     'CustomFieldsMixin',
     'ExportTemplatesMixin',
+    'ObjectLockingMixin',
     'ObjectValidationMixin',
     'discard_events_on_rollback',
     'get_duplicate_objects_response',
     'get_invalid_entries_response',
+    'get_missing_objects_errors',
     'get_missing_objects_response',
     'get_non_list_response',
     '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,
             'errors': {
                 '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:
         return None
 
@@ -278,6 +271,37 @@ def discard_events_on_rollback(sender, using=None):
         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:
     """
     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)
 
 
-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
     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 = []
         error_statuses = set()
         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
             # 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).

+ 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.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.test import RequestFactory, TestCase, TransactionTestCase, override_settings, tag
+from django.test.utils import CaptureQueriesContext
 from django.urls import reverse
 from django.utils.timezone import now
 from rest_framework.exceptions import ValidationError
@@ -16,6 +17,7 @@ from rest_framework.test import APIClient
 
 from core.models import DataFile, DataSource, ObjectType
 from dcim.api.serializers import InterfaceSerializer, RackSerializer
+from dcim.api.views import SiteViewSet
 from dcim.models import Device, Rack, Site
 from extras.models import ExportTemplate
 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.pagination import NetBoxPagination
 from netbox.api.serializers import ValidatedModelSerializer, WritableNestedSerializer
+from netbox.api.viewsets.mixins import BulkUpdateModelMixin
 from users.constants import TOKEN_PREFIX
 from users.models import Token, User
 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(Site.objects.count(), 2)
         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())