feat: wire-in key id on registration and key rotation flows

src/Api/Auth/Controllers/AccountsController.cs

@@ -526,7 +526,15 @@ public async Task<KeysResponseModel> PostKeys([FromBody] KeysRequestModel model)
             {
                 throw new BadRequestException("AccountKeys are only supported for V2 encryption.");
             }
-            await _userRepository.SetV2AccountCryptographicStateAsync(user.Id, accountKeysData);
+            // A client that predates the key id field sends none. The account then picks one up from
+            // the backfill endpoint on a later sync rather than here.
+            var userKeyId = KeyId.FromHexEncodedString(model.UserKeyId);
+            var updateUserDataTasks = userKeyId == null
+                ? null
+                : new UpdateUserData[] { _userRepository.SetUserKeyId(user.Id, userKeyId) };
+
+            await _userRepository.SetV2AccountCryptographicStateAsync(user.Id, accountKeysData,
+                updateUserDataTasks);
             return new KeysResponseModel(accountKeysData, user.Key);
         }
         else

src/Api/KeyManagement/Controllers/AccountsKeyManagementController.cs

@@ -51,6 +51,7 @@ private readonly IRotationValidator<IEnumerable<WebAuthnLoginRotateKeyRequestMod
     private readonly IKeyRotationDataQuery _keyRotationDataQuery;
     private readonly ISetKeyConnectorKeyCommand _setKeyConnectorKeyCommand;
     private readonly IConvertUserToKeyConnectorCommand _convertUserToKeyConnectorCommand;
+    private readonly ISetUserKeyIdCommand _setUserKeyIdCommand;
 
     public AccountsKeyManagementController(IUserService userService,
         IOrganizationUserRepository organizationUserRepository,
@@ -70,7 +71,8 @@ public AccountsKeyManagementController(IUserService userService,
             webAuthnKeyValidator,
         IRotationValidator<IEnumerable<OtherDeviceKeysUpdateRequestModel>, IEnumerable<Device>> deviceValidator,
         ISetKeyConnectorKeyCommand setKeyConnectorKeyCommand,
-        IConvertUserToKeyConnectorCommand convertUserToKeyConnectorCommand)
+        IConvertUserToKeyConnectorCommand convertUserToKeyConnectorCommand,
+        ISetUserKeyIdCommand setUserKeyIdCommand)
     {
         _userService = userService;
         _regenerateUserAsymmetricKeysCommand = regenerateUserAsymmetricKeysCommand;
@@ -88,6 +90,21 @@ public AccountsKeyManagementController(IUserService userService,
         _keyRotationDataQuery = keyRotationDataQuery;
         _setKeyConnectorKeyCommand = setKeyConnectorKeyCommand;
         _convertUserToKeyConnectorCommand = convertUserToKeyConnectorCommand;
+        _setUserKeyIdCommand = setUserKeyIdCommand;
+    }
+
+    /// <summary>
+    /// Reports the key id of the caller's current user key to the server.
+    /// </summary>
+    /// <remarks>
+    /// This is meant for backfilling the user-key id for existing users for whom
+    /// the key id is not yet recorded.
+    /// </remarks>
+    [HttpPost("key-management/user-key-id")]
+    public async Task PostUserKeyIdAsync([FromBody] SetUserKeyIdRequestModel request)
+    {
+        var user = await _userService.GetUserByPrincipalAsync(User) ?? throw new UnauthorizedAccessException();
+        await _setUserKeyIdCommand.SetUserKeyIdAsync(user, request.ToKeyId());
     }
 
     [HttpPost("key-management/regenerate-keys")]
@@ -281,7 +298,7 @@ await _organizationUserValidator.ValidateAsync(user,
             Ciphers = await _cipherValidator.ValidateAsync(user, request.AccountData.Ciphers),
             Folders = await _folderValidator.ValidateAsync(user, request.AccountData.Folders),
             Sends = await _sendValidator.ValidateAsync(user, request.AccountData.Sends),
-            NewUserKeyId = request.NewUserKeyId != null ? KeyId.FromHexEncodedString(request.NewUserKeyId) : null
+            NewUserKeyId = KeyId.FromHexEncodedString(request.NewUserKeyId)
         };
     }
 }

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

@@ -30,6 +30,14 @@ public class SetKeyConnectorKeyRequestModel : IValidatableObject
     public string? KeyConnectorKeyWrappedUserKey { get; set; }
     public AccountKeysRequestModel? AccountKeys { get; set; }
 
