[PM-39455] Set Upgrade Tokens for Organizations

src/Api/KeyManagement/Controllers/AccountsKeyManagementController.cs

@@ -145,7 +145,8 @@ await _organizationUserValidator.ValidateAsync(user,
                 WebAuthnKeys =
                     await _webauthnKeyValidator.ValidateAsync(user, model.AccountUnlockData.PasskeyUnlockData),
                 DeviceKeys = await _deviceValidator.ValidateAsync(user, model.AccountUnlockData.DeviceKeyUnlockData),
-                V2UpgradeToken = model.AccountUnlockData.V2UpgradeToken?.ToData(),
+                // A manual key rotation always logs the user out, so no upgrade token is needed.
+                V2UpgradeToken = null,
                 Ciphers = await _cipherValidator.ValidateAsync(user, model.AccountData.Ciphers),
                 Folders = await _folderValidator.ValidateAsync(user, model.AccountData.Folders),
                 Sends = await _sendValidator.ValidateAsync(user, model.AccountData.Sends),

src/Core/AdminConsole/Entities/OrganizationUser.cs

@@ -46,6 +46,12 @@ public class OrganizationUser : ITableObject<Guid>, IExternal, IOrganizationUser
     /// is not enrolled in account recovery.
     /// </summary>
     public string? ResetPasswordKey { get; set; }
+    /// <summary>
+    /// The V2 user key wrapped with the V1 user key, as JSON. Set during a V1 to V2 upgrade rotation.
+    /// An Organization admin reaches the V1 user key through account recovery, so the admin can unwrap
+    /// the V2 user key from this token and update <see cref="ResetPasswordKey"/> to it. NULL at all other times.
+    /// </summary>
+    public string? V2UpgradeToken { get; set; }
     /// <inheritdoc cref="OrganizationUserStatusType"/>
     public OrganizationUserStatusType Status { get; set; }
     /// <summary>

src/Core/KeyManagement/UserKey/Implementations/RotateUserAccountKeysCommand.cs

@@ -90,7 +90,11 @@ public async Task<IdentityResult> PasswordChangeAndRotateUserAccountKeysAsync(Us
         model.ValidateForUser(user);
 
         List<UpdateEncryptedDataForKeyRotation> saveEncryptedDataActions = [];
-        var shouldPersistV2UpgradeToken = await BaseRotateUserAccountKeysAsync(model.BaseData, user, saveEncryptedDataActions);
+
+        // A manual key rotation always logs the user out, so a V2 upgrade token is never needed here.
+        // Discard anything the client submitted, which also clears a token left over from an earlier upgrade.
+        model.BaseData.V2UpgradeToken = null;
+        await BaseRotateUserAccountKeysAsync(model.BaseData, user, saveEncryptedDataActions);
 
         // Delegate the master password mutation (hash, wrapped user key, hint, time markers) to
         // MasterPasswordService.
@@ -113,7 +117,7 @@ public async Task<IdentityResult> PasswordChangeAndRotateUserAccountKeysAsync(Us
 
         await _userRepository.UpdateUserKeyAndEncryptedDataV2Async(user, saveEncryptedDataActions);
 
-        await HandlePushNotificationAsync(shouldPersistV2UpgradeToken, user);
+        await HandlePushNotificationAsync(shouldPersistV2UpgradeToken: false, user);
         return IdentityResult.Success;
     }
 
@@ -334,6 +338,17 @@ private async Task<bool> BaseRotateUserAccountKeysAsync(BaseRotateUserAccountKey
             user.SecurityStamp = Guid.NewGuid().ToString();
         }
 
+        // Each membership enrolled in account recovery gets a copy of the token. Account recovery gives the admin
+        // the V1 user key, so the admin can unwrap the V2 user key from the token. The admin then re-wraps the
+        // account recovery key with it, and the member sees no prompt. Without an account recovery key the token
+        // is of no use.
+        foreach (var organizationUser in baseModel.OrganizationUsers)
+        {
+            organizationUser.V2UpgradeToken = organizationUser.IsEnrolledInAccountRecovery()
+                ? user.V2UpgradeToken
+                : null;
+        }
+
         await UpdateAccountKeysAsync(baseModel, user, saveEncryptedDataActions);
         UpdateBaseUnlockMethods(baseModel, user, saveEncryptedDataActions);
         UpdateUserData(baseModel, user, saveEncryptedDataActions);

src/Infrastructure.EntityFramework/AdminConsole/Repositories/OrganizationUserRepository.cs

@@ -808,6 +808,7 @@ public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation(
                 var updateOrganizationUser =
                     newOrganizationUsers.First(newOrganizationUser => newOrganizationUser.Id == organizationUser.Id);
                 organizationUser.ResetPasswordKey = updateOrganizationUser.ResetPasswordKey;
+                organizationUser.V2UpgradeToken = updateOrganizationUser.V2UpgradeToken;
             }
 
             await dbContext.SaveChangesAsync();

src/Sql/dbo/AdminConsole/Stored Procedures/OrganizationUser_CreateWithCollections.sql

@@ -15,12 +15,13 @@ CREATE PROCEDURE [dbo].[OrganizationUser_CreateWithCollections]
     @AccessSecretsManager BIT = 0,
     @RevocationReason TINYINT = NULL,
     @StatusNew SMALLINT = NULL,
-    @AccessPam BIT = 0
+    @AccessPam BIT = 0,
+    @V2UpgradeToken VARCHAR(MAX) = NULL
 AS
 BEGIN
     SET NOCOUNT ON
 
-    EXEC [dbo].[OrganizationUser_Create] @Id, @OrganizationId, @UserId, @Email, @Key, @Status, @Type, @ExternalId, @CreationDate, @RevisionDate, @Permissions, @ResetPasswordKey, @AccessSecretsManager, @RevocationReason, @StatusNew, @AccessPam
+    EXEC [dbo].[OrganizationUser_Create] @Id, @OrganizationId, @UserId, @Email, @Key, @Status, @Type, @ExternalId, @CreationDate, @RevisionDate, @Permissions, @ResetPasswordKey, @AccessSecretsManager, @RevocationReason, @StatusNew, @AccessPam, @V2UpgradeToken
 
     ;WITH [AvailableCollectionsCTE] AS(
         SELECT

src/Sql/dbo/AdminConsole/Stored Procedures/OrganizationUser_UpdateWithCollections.sql

@@ -15,12 +15,13 @@
     @AccessSecretsManager BIT = 0,
     @RevocationReason TINYINT = NULL,
     @StatusNew SMALLINT = NULL,
-    @AccessPam BIT = 0
+    @AccessPam BIT = 0,
+    @V2UpgradeToken VARCHAR(MAX) = NULL
 AS
 BEGIN
     SET NOCOUNT ON
 
-    EXEC [dbo].[OrganizationUser_Update] @Id, @OrganizationId, @UserId, @Email, @Key, @Status, @Type, @ExternalId, @CreationDate, @RevisionDate, @Permissions, @ResetPasswordKey, @AccessSecretsManager, @RevocationReason, @StatusNew, @AccessPam
+    EXEC [dbo].[OrganizationUser_Update] @Id, @OrganizationId, @UserId, @Email, @Key, @Status, @Type, @ExternalId, @CreationDate, @RevisionDate, @Permissions, @ResetPasswordKey, @AccessSecretsManager, @RevocationReason, @StatusNew, @AccessPam, @V2UpgradeToken
 
     -- Bump RevisionDate on all affected collections
     ;WITH [AffectedCollectionsCTE] AS (

src/Sql/dbo/Stored Procedures/OrganizationUser_Create.sql

@@ -14,7 +14,8 @@
     @AccessSecretsManager BIT = 0,
     @RevocationReason TINYINT = NULL,
     @StatusNew SMALLINT = NULL,
-    @AccessPam BIT = 0
+    @AccessPam BIT = 0,
+    @V2UpgradeToken VARCHAR(MAX) = NULL
 AS
 BEGIN
     SET NOCOUNT ON
@@ -36,7 +37,8 @@ BEGIN
         [AccessSecretsManager],
         [RevocationReason],
         [StatusNew],
-        [AccessPam]
+        [AccessPam],
+        [V2UpgradeToken]
     )
     VALUES
     (
@@ -55,6 +57,7 @@ BEGIN
         @AccessSecretsManager,
         @RevocationReason,
         @StatusNew,
-        @AccessPam
+        @AccessPam,
+        @V2UpgradeToken
     )
 END

src/Sql/dbo/Stored Procedures/OrganizationUser_Update.sql

@@ -14,7 +14,8 @@
     @AccessSecretsManager BIT = 0,
     @RevocationReason TINYINT = NULL,
     @StatusNew SMALLINT = NULL,
-    @AccessPam BIT = 0
+    @AccessPam BIT = 0,
+    @V2UpgradeToken VARCHAR(MAX) = NULL
 AS
 BEGIN
     SET NOCOUNT ON
@@ -36,7 +37,8 @@ BEGIN
         [AccessSecretsManager] = @AccessSecretsManager,
         [RevocationReason] = @RevocationReason,
         [StatusNew] = @StatusNew,
-        [AccessPam] = @AccessPam
+        [AccessPam] = @AccessPam,
+        [V2UpgradeToken] = @V2UpgradeToken
     WHERE
         [Id] = @Id
 

src/Sql/dbo/Stored Procedures/OrganizationUser_UpdateDataForKeyRotation.sql

@@ -8,24 +8,28 @@ BEGIN
     -- Parse the JSON string and insert into a temporary table
     DECLARE @OrganizationUserInput AS TABLE (
         [Id] UNIQUEIDENTIFIER,
-        [ResetPasswordKey] VARCHAR(MAX)
+        [ResetPasswordKey] VARCHAR(MAX),
+        [V2UpgradeToken] VARCHAR(MAX)
     )
 
     INSERT INTO @OrganizationUserInput
     SELECT
         [Id],
-        [ResetPasswordKey]
+        [ResetPasswordKey],
+        [V2UpgradeToken]
     FROM OPENJSON(@OrganizationUserJson)
     WITH (
         [Id] UNIQUEIDENTIFIER '$.Id',
-        [ResetPasswordKey] VARCHAR(MAX) '$.ResetPasswordKey'
+        [ResetPasswordKey] VARCHAR(MAX) '$.ResetPasswordKey',
+        [V2UpgradeToken] VARCHAR(MAX) '$.V2UpgradeToken'
     )
 
     -- Perform the update
     UPDATE
         [dbo].[OrganizationUser]
     SET
-        [ResetPasswordKey] = OUI.[ResetPasswordKey]
+        [ResetPasswordKey] = OUI.[ResetPasswordKey],
+        [V2UpgradeToken] = OUI.[V2UpgradeToken]
     FROM
         [dbo].[OrganizationUser] OU
     INNER JOIN

src/Sql/dbo/Stored Procedures/OrganizationUser_UpdateMany.sql

@@ -23,7 +23,8 @@ BEGIN
         [AccessSecretsManager] BIT,
         [RevocationReason] TINYINT NULL,
         [StatusNew] SMALLINT NULL,
-        [AccessPam] BIT
+        [AccessPam] BIT,
+        [V2UpgradeToken] VARCHAR(MAX) NULL
     )
 
     INSERT INTO @OrganizationUserInput
@@ -43,7 +44,8 @@ BEGIN
         [AccessSecretsManager],
         [RevocationReason],
         [StatusNew],
-        [AccessPam]
+        [AccessPam],
+        [V2UpgradeToken]
     FROM OPENJSON(@jsonData)
     WITH (
         [Id] UNIQUEIDENTIFIER '$.Id',
@@ -61,7 +63,8 @@ BEGIN
         [AccessSecretsManager] BIT '$.AccessSecretsManager',
         [RevocationReason] TINYINT '$.RevocationReason',
         [StatusNew] SMALLINT '$.StatusNew',
-        [AccessPam] BIT '$.AccessPam'
+        [AccessPam] BIT '$.AccessPam',
+        [V2UpgradeToken] VARCHAR(MAX) '$.V2UpgradeToken'
     )
 
     -- Perform the update
@@ -82,7 +85,8 @@ BEGIN
         [AccessSecretsManager] = OUI.[AccessSecretsManager],
         [RevocationReason] = OUI.[RevocationReason],
         [StatusNew] = OUI.[StatusNew],
-        [AccessPam] = ISNULL(OUI.[AccessPam], 0)
+        [AccessPam] = ISNULL(OUI.[AccessPam], 0),
+        [V2UpgradeToken] = OUI.[V2UpgradeToken]
     FROM
         [dbo].[OrganizationUser] OU
     INNER JOIN

src/Sql/dbo/Tables/OrganizationUser.sql

@@ -15,6 +15,7 @@
     [RevocationReason]              TINYINT             NULL,
     [StatusNew]                     SMALLINT            NULL,
     [AccessPam]                     BIT                 NOT NULL CONSTRAINT [DF_OrganizationUser_Pam] DEFAULT (0),
+    [V2UpgradeToken]                VARCHAR (MAX)       NULL,
     CONSTRAINT [PK_OrganizationUser] PRIMARY KEY CLUSTERED ([Id] ASC),
     CONSTRAINT [FK_OrganizationUser_Organization] FOREIGN KEY ([OrganizationId]) REFERENCES [dbo].[Organization] ([Id]) ON DELETE CASCADE,
     CONSTRAINT [FK_OrganizationUser_User] FOREIGN KEY ([UserId]) REFERENCES [dbo].[User] ([Id])

test/Api.IntegrationTest/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs

@@ -1,6 +1,7 @@
 #nullable enable
 using System.Net;
 using System.Text.Json;
+using Bit.Api.AdminConsole.Models.Request.Organizations;
 using Bit.Api.IntegrationTest.Factories;
 using Bit.Api.IntegrationTest.Helpers;
 using Bit.Api.KeyManagement.Models.Requests;
@@ -610,7 +611,7 @@ public async Task RotateUpgradeToV2UserAccountKeysAsync_Success(RotateUserAccoun
 
     [Theory]
     [BitAutoData]
-    public async Task RotateUserAccountKeys_V1Crypto_WithV2UpgradeToken_PersistsToken_AndDoesNotLogout(
+    public async Task RotateUserAccountKeys_V1Crypto_WithV2UpgradeToken_IgnoresToken_AndLogsOut(
         RotateUserAccountKeysAndDataRequestModel request)
     {
         var user = await SetupUserForKeyRotationAsync();
@@ -628,12 +629,16 @@ public async Task RotateUserAccountKeys_V1Crypto_WithV2UpgradeToken_PersistsToke
 
         var userNewState = await _userRepository.GetByEmailAsync(_ownerEmail);
         Assert.NotNull(userNewState);
-        Assert.NotNull(userNewState.V2UpgradeToken);
-        Assert.Contains($"\"WrappedUserKey1\":\"{_mockEncryptedType7String}\"", userNewState.V2UpgradeToken);
-        Assert.Contains($"\"WrappedUserKey2\":\"{_mockEncryptedString}\"", userNewState.V2UpgradeToken);
-        Assert.Equal(user.SecurityStamp, userNewState.SecurityStamp);
+
+        // A manual rotation always logs out, so the submitted token is ignored
+        Assert.Null(userNewState.V2UpgradeToken);
+
+        // Security stamp must change (logout occurred)
+        Assert.NotEqual(user.SecurityStamp, userNewState.SecurityStamp);
+
+        // Standard logout push sent without a reason (full logout, not KeyRotation)
         await _pushNotificationService.Received(1)
-            .PushLogOutAsync(userNewState.Id, false, PushNotificationLogOutReason.KeyRotation);
+            .PushLogOutAsync(userNewState.Id, false, null);
     }
 
     [Theory]
@@ -729,19 +734,18 @@ public async Task RotateUserAccountKeys_WithExistingV2UpgradeToken_WithoutNewTok
 
     [Theory]
     [BitAutoData]
-    public async Task RotateUserAccountKeys_WithExistingV2UpgradeToken_WithNewToken_ReplacesToken_AndDoesNotLogout(
+    public async Task RotateUserAccountKeys_WithExistingV2UpgradeToken_WithNewToken_ClearsToken_AndLogsOut(
         RotateUserAccountKeysAndDataRequestModel request)
     {
         // Arrange
         var user = await SetupUserForKeyRotationAsync();
 
         // Add existing old token to user BEFORE rotation
-        var oldToken = new V2UpgradeTokenData
+        user.V2UpgradeToken = new V2UpgradeTokenData
         {
             WrappedUserKey1 = _mockEncryptedType7String2,
             WrappedUserKey2 = _mockEncryptedType2String2
-        };
-        user.V2UpgradeToken = oldToken.ToJson();
+        }.ToJson();
         await _userRepository.ReplaceAsync(user);
 
         // Setup request WITH new V2UpgradeToken
@@ -761,20 +765,16 @@ public async Task RotateUserAccountKeys_WithExistingV2UpgradeToken_WithNewToken_
         // Assert
         var userNewState = await _userRepository.GetByEmailAsync(_ownerEmail);
         Assert.NotNull(userNewState);
-        Assert.NotNull(userNewState.V2UpgradeToken);
 
-        // Verify new token is present
-        Assert.Contains($"\"WrappedUserKey1\":\"{_mockEncryptedType7String}\"", userNewState.V2UpgradeToken);
-        Assert.Contains($"\"WrappedUserKey2\":\"{_mockEncryptedString}\"", userNewState.V2UpgradeToken);
+        // Neither the old nor the new token survives a manual rotation
+        Assert.Null(userNewState.V2UpgradeToken);
 
-        // Verify old token is NOT present
-        Assert.DoesNotContain(oldToken.WrappedUserKey1, userNewState.V2UpgradeToken);
-        Assert.DoesNotContain(oldToken.WrappedUserKey2, userNewState.V2UpgradeToken);
+        // Security stamp must change (logout occurred)
+        Assert.NotEqual(user.SecurityStamp, userNewState.SecurityStamp);
 
-        // Verify NO logout (SecurityStamp should be the same for key rotation with token)
-        Assert.Equal(user.SecurityStamp, userNewState.SecurityStamp);
+        // Standard logout push sent without a reason (full logout, not KeyRotation)
         await _pushNotificationService.Received(1)
-            .PushLogOutAsync(userNewState.Id, false, PushNotificationLogOutReason.KeyRotation);
+            .PushLogOutAsync(userNewState.Id, false, null);
     }
 
     [Theory]
@@ -901,6 +901,71 @@ public async Task RotateUserKeysAsync_V1ToV2Rotation_Success(RotateUserKeysReque
             signatureKeyPair.VerifyingKey);
     }
 
+    [Theory]
+    [BitAutoData]
+    public async Task RotateUserKeysAsync_V1ToV2Rotation_OrganizationUserEnrolledInAccountRecovery_SetsTokenOnMembership(
+        RotateUserKeysRequestModel request)
+    {
+        // Arrange
+        var (organization, organizationUser) = await OrganizationTestHelpers.SignUpAsync(_factory,
+            PlanType.EnterpriseAnnually, _ownerEmail, passwordManagerSeats: 10,
+            paymentMethod: PaymentMethodType.Card);
+        var user = await SetupUserForKeyRotationAsync();
+
+        organizationUser.ResetPasswordKey = _mockEncryptedString;
+        await _organizationUserRepository.ReplaceAsync(organizationUser);
+
+        SetupMasterPasswordRotateUserAccount(request, user, upgradeToken: true);
+        request.UnlockData.OrganizationAccountRecoveryUnlockData =
+        [
+            new ResetPasswordWithOrgIdRequestModel
+            {
+                OrganizationId = organization.Id,
+                ResetPasswordKey = _mockEncryptedString
+            }
+        ];
+
+        // Act
+        var response = await _client.PostAsJsonAsync("/accounts/key-management/rotate-user-keys", request);
+        response.EnsureSuccessStatusCode();
+
+        // Assert - The admin reaches the same token through the membership
+        var userNewState = await _userRepository.GetByEmailAsync(_ownerEmail);
+        Assert.NotNull(userNewState);
+        Assert.NotNull(userNewState.V2UpgradeToken);
+
+        var membership = await _organizationUserRepository.GetByIdAsync(organizationUser.Id);
+        Assert.NotNull(membership);
+        Assert.Equal(userNewState.V2UpgradeToken, membership.V2UpgradeToken);
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task RotateUserKeysAsync_V1ToV2Rotation_OrganizationUserNotEnrolledInAccountRecovery_DoesNotSetTokenOnMembership(
+        RotateUserKeysRequestModel request)
+    {
+        // Arrange - The member is not enrolled in account recovery, so the row has no account recovery key
+        var (_, organizationUser) = await OrganizationTestHelpers.SignUpAsync(_factory,
+            PlanType.EnterpriseAnnually, _ownerEmail, passwordManagerSeats: 10,
+            paymentMethod: PaymentMethodType.Card);
+        var user = await SetupUserForKeyRotationAsync();
+
+        SetupMasterPasswordRotateUserAccount(request, user, upgradeToken: true);
+
+        // Act
+        var response = await _client.PostAsJsonAsync("/accounts/key-management/rotate-user-keys", request);
+        response.EnsureSuccessStatusCode();
+
+        // Assert - The user keeps the token, but the membership does not get a copy
+        var userNewState = await _userRepository.GetByEmailAsync(_ownerEmail);
+        Assert.NotNull(userNewState);
+        Assert.NotNull(userNewState.V2UpgradeToken);
+
+        var membership = await _organizationUserRepository.GetByIdAsync(organizationUser.Id);
+        Assert.NotNull(membership);
+        Assert.Null(membership.V2UpgradeToken);
+    }
+
     [Theory]
     [BitAutoData]
     public async Task RotateUserKeysAsync_MasterPasswordUnlockDataKeyIdDoesNotMatchNewUserKeyId_BadRequest(