From 6b66b857f3412f90dde47717af9023502c0b53c5 Mon Sep 17 00:00:00 2001
From: tedbow <tedbow@240860.no-reply.drupal.org>
Date: Mon, 29 Nov 2021 17:31:10 +0000
Subject: [PATCH] Issue #3250386 by tedbow, phenaproxima: Ensure that the stage
 is uniquely idenitfiable across multiple sessions

---
 automatic_updates.routing.yml                 |   2 +-
 package_manager/src/Stage.php                 | 127 ++++++++++++++----
 .../tests/src/Kernel/StageOwnershipTest.php   | 119 +++++++++++++---
 src/BatchProcessor.php                        |  24 ++--
 src/Form/UpdateReady.php                      |  20 ++-
 src/Form/UpdaterForm.php                      |   2 +-
 src/Updater.php                               |  17 ++-
 .../AutomaticUpdatesFunctionalTestBase.php    |   8 ++
 tests/src/Functional/UpdateLockTest.php       |   7 +-
 tests/src/Functional/UpdaterFormTest.php      |   4 +-
 tests/src/Kernel/UpdaterTest.php              |   4 +-
 11 files changed, 266 insertions(+), 68 deletions(-)

diff --git a/automatic_updates.routing.yml b/automatic_updates.routing.yml
index a886f78046..848b3f8ce5 100644
--- a/automatic_updates.routing.yml
+++ b/automatic_updates.routing.yml
@@ -6,7 +6,7 @@ automatic_updates.update_readiness:
   requirements:
     _permission: 'administer software updates'
 automatic_updates.confirmation_page:
-  path: '/admin/automatic-update-ready'
+  path: '/admin/automatic-update-ready/{stage_id}'
   defaults:
     _form: '\Drupal\automatic_updates\Form\UpdateReady'
     _title: 'Ready to update'
diff --git a/package_manager/src/Stage.php b/package_manager/src/Stage.php
index 43fd42f6bd..af15bae087 100644
--- a/package_manager/src/Stage.php
+++ b/package_manager/src/Stage.php
@@ -2,6 +2,7 @@
 
 namespace Drupal\package_manager;
 
+use Drupal\Component\Utility\Crypt;
 use Drupal\Core\TempStore\SharedTempStoreFactory;
 use Drupal\package_manager\Event\PostApplyEvent;
 use Drupal\package_manager\Event\PostCreateEvent;
@@ -25,9 +26,22 @@ use Symfony\Contracts\EventDispatcher\EventDispatcherInterface;
  * directory, use Composer to require packages into it, sync changes from the
  * staging directory back into the active code base, and then delete the
  * staging directory.