+    /// <summary>
+    /// Key id of the user key wrapped by <see cref="KeyConnectorKeyWrappedUserKey"/>, when the client
+    /// supplied it. Absent for clients that predate the field, so it is deliberately not part of
+    /// <see cref="IsV2Request"/>.
+    /// </summary>
+    [KeyId]
+    public string? ContainedKeyId { get; init; }
+
     [Required]
     public required string OrgIdentifier { get; init; }
 
@@ -106,7 +114,8 @@ public KeyConnectorKeysData ToKeyConnectorKeysData()
         {
             KeyConnectorKeyWrappedUserKey = KeyConnectorKeyWrappedUserKey,
             AccountKeys = AccountKeys,
-            OrgIdentifier = OrgIdentifier
+            OrgIdentifier = OrgIdentifier,
+            ContainedKeyId = KeyId.FromHexEncodedString(ContainedKeyId)
         };
     }
 }

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

@@ -0,0 +1,17 @@
+using System.ComponentModel.DataAnnotations;
+using Bit.Core.KeyManagement.Models.Data;
+using Bit.Core.Utilities;
+
+namespace Bit.Api.KeyManagement.Models.Requests;
+
+public class SetUserKeyIdRequestModel
+{
+    /// <summary>
+    /// Hex-encoded key id of the user's current user key.
+    /// </summary>
+    [Required(AllowEmptyStrings = false)]
+    [KeyId]
+    public required string UserKeyId { get; init; }
+
+    public KeyId ToKeyId() => KeyId.FromHexEncodedString(UserKeyId)!;
+}

src/Core/Auth/Models/Api/Request/Accounts/KeysRequestModel.cs

 public async Task<KeysResponseModel> PostKeys([FromBody] KeysRequestModel model) 
 public async Task<KeysResponseModel> PostKeys([FromBody] KeysRequestModel model) 
@@ -18,6 +18,13 @@ public class KeysRequestModel
     public string EncryptedPrivateKey { get; set; }
     public AccountKeysRequestModel AccountKeys { get; set; }
 
+    /// <summary>
+    /// Key id of the user key these account keys belong to, when the client supplied it. Absent for
+    /// clients that predate the field. Only honored on the V2 path.
+    /// </summary>
+    [KeyId]
+    public string UserKeyId { get; set; }
+
     [Obsolete("Use SetAccountKeysForUserCommand instead")]
     public User ToUser(User existingUser)
     {

src/Core/Auth/Models/Api/Request/Accounts/RegisterFinishRequestModel.cs

@@ -89,6 +89,8 @@ public User ToUser(bool isV2Encryption)
                 KdfParallelism = MasterPasswordUnlock?.Kdf.Parallelism ?? KdfParallelism,
                 MasterPasswordSalt = MasterPasswordUnlock?.Salt,
                 Key = MasterPasswordUnlock?.MasterKeyWrappedUserKey ?? UserSymmetricKey
+                // Note: V1 register flows do not set the UserKeyId; those accounts report it later
+                // through the backfill endpoint.
             };
 
             user = UserAsymmetricKeys?.ToUser(user)!;
@@ -132,6 +134,7 @@ public RegisterFinishData ToData()
             MasterKeyWrappedUserKey = unlockData?.MasterKeyWrappedUserKey ?? UserSymmetricKey ?? throw new BadRequestException("MasterKeyWrappedUserKey couldn't be found on either the MasterPasswordUnlockData or the UserSymmetricKey property passed in."),
             MasterPasswordAuthenticationHash = authenticationData?.MasterPasswordAuthenticationHash ?? MasterPasswordHash ?? throw new BadRequestException("MasterPasswordHash couldn't be found on either the MasterPasswordAuthenticationData or the MasterPasswordHash property passed in."),
             Salt = unlockData?.Salt,
+            UserKeyId = unlockData?.ContainedKeyId,
         };
     }
 

src/Core/Auth/UserFeatures/UserMasterPassword/FinishSsoJitProvisionMasterPasswordCommand.cs

@@ -62,13 +62,21 @@ public async Task FinishProvisionAsync(User user,
             throw new BadRequestException("User not found within organization.");
         }
 
-        var updateUserData =
+        var updateUserDataTasks = new List<UpdateUserData>
+        {
             _masterPasswordService.BuildUpdateUserDelegateSetInitialMasterPassword(
                 user,
-                masterPasswordDataModel.ToSetInitialPasswordData());
+                masterPasswordDataModel.ToSetInitialPasswordData())
+        };
+
+        var containedKeyId = masterPasswordDataModel.MasterPasswordUnlock.ContainedKeyId;
+        if (containedKeyId is not null)
+        {
+            updateUserDataTasks.Add(_userRepository.SetUserKeyId(user.Id, containedKeyId));
+        }
 
         await _userRepository.SetV2AccountCryptographicStateAsync(user.Id, masterPasswordDataModel.AccountKeys,
-            [updateUserData]);
+            updateUserDataTasks);
 
         await _eventService.LogUserEventAsync(user.Id, EventType.User_ChangedPassword);
 

