Use token-based ownership lock and assert release path in tests

Author: mn-ram

openwisp_controller/connection/tasks.py

@@ -1,5 +1,6 @@
 import logging
 import time
+import uuid
 
 import swapper
 from celery import shared_task
@@ -23,17 +24,27 @@
 def _acquire_update_config_lock(device_id):
     """
     Attempts to atomically acquire a per-device lock using the Django cache.
-    Returns True if the lock was acquired, False if another task already holds it.
+    Returns a unique token string if the lock was acquired, None otherwise.
+    The token must be passed to _release_update_config_lock to ensure
+    only the lock owner can release it.
     """
     lock_key = _UPDATE_CONFIG_LOCK_KEY.format(device_id=device_id)
+    token = str(uuid.uuid4())
     # cache.add is atomic: returns True only if the key doesn't already exist
-    return cache.add(lock_key, True, timeout=_UPDATE_CONFIG_LOCK_TIMEOUT)
+    if cache.add(lock_key, token, timeout=_UPDATE_CONFIG_LOCK_TIMEOUT):
+        return token
+    return None
 
 
-def _release_update_config_lock(device_id):
-    """Releases the per-device update_config lock."""
+def _release_update_config_lock(device_id, token):
+    """
+    Releases the per-device update_config lock only if the caller
+    owns it (i.e. the stored token matches).
+    """
     lock_key = _UPDATE_CONFIG_LOCK_KEY.format(device_id=device_id)
-    cache.delete(lock_key)
+    stored_token = cache.get(lock_key)
+    if stored_token == token:
+        cache.delete(lock_key)
 
 
 @shared_task
@@ -56,7 +67,8 @@ def update_config(device_id):
     except ObjectDoesNotExist as e:
         logger.warning(f'update_config("{device_id}") failed: {e}')
         return
-    if not _acquire_update_config_lock(device_id):
+    lock_token = _acquire_update_config_lock(device_id)
+    if not lock_token:
         logger.info(
             f"update_config for device {device_id} is already in progress, skipping"
         )
@@ -70,7 +82,7 @@ def update_config(device_id):
             logger.info(f"Updating {device} (pk: {device_id})")
             device_conn.update_config()
     finally:
-        _release_update_config_lock(device_id)
+        _release_update_config_lock(device_id, lock_token)
 
 
 # task timeout is SSH_COMMAND_TIMEOUT plus a 20% margin

openwisp_controller/connection/tests/test_models.py

@@ -1036,7 +1036,7 @@ def test_device_update_config_in_progress(
 
         with mock.patch(
             "openwisp_controller.connection.tasks._acquire_update_config_lock",
-            return_value=False,
+            return_value=None,
         ):
             conf.config = {"general": {"timezone": "UTC"}}
             conf.full_clean()
@@ -1057,28 +1057,43 @@ def test_device_update_config_not_in_progress(
 
         with mock.patch(
             "openwisp_controller.connection.tasks._acquire_update_config_lock",
-            return_value=True,
+            return_value="fake-lock-token",
         ), mock.patch(
             "openwisp_controller.connection.tasks._release_update_config_lock",
-        ):
+        ) as mocked_release:
             conf.config = {"general": {"timezone": "UTC"}}
             conf.full_clean()
             conf.save()
             mocked_get_working_connection.assert_called_once()
             mocked_update_config.assert_called_once()
+            mocked_release.assert_called_once()
 
     def test_acquire_update_config_lock(self):
         """Test that the lock can be acquired and prevents duplicate acquisition."""
         device_id = "test-device-id"
-        # First acquisition should succeed
-        self.assertTrue(_acquire_update_config_lock(device_id))
+        # First acquisition should succeed and return a token
+        token = _acquire_update_config_lock(device_id)
+        self.assertIsNotNone(token)
         # Second acquisition should fail (lock already held)
-        self.assertFalse(_acquire_update_config_lock(device_id))
-        # After releasing, acquisition should succeed again
-        _release_update_config_lock(device_id)
-        self.assertTrue(_acquire_update_config_lock(device_id))
+        self.assertIsNone(_acquire_update_config_lock(device_id))
+        # After releasing with correct token, acquisition should succeed again
+        _release_update_config_lock(device_id, token)
+        token2 = _acquire_update_config_lock(device_id)
+        self.assertIsNotNone(token2)
         # Cleanup
-        _release_update_config_lock(device_id)
+        _release_update_config_lock(device_id, token2)
+
+    def test_release_update_config_lock_wrong_token(self):
+        """Only the lock owner can release the lock."""
+        device_id = "test-device-id"
+        token = _acquire_update_config_lock(device_id)
+        self.assertIsNotNone(token)
+        # Releasing with wrong token should not delete the lock
+        _release_update_config_lock(device_id, "wrong-token")
+        # Lock should still be held
+        self.assertIsNone(_acquire_update_config_lock(device_id))
+        # Releasing with correct token should work
+        _release_update_config_lock(device_id, token)
 
     @mock.patch(_connect_path)
     def test_schedule_command_called(self, connect_mocked):