+ *
+ * Only one staging area can exist at any given time, and the stage is owned by
+ * the user or session that originally created it. Only the owner can perform
+ * operations on the staging area, and the stage must be "claimed" by its owner
+ * before any such operations are done. A stage is claimed by presenting a
+ * unique token that is generated when the stage is created.
  */
 class Stage {
 
+  /**
+   * The tempstore key under which to store the locking info for this stage.
+   *
+   * @var string
+   */
+  protected const TEMPSTORE_LOCK_KEY = 'lock';
+
   /**
    * The path locator service.
    *
@@ -70,20 +84,21 @@ class Stage {
    */
   protected $eventDispatcher;
 
-
   /**
-   * The tempstore key under which to store the active status of this stage.
+   * The shared temp store.
    *
-   * @var string
+   * @var \Drupal\Core\TempStore\SharedTempStore
    */
-  protected const TEMPSTORE_ACTIVE_KEY = 'active';
+  protected $tempStore;
 
   /**
-   * The shared temp store.
+   * The lock info for the stage.
    *
-   * @var \Drupal\Core\TempStore\SharedTempStore
+   * Consists of a unique random string and the current class name.
+   *
+   * @var string[]
    */
-  protected $tempStore;
+  private $lock;
 
   /**
    * Constructs a new Stage object.
@@ -120,24 +135,24 @@ class Stage {
    *   TRUE if the staging area can be created, otherwise FALSE.
    */
   final public function isAvailable(): bool {
-    return empty($this->tempStore->getMetadata(static::TEMPSTORE_ACTIVE_KEY));
-  }
-
-  /**
-   * Determines if the current user or session is the owner of the staging area.
-   *
-   * @return bool
-   *   TRUE if the current session or user is the owner of the staging area,
-   *   otherwise FALSE.
-   */
-  final public function isOwnedByCurrentUser(): bool {
-    return !empty($this->tempStore->getIfOwner(static::TEMPSTORE_ACTIVE_KEY));
+    return empty($this->tempStore->getMetadata(static::TEMPSTORE_LOCK_KEY));
   }
 
   /**
    * Copies the active code base into the staging area.
+   *
+   * This will automatically claim the stage, so external code does NOT need to
+   * call ::claim(). However, if it was created during another request, the
+   * stage must be claimed before operations can be performed on it.
+   *
+   * @return string
+   *   Unique ID for the stage, which can be used to claim the stage before
+   *   performing other operations on it. Calling code should store this ID for
+   *   as long as the stage needs to exist.
+   *
+   * @see ::claim()
    */
-  public function create(): void {
+  public function create(): string {
     if (!$this->isAvailable()) {
       throw new StageException([], 'Cannot create a new stage because one already exists.');
     }
@@ -147,7 +162,9 @@ class Stage {
     // to create a staging area at around the same time. If an error occurs
     // while the event is being processed, the stage is marked as available.
     // @see ::dispatch()
-    $this->tempStore->set(static::TEMPSTORE_ACTIVE_KEY, TRUE);
+    $id = Crypt::randomBytesBase64();
+    $this->tempStore->set(static::TEMPSTORE_LOCK_KEY, [$id, static::class]);
+    $this->claim($id);
 
     $active_dir = $this->pathLocator->getActiveDirectory();
     $stage_dir = $this->pathLocator->getStageDirectory();
@@ -157,6 +174,7 @@ class Stage {
 
     $this->beginner->begin($active_dir, $stage_dir, $event->getExcludedPaths());
     $this->dispatch(new PostCreateEvent($this));
+    return $id;
   }
 
   /**
@@ -218,11 +236,18 @@ class Stage {
     if (is_dir($stage_dir)) {
       $this->cleaner->clean($stage_dir);
     }
-    // We're all done, so mark the stage as available.
-    $this->tempStore->delete(static::TEMPSTORE_ACTIVE_KEY);
+    $this->markAsAvailable();
     $this->dispatch(new PostDestroyEvent($this));
   }
 
+  /**
+   * Marks the stage as available.
+   */
+  protected function markAsAvailable(): void {
+    $this->tempStore->delete(static::TEMPSTORE_LOCK_KEY);
+    $this->lock = NULL;
+  }
+
   /**
    * Dispatches an event and handles any errors that it collects.
    *
@@ -249,7 +274,7 @@ class Stage {
       // available.
       // @see ::create()
       if ($event instanceof PreCreateEvent) {
-        $this->tempStore->delete(static::TEMPSTORE_ACTIVE_KEY);
+        $this->markAsAvailable();
       }
 
       // Wrap the exception to preserve the backtrace, and re-throw it.
@@ -284,14 +309,64 @@ class Stage {
     return ComposerUtility::createForDirectory($dir);
   }
 
+  /**
+   * Attempts to claim the stage.
+   *
+   * Once a stage has been created, no operations can be performed on it until
+   * it is claimed. This is to ensure that stage operations across multiple
+   * requests are being done by the same code, running under the same user or
+   * session that created the stage in the first place. To claim a stage, the
+   * calling code must provide the unique identifier that was generated when the
+   * stage was created.
+   *
+   * The stage is claimed when it is created, so external code does NOT need to
+   * call this method after calling ::create() in the same request.
+   *
+   * @param string $unique_id
+   *   The unique ID that was returned by ::create().
+   *
+   * @return $this
+   *
+   * @throws \Drupal\package_manager\StageException
+   *   If the stage cannot be claimed. This can happen if the current user or
+   *   session did not originally create the stage, if $unique_id doesn't match
+   *   the unique ID that was generated when the stage was created, or the
+   *   current class is not the same one that was used to create the stage.
+   *
+   * @see ::create()
+   */
+  final public function claim(string $unique_id): self {
+    if ($this->isAvailable()) {
+      throw new StageException([], 'Cannot claim the stage because no stage has been created.');
+    }
+
+    $stored_lock = $this->tempStore->getIfOwner(self::TEMPSTORE_LOCK_KEY);
+    if (!$stored_lock) {
+      throw new StageException([], 'Cannot claim the stage because it is not owned by the current user or session.');
+    }
+
+    if ($stored_lock === [$unique_id, static::class]) {
+      $this->lock = $stored_lock;
+      return $this;
+    }
+    throw new StageException([], 'Cannot claim the stage because the current lock does not match the stored lock.');
+  }
+
   /**
    * Ensures that the current user or session owns the staging area.
    *
+   * @throws \LogicException
+   *   If ::claim() has not been previously called.
    * @throws \Drupal\package_manager\StageException
    *   If the current user or session does not own the staging area.
    */
-  protected function checkOwnership(): void {
-    if (!$this->isOwnedByCurrentUser()) {
+  final protected function checkOwnership(): void {
+    if (empty($this->lock)) {
+      throw new \LogicException('Stage must be claimed before performing any operations on it.');
+    }
+
+    $stored_lock = $this->tempStore->getIfOwner(static::TEMPSTORE_LOCK_KEY);
+    if ($stored_lock !== $this->lock) {
       throw new StageException([], 'Stage is not owned by the current user or session.');
     }
   }
diff --git a/package_manager/tests/src/Kernel/StageOwnershipTest.php b/package_manager/tests/src/Kernel/StageOwnershipTest.php
index b763c6ef2d..9f05fea9f9 100644
--- a/package_manager/tests/src/Kernel/StageOwnershipTest.php
+++ b/package_manager/tests/src/Kernel/StageOwnershipTest.php
@@ -2,7 +2,9 @@
 
 namespace Drupal\Tests\package_manager\Kernel;
 
+use Drupal\package_manager\Event\PreCreateEvent;
 use Drupal\package_manager\StageException;
+use Drupal\package_manager_test_validation\TestSubscriber;
 use Drupal\Tests\user\Traits\UserCreationTrait;
 
 /**
@@ -17,7 +19,11 @@ class StageOwnershipTest extends PackageManagerKernelTestBase {
   /**
    * {@inheritdoc}
    */
-  protected static $modules = ['system', 'user'];
+  protected static $modules = [
+    'system',
+    'user',
+    'package_manager_test_validation',
+  ];
 
   /**
    * {@inheritdoc}
@@ -64,20 +70,13 @@ class StageOwnershipTest extends PackageManagerKernelTestBase {
    *   ownership and status of the other stage.
    */
   private function assertOwnershipIsEnforced(TestStage $will_create, TestStage $never_create): void {
-    // Before the staging area is created, isOwnedByCurrentUser() should return
-    // FALSE and isAvailable() should return TRUE.
-    $this->assertFalse($will_create->isOwnedByCurrentUser());
-    $this->assertFalse($never_create->isOwnedByCurrentUser());
+    // Before the staging area is created, isAvailable() should return TRUE.
     $this->assertTrue($will_create->isAvailable());
     $this->assertTrue($never_create->isAvailable());
 
-    $will_create->create();
-    // Only the staging area that was actually created should be owned by the
-    // current user...
-    $this->assertTrue($will_create->isOwnedByCurrentUser());
-    $this->assertFalse($never_create->isOwnedByCurrentUser());
-    // ...but both staging areas should be considered unavailable (i.e., cannot
-    // be created until the existing one is destroyed first).
+    $stage_id = $will_create->create();
+    // Both staging areas should be considered unavailable (i.e., cannot be
+    // created until the existing one is destroyed first).
     $this->assertFalse($will_create->isAvailable());
     $this->assertFalse($never_create->isAvailable());
 
@@ -93,6 +92,13 @@ class StageOwnershipTest extends PackageManagerKernelTestBase {
       }
     }
 
+    try {
+      $never_create->claim($stage_id);
+    }
+    catch (StageException $exception) {
+      $this->assertSame('Cannot claim the stage because it is not owned by the current user or session.', $exception->getMessage());
+    }
+
     // Only the stage's owner should be able to move it through its life cycle.
     $callbacks = [
       'require' => [
@@ -106,14 +112,95 @@ class StageOwnershipTest extends PackageManagerKernelTestBase {
         $never_create->$method(...$arguments);
         $this->fail("Able to call '$method' on a stage that was never created.");
       }
-      catch (StageException $exception) {
-        $this->assertSame('Stage is not owned by the current user or session.', $exception->getMessage());
+      catch (\LogicException $exception) {
+        $this->assertSame('Stage must be claimed before performing any operations on it.', $exception->getMessage());
       }
       // The call should succeed on the created stage.
       $will_create->$method(...$arguments);
     }
   }
 
+  /**
+   * Tests behavior of claiming a stage.
+   */
+  public function testClaim(): void {
+    // Log in as a user so that any stage instances created during the session
+    // should be able to successfully call ::claim().
+    $user_2 = $this->createUser([], NULL, FALSE, ['uid' => 2]);
+    $this->setCurrentUser($user_2);
+    $creator_stage = $this->createStage();
+
+    // Ensure that exceptions thrown during ::create() will not lock the stage.
+    $error = new \Exception('I am going to stop stage creation.');
+    TestSubscriber::setException($error, PreCreateEvent::class);
+    try {
+      $creator_stage->create();
+      $this->fail('Was able to create the stage despite throwing an exception in pre-create.');
+    }
+    catch (\RuntimeException $exception) {
+      $this->assertSame($error->getMessage(), $exception->getMessage());
+    }
+
+    // The stage should be available, and throw if we try to claim it.
+    $this->assertTrue($creator_stage->isAvailable());
+    try {
+      $creator_stage->claim('any-id-would-fail');
+      $this->fail('Was able to claim a stage that has not been created.');
+    }
+    catch (StageException $exception) {
+      $this->assertSame('Cannot claim the stage because no stage has been created.', $exception->getMessage());
+    }
+    TestSubscriber::setException(NULL, PreCreateEvent::class);
+
+    // Even if we own the stage, we should not be able to claim it with an
+    // incorrect ID.
+    $stage_id = $creator_stage->create();
+    try {
+      $this->createStage()->claim('not-correct-id');
+      $this->fail('Was able to claim an owned stage with an incorrect ID.');
+    }
+    catch (StageException $exception) {
+      $this->assertSame('Cannot claim the stage because the current lock does not match the stored lock.', $exception->getMessage());
+    }
+
+    // A stage that is successfully claimed should be able to call any method
+    // for its life cycle.
+    $callbacks = [
+      'require' => [
+        ['vendor/lib:0.0.1'],
+      ],
+      'apply' => [],
+      'destroy' => [],
+    ];
+    foreach ($callbacks as $method => $arguments) {
+      // Create a new stage instance for each method.
+      $this->createStage()->claim($stage_id)->$method(...$arguments);
+    }
+
+    // The stage cannot be claimed after it's been destroyed.
+    try {
+      $this->createStage()->claim($stage_id);
+      $this->fail('Was able to claim an owned stage after it was destroyed.');
+    }
+    catch (StageException $exception) {
+      $this->assertSame('Cannot claim the stage because no stage has been created.', $exception->getMessage());
+    }
+
+    // Create a new stage and then log in as a different user.
+    $new_stage_id = $this->createStage()->create();
+    $user_3 = $this->createUser([], NULL, FALSE, ['uid' => 3]);
+    $this->setCurrentUser($user_3);
+
+    // Even if they use the correct stage ID, the current user cannot claim a
+    // stage they didn't create.
+    try {
+      $this->createStage()->claim($new_stage_id);
+    }
+    catch (StageException $exception) {
+      $this->assertSame('Cannot claim the stage because it is not owned by the current user or session.', $exception->getMessage());
+    }
+  }
+
   /**
    * Tests a stage being destroyed by a user who doesn't own it.
    */
@@ -126,8 +213,8 @@ class StageOwnershipTest extends PackageManagerKernelTestBase {
       $not_owned->destroy();
       $this->fail("Able to destroy a stage that we don't own.");
     }
-    catch (StageException $exception) {
-      $this->assertSame('Stage is not owned by the current user or session.', $exception->getMessage());
+    catch (\LogicException $exception) {
+      $this->assertSame('Stage must be claimed before performing any operations on it.', $exception->getMessage());
     }
     // We should be able to destroy the stage if we ignore ownership.
     $not_owned->destroy(TRUE);
diff --git a/src/BatchProcessor.php b/src/BatchProcessor.php
index fbf90ab015..0b68b0a6b1 100644
--- a/src/BatchProcessor.php
+++ b/src/BatchProcessor.php
@@ -66,7 +66,8 @@ class BatchProcessor {
    */
   public static function begin(array $project_versions, array &$context): void {
     try {
-      static::getUpdater()->begin($project_versions);
+      $stage_unique = static::getUpdater()->begin($project_versions);
+      $context['results']['stage_id'] = $stage_unique;
     }
     catch (\Throwable $e) {
       static::handleException($e, $context);
@@ -83,7 +84,8 @@ class BatchProcessor {
    */
   public static function stage(array &$context): void {
     try {
-      static::getUpdater()->stage();
+      $stage_id = $context['results']['stage_id'];
+      static::getUpdater()->claim($stage_id)->stage();
     }
     catch (\Throwable $e) {
       static::handleException($e, $context);
@@ -93,14 +95,16 @@ class BatchProcessor {
   /**
    * Calls the updater's commit() method.
    *
+   * @param string $stage_id
+   *   The stage ID.
    * @param array $context
    *   The current context of the batch job.
    *
    * @see \Drupal\automatic_updates\Updater::apply()
    */
-  public static function commit(array &$context): void {
+  public static function commit(string $stage_id, array &$context): void {
     try {
-      static::getUpdater()->apply();
+      static::getUpdater()->claim($stage_id)->apply();
     }
     catch (\Throwable $e) {
       static::handleException($e, $context);
@@ -110,14 +114,16 @@ class BatchProcessor {
   /**
    * Calls the updater's clean() method.
    *
+   * @param string $stage_id
+   *   The stage ID.
    * @param array $context
    *   The current context of the batch job.
    *
    * @see \Drupal\automatic_updates\Updater::clean()
    */
-  public static function clean(array &$context): void {
+  public static function clean(string $stage_id, array &$context): void {
     try {
-      static::getUpdater()->clean();
+      static::getUpdater()->claim($stage_id)->clean();
     }
     catch (\Throwable $e) {
       static::handleException($e, $context);
@@ -136,8 +142,10 @@ class BatchProcessor {
    */
   public static function finishStage(bool $success, array $results, array $operations): ?RedirectResponse {
     if ($success) {
-      return new RedirectResponse(Url::fromRoute('automatic_updates.confirmation_page', [],
-        ['absolute' => TRUE])->toString());
+      $url = Url::fromRoute('automatic_updates.confirmation_page', [
+        'stage_id' => $results['stage_id'],
+      ]);
+      return new RedirectResponse($url->setAbsolute()->toString());
     }
     static::handleBatchError($results);
     return NULL;
diff --git a/src/Form/UpdateReady.php b/src/Form/UpdateReady.php
index ece828ee94..586c1d0e46 100644
--- a/src/Form/UpdateReady.php
+++ b/src/Form/UpdateReady.php
@@ -9,6 +9,7 @@ use Drupal\Core\Form\FormBase;
 use Drupal\Core\Form\FormStateInterface;
 use Drupal\Core\Messenger\MessengerInterface;
 use Drupal\Core\State\StateInterface;
+use Drupal\package_manager\StageException;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 
 /**
@@ -70,12 +71,20 @@ class UpdateReady extends FormBase {
   /**
    * {@inheritdoc}
    */
-  public function buildForm(array $form, FormStateInterface $form_state) {
-    if (!$this->updater->isOwnedByCurrentUser()) {
-      $this->messenger->addError('Cannot continue the update because another Composer operation is currently in progress.');
+  public function buildForm(array $form, FormStateInterface $form_state, string $stage_id = NULL) {
+    try {
+      $this->updater->claim($stage_id);
+    }
+    catch (StageException $e) {
+      $this->messenger()->addError($this->t('Cannot continue the update because another Composer operation is currently in progress.'));
       return $form;
     }
 
+    $form['stage_id'] = [
+      '#type' => 'value',
+      '#value' => $stage_id,
+    ];
+
     $form['backup'] = [
       '#prefix' => '<strong>',
       '#markup' => $this->t('Back up your database and site before you continue. <a href=":backup_url">Learn how</a>.', [':backup_url' => 'https://www.drupal.org/node/22281']),
@@ -108,11 +117,12 @@ class UpdateReady extends FormBase {
       $this->state->set('system.maintenance_mode', TRUE);
       // @todo unset after updater. After db update?
     }
+    $stage_id = $form_state->getValue('stage_id');
     $batch = (new BatchBuilder())
       ->setTitle($this->t('Apply updates'))
       ->setInitMessage($this->t('Preparing to apply updates'))
-      ->addOperation([BatchProcessor::class, 'commit'])
-      ->addOperation([BatchProcessor::class, 'clean'])
+      ->addOperation([BatchProcessor::class, 'commit'], [$stage_id])
+      ->addOperation([BatchProcessor::class, 'clean'], [$stage_id])
       ->setFinishCallback([BatchProcessor::class, 'finishCommit'])
       ->toArray();
 
diff --git a/src/Form/UpdaterForm.php b/src/Form/UpdaterForm.php
index 43ad1a01cd..83131b4e12 100644
--- a/src/Form/UpdaterForm.php
+++ b/src/Form/UpdaterForm.php
@@ -234,7 +234,7 @@ class UpdaterForm extends FormBase {
    * Submit function to delete an existing in-progress update.
    */
   public function deleteExistingUpdate(): void {
-    $this->updater->clean();
+    $this->updater->clean(TRUE);
     $this->messenger()->addMessage($this->t("Staged update deleted"));
   }
 
diff --git a/src/Updater.php b/src/Updater.php
index 4a13714234..4a18cde593 100644
--- a/src/Updater.php
+++ b/src/Updater.php
@@ -47,10 +47,13 @@ class Updater extends Stage {
    * @param string[] $project_versions
    *   The versions of the packages to update to, keyed by package name.
    *
+   * @return string
+   *   The unique ID of the stage.
+   *
    * @throws \InvalidArgumentException
    *   Thrown if no project version for Drupal core is provided.
    */
-  public function begin(array $project_versions): void {
+  public function begin(array $project_versions): string {
     if (count($project_versions) !== 1 || !array_key_exists('drupal', $project_versions)) {
       throw new \InvalidArgumentException("Currently only updates to Drupal core are supported.");
     }
@@ -65,7 +68,7 @@ class Updater extends Stage {
       $package_versions['dev'][$package] = $project_versions['drupal'];
     }
     $this->state->set(static::PACKAGES_KEY, $package_versions);
-    $this->create();
+    return $this->create();
   }
 
   /**
@@ -87,6 +90,8 @@ class Updater extends Stage {
    * Stages the update.
    */
   public function stage(): void {
+    $this->checkOwnership();
+
     // Convert an associative array of package versions, keyed by name, to
     // command-line arguments in the form `vendor/name:version`.
     $map = function (array $versions): array {
@@ -107,9 +112,13 @@ class Updater extends Stage {
 
   /**
    * Cleans the current update.
+   *
+   * @param bool $force
+   *   (optional) If TRUE, the staging area will be destroyed even if it is not
+   *   owned by the current user or session. Defaults to FALSE.
    */
-  public function clean(): void {
-    $this->destroy();
+  public function clean(bool $force = FALSE): void {
+    $this->destroy($force);
     $this->state->delete(static::PACKAGES_KEY);
   }
 
diff --git a/tests/src/Functional/AutomaticUpdatesFunctionalTestBase.php b/tests/src/Functional/AutomaticUpdatesFunctionalTestBase.php
index bd06e15eb6..f6e07aaa11 100644
--- a/tests/src/Functional/AutomaticUpdatesFunctionalTestBase.php
+++ b/tests/src/Functional/AutomaticUpdatesFunctionalTestBase.php
@@ -81,4 +81,12 @@ abstract class AutomaticUpdatesFunctionalTestBase extends BrowserTestBase {
     $this->checkForMetaRefresh();
   }
 
+  /**
+   * Asserts that we are on the "update ready" form.
+   */
+  protected function assertUpdateReady(): void {
+    $this->assertSession()
+      ->addressMatches('/\/admin\/automatic-update-ready\/[a-zA-Z0-9_\-]+$/');
+  }
+
 }
diff --git a/tests/src/Functional/UpdateLockTest.php b/tests/src/Functional/UpdateLockTest.php
index b5411aacaa..7bf9378797 100644
--- a/tests/src/Functional/UpdateLockTest.php
+++ b/tests/src/Functional/UpdateLockTest.php
@@ -51,8 +51,9 @@ class UpdateLockTest extends AutomaticUpdatesFunctionalTestBase {
     $this->drupalGet('/admin/modules/automatic-update');
     $page->pressButton('Update');
     $this->checkForMetaRefresh();
+    $this->assertUpdateReady();
     $assert_session->buttonExists('Continue');
-    $assert_session->addressEquals('/admin/automatic-update-ready');
+    $url = parse_url($this->getSession()->getCurrentUrl(), PHP_URL_PATH);
 
     // Another user cannot show up and try to start an update, since the other
     // user already started one.
@@ -63,13 +64,13 @@ class UpdateLockTest extends AutomaticUpdatesFunctionalTestBase {
 
     // If the current user did not start the update, they should not be able to
     // continue it, either.
-    $this->drupalGet('/admin/automatic-update-ready');
+    $this->drupalGet($url);
     $assert_session->pageTextContains('Cannot continue the update because another Composer operation is currently in progress.');
     $assert_session->buttonNotExists('Continue');
 
     // The user who started the update should be able to continue it.
     $this->drupalLogin($user_1);
-    $this->drupalGet('/admin/automatic-update-ready');
+    $this->drupalGet($url);
     $assert_session->pageTextNotContains('Cannot continue the update because another Composer operation is currently in progress.');
     $assert_session->buttonExists('Continue');
   }
diff --git a/tests/src/Functional/UpdaterFormTest.php b/tests/src/Functional/UpdaterFormTest.php
index 92371ae2aa..4ce78c8994 100644
--- a/tests/src/Functional/UpdaterFormTest.php
+++ b/tests/src/Functional/UpdaterFormTest.php
@@ -247,7 +247,7 @@ class UpdaterFormTest extends AutomaticUpdatesFunctionalTestBase {
     $this->assertUpdateStagedTimes(1);
 
     // Confirm we are on the confirmation page.
-    $assert_session->addressEquals('/admin/automatic-update-ready');
+    $this->assertUpdateReady();
     $assert_session->buttonExists('Continue');
 
     // Return to the start page.
@@ -264,7 +264,7 @@ class UpdaterFormTest extends AutomaticUpdatesFunctionalTestBase {
     $this->checkForMetaRefresh();
 
     // Confirm we are on the confirmation page.
-    $assert_session->addressEquals('/admin/automatic-update-ready');
+    $this->assertUpdateReady();
     $this->assertUpdateStagedTimes(2);
     $assert_session->buttonExists('Continue');
   }
diff --git a/tests/src/Kernel/UpdaterTest.php b/tests/src/Kernel/UpdaterTest.php
index 86d6bc494d..4d2a172f80 100644
--- a/tests/src/Kernel/UpdaterTest.php
+++ b/tests/src/Kernel/UpdaterTest.php
@@ -55,7 +55,7 @@ class UpdaterTest extends AutomaticUpdatesKernelTestBase {
     $locator->getStageDirectory()->willReturn('/tmp');
     $this->container->set('package_manager.path_locator', $locator->reveal());
 
-    $this->container->get('automatic_updates.updater')->begin([
+    $id = $this->container->get('automatic_updates.updater')->begin([
       'drupal' => '9.8.1',
     ]);
     // Rebuild the container to ensure the project versions are kept in state.
@@ -85,7 +85,7 @@ class UpdaterTest extends AutomaticUpdatesKernelTestBase {
       '--update-with-all-dependencies',
       '--dev',
     ];
-    $this->container->get('automatic_updates.updater')->stage();
+    $this->container->get('automatic_updates.updater')->claim($id)->stage();
 
     /** @var \Drupal\package_manager_bypass\InvocationRecorderBase $stager */
     $stager = $this->container->get('package_manager.stager');
-- 
GitLab