src/Core/Entities/User.cs

@@ -114,18 +114,6 @@ public class User : ITableObject<Guid>, IStorableSubscriber, IRevisable, ITwoFac
     public string? V2UpgradeToken { get; set; }
     [MaxLength(256)]
     public string? MasterPasswordSalt { get; set; }
-
-    public KeyId? GetUserKeyId()
-    {
-        // Todo: Database Implementation in follow-up PR
-        return null;
-    }
-
-    public void SetUserKeyId(KeyId keyId)
-    {
-        return; // Todo: Database Implementation in follow-up PR
-    }
-
     public DateTime? LastApiKeyRotationDate { get; set; }
     /// <summary>
     /// A hex-endcoded key-id of the user's current user-key.
@@ -137,8 +125,17 @@ public void SetUserKeyId(KeyId keyId)
     /// A key rotation will set a new key id. Account registrations will carry a key id.
     /// </summary>
     [MaxLength(32)]
+    [KeyId]
     public string? UserKeyId { get; set; }
 
+    public void SetUserKeyId(KeyId? userKeyId)
+    {
+        UserKeyId = userKeyId?.ToString();
+    }
+
+    public KeyId? GetUserKeyId() =>
+        KeyId.FromHexEncodedString(string.IsNullOrEmpty(UserKeyId) ? null : UserKeyId);
+
     public string GetMasterPasswordSalt()
     {
         return MasterPasswordSalt ?? Email.ToLowerInvariant().Trim();

src/Core/KeyManagement/Commands/Interfaces/ISetUserKeyIdCommand.cs

@@ -0,0 +1,22 @@
+using Bit.Core.Entities;
+using Bit.Core.KeyManagement.Models.Data;
+
+namespace Bit.Core.KeyManagement.Commands.Interfaces;
+
+public interface ISetUserKeyIdCommand
+{
+    /// <summary>
+    /// Stores the key id of a user's current user key.
+    /// </summary>
+    /// <remarks>
+    /// This is a backfill primitive for accounts that pre-date the key id being reported alongside
+    /// key material. It therefore only accepts a value when the account does not already have one —
+    /// changing an existing key id must happen through a key rotation.
+    /// </remarks>
+    /// <param name="user">The user whose key id is being recorded.</param>
+    /// <param name="userKeyId">Key id of the user's current user key.</param>
+    /// <exception cref="Bit.Core.Exceptions.BadRequestException">
+    /// Thrown when the account already has a key id.
+    /// </exception>
+    Task SetUserKeyIdAsync(User user, KeyId userKeyId);
+}

src/Core/KeyManagement/Commands/SetKeyConnectorKeyCommand.cs

@@ -46,11 +46,18 @@ public async Task SetKeyConnectorKeyForUserAsync(User user, KeyConnectorKeysData
             throw new BadRequestException("Cannot use Key Connector");
         }
 
-        var setKeyConnectorUserKeyTask =
-            _userRepository.SetKeyConnectorUserKey(user.Id, keyConnectorKeysData.KeyConnectorKeyWrappedUserKey);
+        var updateUserDataTasks = new List<UpdateUserData>
+        {
+            _userRepository.SetKeyConnectorUserKey(user.Id, keyConnectorKeysData.KeyConnectorKeyWrappedUserKey)
+        };
+
+        if (keyConnectorKeysData.ContainedKeyId is not null)
+        {
+            updateUserDataTasks.Add(_userRepository.SetUserKeyId(user.Id, keyConnectorKeysData.ContainedKeyId));
+        }
 
         await _userRepository.SetV2AccountCryptographicStateAsync(user.Id,
-            keyConnectorKeysData.AccountKeys.ToAccountKeysData(), [setKeyConnectorUserKeyTask]);
+            keyConnectorKeysData.AccountKeys.ToAccountKeysData(), updateUserDataTasks);
 
         await _eventService.LogUserEventAsync(user.Id, EventType.User_MigratedKeyToKeyConnector);
 

src/Core/KeyManagement/Commands/SetUserKeyIdCommand.cs

