[PM-31052][PM-32469] Add V2UpgradeToken for key rotation without logout (#6995)

* User V2UpgradeToken for key rotation without logout

* reset old v2 upgrade token on manual key rotation

* sql migration fix

* missing table column

* missing view update

* tests for V2UpgradeToken clearing on manual key rotation

* V2 to V2 rotation causes logout. Updated wrapped key 1 to be a valid V2 encrypted string in tests.

* integration tests failures - increase assert recent for date time type from 2 to 5 seconds (usually for UpdatedAt assertions)

* repository test coverage

* migration script update

* new EF migration scripts

* broken EF migration scripts fixed

* refresh views due to User table alternation

src/Api/KeyManagement/Controllers/AccountsKeyManagementController.cs

@@ -121,6 +121,7 @@ public class AccountsKeyManagementController : Controller
             OrganizationUsers = await _organizationUserValidator.ValidateAsync(user, model.AccountUnlockData.OrganizationAccountRecoveryUnlockData),
             WebAuthnKeys = await _webauthnKeyValidator.ValidateAsync(user, model.AccountUnlockData.PasskeyUnlockData),
             DeviceKeys = await _deviceValidator.ValidateAsync(user, model.AccountUnlockData.DeviceKeyUnlockData),
+            V2UpgradeToken = model.AccountUnlockData.V2UpgradeToken?.ToData(),
 
             Ciphers = await _cipherValidator.ValidateAsync(user, model.AccountData.Ciphers),
             Folders = await _folderValidator.ValidateAsync(user, model.AccountData.Folders),

src/Api/KeyManagement/Models/Requests/UnlockDataRequestModel.cs

@@ -14,4 +14,5 @@ public class UnlockDataRequestModel
     public required IEnumerable<ResetPasswordWithOrgIdRequestModel> OrganizationAccountRecoveryUnlockData { get; set; }
     public required IEnumerable<WebAuthnLoginRotateKeyRequestModel> PasskeyUnlockData { get; set; }
     public required IEnumerable<OtherDeviceKeysUpdateRequestModel> DeviceKeyUnlockData { get; set; }
+    public V2UpgradeTokenRequestModel? V2UpgradeToken { get; set; }
 }

src/Api/KeyManagement/Models/Requests/V2UpgradeTokenRequestModel.cs

@@ -0,0 +1,35 @@
+using System.ComponentModel.DataAnnotations;
+using Bit.Core.KeyManagement.Models.Data;
+using Bit.Core.Utilities;
+
+namespace Bit.Api.KeyManagement.Models.Requests;
+
+/// <summary>
+/// Request model for V2 upgrade token submitted during key rotation.
+/// Contains wrapped user keys allowing clients to unlock after V1→V2 upgrade.
+/// </summary>
+public class V2UpgradeTokenRequestModel
+{
+    /// <summary>
+    /// User Key V2 Wrapped User Key V1.
+    /// </summary>
+    [Required]
+    [EncryptedString]
+    public required string WrappedUserKey1 { get; init; }
+
+    /// <summary>
+    /// User Key V1 Wrapped User Key V2.
+    /// </summary>
+    [Required]
+    [EncryptedString]
+    public required string WrappedUserKey2 { get; init; }
+
+    public V2UpgradeTokenData ToData()
+    {
+        return new V2UpgradeTokenData
+        {
+            WrappedUserKey1 = WrappedUserKey1,
+            WrappedUserKey2 = WrappedUserKey2
+        };
+    }
+}

src/Api/Vault/Models/Response/SyncResponseModel.cs

@@ -87,7 +87,14 @@ public class SyncResponseModel() : ResponseModel("sync")
                     Salt = user.Email.ToLowerInvariant()
                 }
                 : null,
-            WebAuthnPrfOptions = webAuthnPrfOptions.Length > 0 ? webAuthnPrfOptions : null
+            WebAuthnPrfOptions = webAuthnPrfOptions.Length > 0 ? webAuthnPrfOptions : null,
+            V2UpgradeToken = V2UpgradeTokenData.FromJson(user.V2UpgradeToken) is { } data
+                ? new V2UpgradeTokenResponseModel
+                {
+                    WrappedUserKey1 = data.WrappedUserKey1,
+                    WrappedUserKey2 = data.WrappedUserKey2
+                }
+                : null
         };
     }
 

src/Core/Constants.cs

@@ -210,6 +210,7 @@ public static class FeatureFlagKeys
     public const string EnableAccountEncryptionV2KeyConnectorRegistration = "enable-account-encryption-v2-key-connector-registration";
     public const string SdkKeyRotation = "pm-30144-sdk-key-rotation";
     public const string UnlockViaSdk = "unlock-via-sdk";