@@ -0,0 +1,28 @@
+using Bit.Core.Entities;
+using Bit.Core.Exceptions;
+using Bit.Core.KeyManagement.Commands.Interfaces;
+using Bit.Core.KeyManagement.Models.Data;
+using Bit.Core.Repositories;
+
+namespace Bit.Core.KeyManagement.Commands;
+
+public class SetUserKeyIdCommand : ISetUserKeyIdCommand
+{
+    private readonly IUserRepository _userRepository;
+
+    public SetUserKeyIdCommand(IUserRepository userRepository)
+    {
+        _userRepository = userRepository;
+    }
+
+    /// <inheritdoc />
+    public async Task SetUserKeyIdAsync(User user, KeyId userKeyId)
+    {
+        if (user.GetUserKeyId() is not null)
+        {
+            throw new BadRequestException("User key id is already set.");
+        }
+
+        await _userRepository.UpdateUserDataAsync([_userRepository.SetUserKeyId(user.Id, userKeyId)]);
+    }
+}

src/Core/KeyManagement/KeyManagementServiceCollectionExtensions.cs

@@ -33,6 +33,7 @@ private static void AddKeyManagementCommands(this IServiceCollection services)
         services.AddScoped<IChangeKdfCommand, ChangeKdfCommand>();
         services.AddScoped<ISetKeyConnectorKeyCommand, SetKeyConnectorKeyCommand>();
         services.AddScoped<IConvertUserToKeyConnectorCommand, ConvertUserToKeyConnectorCommand>();
+        services.AddScoped<ISetUserKeyIdCommand, SetUserKeyIdCommand>();
     }
 
     private static void AddKeyManagementQueries(this IServiceCollection services)

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

@@ -10,5 +10,9 @@ public class KeyConnectorKeysData
 
     public required string OrgIdentifier { get; init; }
 
+    /// <summary>
+    /// Key id of the user key wrapped by <see cref="KeyConnectorKeyWrappedUserKey"/>, when the client
+    /// supplied it.
+    /// </summary>
     public KeyId? ContainedKeyId { get; init; }
 }

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

@@ -9,6 +9,11 @@ public class RegisterFinishData
     public required string MasterPasswordAuthenticationHash { get; init; }
     public string? Salt { get; init; }
 
+    /// <summary>
+    /// Key id of the new account's user key, when the client supplied it.
+    /// </summary>
+    public KeyId? UserKeyId { get; init; }
+
     public bool IsV2Encryption()
     {
         return UserAccountKeysData.IsV2Encryption();
@@ -26,11 +31,12 @@ public override bool Equals(object? obj)
                MasterKeyWrappedUserKey == other.MasterKeyWrappedUserKey &&
                MasterPasswordAuthenticationHash == other.MasterPasswordAuthenticationHash &&
                Salt == other.Salt &&
+               Equals(UserKeyId, other.UserKeyId) &&
                IsV2Encryption() == other.IsV2Encryption();
     }
 
     public override int GetHashCode()
     {
-        return HashCode.Combine(UserAccountKeysData, Kdf, MasterKeyWrappedUserKey, MasterPasswordAuthenticationHash, Salt);
+        return HashCode.Combine(UserAccountKeysData, Kdf, MasterKeyWrappedUserKey, MasterPasswordAuthenticationHash, Salt, UserKeyId);
     }
 }

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

@@ -319,6 +319,7 @@ private async Task<bool> BaseRotateUserAccountKeysAsync(BaseRotateUserAccountKey
         var now = DateTime.UtcNow;
         user.RevisionDate = user.AccountRevisionDate = now;
         user.LastKeyRotationDate = now;
+        user.SetUserKeyId(baseModel.NewUserKeyId);
 
         // 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.

src/Core/Services/Implementations/UserService.cs

@@ -326,8 +326,16 @@ public async Task<IdentityResult> CreateUserAsync(User user, RegisterFinishData
         var result = await CreateAsync(user, registerFinishData.MasterPasswordAuthenticationHash);
         if (result.Succeeded)
         {
-            var setRegisterFinishUserDataTask = _userRepository.UpdateMasterPasswordUnlockData(user.Id, registerFinishData);
-            await _userRepository.SetV2AccountCryptographicStateAsync(user.Id, registerFinishData.UserAccountKeysData, [setRegisterFinishUserDataTask]);
+            var updateUserDataActions = new List<UpdateUserData>
+            {
+                _userRepository.UpdateMasterPasswordUnlockData(user.Id, registerFinishData)
+            };
+            if (registerFinishData.UserKeyId is not null)
+            {
+                updateUserDataActions.Add(_userRepository.SetUserKeyId(user.Id, registerFinishData.UserKeyId));
+            }
+
+            await _userRepository.SetV2AccountCryptographicStateAsync(user.Id, registerFinishData.UserAccountKeysData, updateUserDataActions);
         }
         return result;
     }