+    public const string NoLogoutOnKeyUpgradeRotation = "pm-31050-no-logout-key-upgrade-rotation";
     public const string EnableAccountEncryptionV2JitPasswordRegistration = "enable-account-encryption-v2-jit-password-registration";
 
     /* Mobile Team */

src/Core/Entities/User.cs

@@ -105,6 +105,11 @@ public class User : ITableObject<Guid>, IStorableSubscriber, IRevisable, ITwoFac
     public DateTime? LastKeyRotationDate { get; set; }
     public DateTime? LastEmailChangeDate { get; set; }
     public bool VerifyDevices { get; set; } = true;
+    /// <summary>
+    /// V2 upgrade token stored as JSON containing two wrapped user keys.
+    /// Allows clients to unlock vault after V1 to V2 key rotation without logout.
+    /// </summary>
+    public string? V2UpgradeToken { get; set; }
     // PM-28827 Uncomment below line.
     // public string? MasterPasswordSalt { get; set; }
 

src/Core/Enums/PushNotificationLogOutReason.cs

@@ -2,5 +2,6 @@
 
 public enum PushNotificationLogOutReason : byte
 {
-    KdfChange = 0
+    KdfChange = 0,
+    KeyRotation = 1
 }

src/Core/KeyManagement/Models/Api/Response/UserDecryptionResponseModel.cs

@@ -15,4 +15,10 @@ public class UserDecryptionResponseModel
     /// </summary>
     [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
     public WebAuthnPrfDecryptionOption[]? WebAuthnPrfOptions { get; set; }
+
+    /// <summary>
+    /// V2 upgrade token returned when available, allowing unlock after V1→V2 upgrade.
+    /// </summary>
+    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+    public V2UpgradeTokenResponseModel? V2UpgradeToken { get; set; }
 }

src/Core/KeyManagement/Models/Api/Response/V2UpgradeTokenResponseModel.cs

@@ -0,0 +1,7 @@
+namespace Bit.Core.KeyManagement.Models.Api.Response;
+
+public class V2UpgradeTokenResponseModel
+{
+    public required string WrappedUserKey1 { get; set; }
+    public required string WrappedUserKey2 { get; set; }
+}

src/Core/KeyManagement/Models/Data/RotateUserAccountKeysData.cs

@@ -20,6 +20,7 @@ public class RotateUserAccountKeysData
     public required IReadOnlyList<OrganizationUser> OrganizationUsers { get; set; }
     public required IEnumerable<WebAuthnLoginRotateKeyData> WebAuthnKeys { get; set; }
     public required IEnumerable<Device> DeviceKeys { get; set; }
+    public V2UpgradeTokenData? V2UpgradeToken { get; set; }
 
     // User vault data encrypted by the userkey
     public required IEnumerable<Cipher> Ciphers { get; set; }

src/Core/KeyManagement/Models/Data/V2UpgradeTokenData.cs

@@ -0,0 +1,31 @@
+using System.Text.Json;
+
+namespace Bit.Core.KeyManagement.Models.Data;
+
+public class V2UpgradeTokenData
+{
+    public required string WrappedUserKey1 { get; init; }
+    public required string WrappedUserKey2 { get; init; }
+
+    public string ToJson()
+    {
+        return JsonSerializer.Serialize(this);
+    }
+
+    public static V2UpgradeTokenData? FromJson(string? json)
+    {
+        if (string.IsNullOrWhiteSpace(json))
+        {
+            return null;
+        }
+
+        try
+        {
+            return JsonSerializer.Deserialize<V2UpgradeTokenData>(json);
+        }
+        catch (JsonException)
+        {
+            return null;
+        }
+    }
+}

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

@@ -90,7 +90,19 @@ public class RotateUserAccountKeysCommand : IRotateUserAccountKeysCommand
         var now = DateTime.UtcNow;
         user.RevisionDate = user.AccountRevisionDate = now;
         user.LastKeyRotationDate = now;
-        user.SecurityStamp = Guid.NewGuid().ToString();
+
+        // V2UpgradeToken is only valid for V1 users transitioning to V2.
+        // For V2 users the token is semantically invalid — discard it and perform a full logout.
+        var shouldPersistV2UpgradeToken = model.V2UpgradeToken != null && !IsV2EncryptionUserAsync(user);
+        if (shouldPersistV2UpgradeToken)
+        {
+            user.V2UpgradeToken = model.V2UpgradeToken!.ToJson();
+        }
+        else
+        {
+            user.V2UpgradeToken = null;
+            user.SecurityStamp = Guid.NewGuid().ToString();
+        }
 
         List<UpdateEncryptedDataForKeyRotation> saveEncryptedDataActions = [];
 
@@ -99,7 +111,17 @@ public class RotateUserAccountKeysCommand : IRotateUserAccountKeysCommand
         UpdateUserData(model, user, saveEncryptedDataActions);
 
         await _userRepository.UpdateUserKeyAndEncryptedDataV2Async(user, saveEncryptedDataActions);
-        await _pushService.PushLogOutAsync(user.Id);
+
+        if (shouldPersistV2UpgradeToken)
+        {
+            await _pushService.PushLogOutAsync(user.Id,
+                reason: PushNotificationLogOutReason.KeyRotation);
+        }
+        else
+        {
+            await _pushService.PushLogOutAsync(user.Id);
+        }
+
         return IdentityResult.Success;
     }
 

src/Infrastructure.EntityFramework/Repositories/UserRepository.cs

@@ -255,6 +255,12 @@ public class UserRepository : Repository<Core.Entities.User, User, Guid>, IUserR
         userEntity.AccountRevisionDate = user.AccountRevisionDate;
         userEntity.RevisionDate = user.RevisionDate;
 
+        userEntity.SignedPublicKey = user.SignedPublicKey;
+        userEntity.SecurityState = user.SecurityState;
+        userEntity.SecurityVersion = user.SecurityVersion;
+
+        userEntity.V2UpgradeToken = user.V2UpgradeToken;
+
         await dbContext.SaveChangesAsync();
 
         //  Update re-encrypted data

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

@@ -44,7 +44,8 @@
     @VerifyDevices BIT = 1,
     @SecurityState VARCHAR(MAX) = NULL,
     @SecurityVersion INT = NULL,
-    @SignedPublicKey VARCHAR(MAX) = NULL
+    @SignedPublicKey VARCHAR(MAX) = NULL,
+    @V2UpgradeToken VARCHAR(MAX) = NULL
 AS
 BEGIN
     SET NOCOUNT ON
@@ -97,7 +98,8 @@ BEGIN
         [SecurityState],
         [SecurityVersion],
         [SignedPublicKey],
-        [MaxStorageGbIncreased]
+        [MaxStorageGbIncreased],
+        [V2UpgradeToken]
     )
     VALUES
     (
@@ -147,6 +149,7 @@ BEGIN
         @SecurityState,
         @SecurityVersion,
         @SignedPublicKey,
-        @MaxStorageGb
+        @MaxStorageGb,
+        @V2UpgradeToken
     )
 END

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

@@ -44,7 +44,8 @@
     @VerifyDevices BIT = 1,
     @SecurityState VARCHAR(MAX) = NULL,
     @SecurityVersion INT = NULL,
-    @SignedPublicKey VARCHAR(MAX) = NULL
+    @SignedPublicKey VARCHAR(MAX) = NULL,
+    @V2UpgradeToken VARCHAR(MAX) = NULL
 AS
 BEGIN
     SET NOCOUNT ON
@@ -97,7 +98,8 @@ BEGIN
         [SecurityState] = @SecurityState,
         [SecurityVersion] = @SecurityVersion,
         [SignedPublicKey] = @SignedPublicKey,
-        [MaxStorageGbIncreased] = @MaxStorageGb
+        [MaxStorageGbIncreased] = @MaxStorageGb,
+        [V2UpgradeToken] = @V2UpgradeToken
     WHERE
         [Id] = @Id
 END

src/Sql/dbo/Tables/User.sql

@@ -42,10 +42,11 @@
     [LastKeyRotationDate]              DATETIME2 (7)    NULL,
     [LastEmailChangeDate]              DATETIME2 (7)    NULL,
     [VerifyDevices]                    BIT              DEFAULT ((1)) NOT NULL,
-    [SecurityState]                    VARCHAR (MAX)   NULL,
+    [SecurityState]                    VARCHAR (MAX)    NULL,
     [SecurityVersion]                  INT              NULL,
-    [SignedPublicKey]                  VARCHAR (MAX)   NULL,
+    [SignedPublicKey]                  VARCHAR (MAX)    NULL,
     [MaxStorageGbIncreased]            SMALLINT         NULL,
+    [V2UpgradeToken]                   VARCHAR(MAX)     NULL,
     CONSTRAINT [PK_User] PRIMARY KEY CLUSTERED ([Id] ASC)
 );
 

src/Sql/dbo/Views/UserView.sql

@@ -46,6 +46,7 @@ SELECT
     [VerifyDevices],
     [SecurityState],
     [SecurityVersion],
-    [SignedPublicKey]
+    [SignedPublicKey],
+    [V2UpgradeToken]
 FROM
     [dbo].[User]