[PM-27280] Support v2 encryption on key-connector signups

src/Api/KeyManagement/Controllers/AccountsKeyManagementController.cs

@@ -13,6 +13,7 @@
 using Bit.Core.Entities;
 using Bit.Core.Exceptions;
 using Bit.Core.KeyManagement.Commands.Interfaces;
+using Bit.Core.KeyManagement.Models.Api.Request;
 using Bit.Core.KeyManagement.Models.Data;
 using Bit.Core.KeyManagement.Queries.Interfaces;
 using Bit.Core.KeyManagement.UserKey;
@@ -47,6 +48,7 @@ private readonly IRotationValidator<IEnumerable<WebAuthnLoginRotateKeyRequestMod
         _webauthnKeyValidator;
     private readonly IRotationValidator<IEnumerable<OtherDeviceKeysUpdateRequestModel>, IEnumerable<Device>> _deviceValidator;
     private readonly IKeyConnectorConfirmationDetailsQuery _keyConnectorConfirmationDetailsQuery;
+    private readonly ISetKeyConnectorKeyCommand _setKeyConnectorKeyCommand;
 
     public AccountsKeyManagementController(IUserService userService,
         IFeatureService featureService,
@@ -62,8 +64,10 @@ public AccountsKeyManagementController(IUserService userService,
             emergencyAccessValidator,
         IRotationValidator<IEnumerable<ResetPasswordWithOrgIdRequestModel>, IReadOnlyList<OrganizationUser>>
             organizationUserValidator,
-        IRotationValidator<IEnumerable<WebAuthnLoginRotateKeyRequestModel>, IEnumerable<WebAuthnLoginRotateKeyData>> webAuthnKeyValidator,
-        IRotationValidator<IEnumerable<OtherDeviceKeysUpdateRequestModel>, IEnumerable<Device>> deviceValidator)
+        IRotationValidator<IEnumerable<WebAuthnLoginRotateKeyRequestModel>, IEnumerable<WebAuthnLoginRotateKeyData>>
+            webAuthnKeyValidator,
+        IRotationValidator<IEnumerable<OtherDeviceKeysUpdateRequestModel>, IEnumerable<Device>> deviceValidator,
+        ISetKeyConnectorKeyCommand setKeyConnectorKeyCommand)
     {
         _userService = userService;
         _featureService = featureService;
@@ -79,6 +83,7 @@ public AccountsKeyManagementController(IUserService userService,
         _webauthnKeyValidator = webAuthnKeyValidator;
         _deviceValidator = deviceValidator;
         _keyConnectorConfirmationDetailsQuery = keyConnectorConfirmationDetailsQuery;
+        _setKeyConnectorKeyCommand = setKeyConnectorKeyCommand;
     }
 
     [HttpPost("key-management/regenerate-keys")]
@@ -146,18 +151,28 @@ public async Task PostSetKeyConnectorKeyAsync([FromBody] SetKeyConnectorKeyReque
             throw new UnauthorizedAccessException();
         }
 
-        var result = await _userService.SetKeyConnectorKeyAsync(model.ToUser(user), model.Key, model.OrgIdentifier);
-        if (result.Succeeded)
+        if (model.IsV2Request())
         {
-            return;
+            // V2 account registration
+            await _setKeyConnectorKeyCommand.SetKeyConnectorKeyForUserAsync(user, model);
         }
-
-        foreach (var error in result.Errors)
+        else
         {
-            ModelState.AddModelError(string.Empty, error.Description);
+            // V1 account registration
+            // TODO removed with https://bitwarden.atlassian.net/browse/PM-27328
+            var result = await _userService.SetKeyConnectorKeyAsync(model.ToUser(user), model.Key, model.OrgIdentifier);
+            if (result.Succeeded)
+            {
+                return;
+            }
+
+            foreach (var error in result.Errors)
+            {
+                ModelState.AddModelError(string.Empty, error.Description);
+            }
+
+            throw new BadRequestException(ModelState);
         }
-
-        throw new BadRequestException(ModelState);
     }
 
     [HttpPost("convert-to-key-connector")]

src/Core/Constants.cs

@@ -213,6 +213,7 @@ public static class FeatureFlagKeys
     public const string DisableType0Decryption = "pm-25174-disable-type-0-decryption";
     public const string ConsolidatedSessionTimeoutComponent = "pm-26056-consolidated-session-timeout-component";
     public const string DataRecoveryTool = "pm-28813-data-recovery-tool";
+    public const string EnableAccountEncryptionV2KeyConnectorRegistration = "enable-account-encryption-v2-key-connector-registration";
 
     /* Mobile Team */
     public const string AndroidImportLoginsFlow = "import-logins-flow";

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

@@ -0,0 +1,13 @@
+using Bit.Core.Entities;
+using Bit.Core.KeyManagement.Models.Api.Request;
+
+namespace Bit.Core.KeyManagement.Commands.Interfaces;
+
+/// <summary>
+/// Creates the user key and account cryptographic state for a new user registering
+/// with Key Connector SSO configuration.
+/// </summary>
+public interface ISetKeyConnectorKeyCommand
+{
+    Task SetKeyConnectorKeyForUserAsync(User user, SetKeyConnectorKeyRequestModel requestModel);
+}

src/Core/KeyManagement/Commands/SetKeyConnectorKeyCommand.cs

@@ -0,0 +1,55 @@
+using Bit.Core.Entities;
+using Bit.Core.Enums;
+using Bit.Core.Exceptions;
+using Bit.Core.KeyManagement.Commands.Interfaces;
+using Bit.Core.KeyManagement.Models.Api.Request;
+using Bit.Core.KeyManagement.Queries.Interfaces;
+using Bit.Core.OrganizationFeatures.OrganizationUsers.Interfaces;
+using Bit.Core.Repositories;
+using Bit.Core.Services;
+
+namespace Bit.Core.KeyManagement.Commands;
+
+public class SetKeyConnectorKeyCommand : ISetKeyConnectorKeyCommand
+{
+    private readonly ICanUseKeyConnectorQuery _canUseKeyConnectorQuery;
+    private readonly IEventService _eventService;
+    private readonly IAcceptOrgUserCommand _acceptOrgUserCommand;
+    private readonly IUserService _userService;
+    private readonly IUserRepository _userRepository;
+
+    public SetKeyConnectorKeyCommand(
+        ICanUseKeyConnectorQuery canUseKeyConnectorQuery,
+        IEventService eventService,
+        IAcceptOrgUserCommand acceptOrgUserCommand,
+        IUserService userService,
+        IUserRepository userRepository)
+    {
+        _canUseKeyConnectorQuery = canUseKeyConnectorQuery;
+        _eventService = eventService;
+        _acceptOrgUserCommand = acceptOrgUserCommand;
+        _userService = userService;
+        _userRepository = userRepository;
+    }
+
+    public async Task SetKeyConnectorKeyForUserAsync(User user, SetKeyConnectorKeyRequestModel requestModel)
+    {
+        // TODO remove validation with https://bitwarden.atlassian.net/browse/PM-27280
+        if (string.IsNullOrEmpty(requestModel.KeyConnectorKeyWrappedUserKey) || requestModel.AccountKeys == null)
+        {
+            throw new BadRequestException("KeyConnectorKeyWrappedUserKey and AccountKeys must be provided");
+        }
+
+        _canUseKeyConnectorQuery.VerifyCanUseKeyConnector(user);
+
+        var setKeyConnectorUserKeyTask =
+            _userRepository.SetKeyConnectorUserKey(user.Id, requestModel.KeyConnectorKeyWrappedUserKey);
+
+        await _userRepository.SetV2AccountCryptographicStateAsync(user.Id, requestModel.AccountKeys.ToAccountKeysData(),
+            [setKeyConnectorUserKeyTask]);
+
+        await _eventService.LogUserEventAsync(user.Id, EventType.User_MigratedKeyToKeyConnector);
+
+        await _acceptOrgUserCommand.AcceptOrgUserByOrgSsoIdAsync(requestModel.OrgIdentifier, user, _userService);
+    }
+}

src/Core/KeyManagement/KeyManagementServiceCollectionExtensions.cs

@@ -21,11 +21,13 @@ private static void AddKeyManagementCommands(this IServiceCollection services)
     {
         services.AddScoped<IRegenerateUserAsymmetricKeysCommand, RegenerateUserAsymmetricKeysCommand>();
         services.AddScoped<IChangeKdfCommand, ChangeKdfCommand>();
+        services.AddScoped<ISetKeyConnectorKeyCommand, SetKeyConnectorKeyCommand>();
     }
 
     private static void AddKeyManagementQueries(this IServiceCollection services)
     {
         services.AddScoped<IUserAccountKeysQuery, UserAccountKeysQuery>();
         services.AddScoped<IKeyConnectorConfirmationDetailsQuery, KeyConnectorConfirmationDetailsQuery>();
+        services.AddScoped<ICanUseKeyConnectorQuery, CanUseKeyConnectorQuery>();
     }
 }

src/Core/KeyManagement/Models/Api/Request/SetKeyConnectorKeyRequestModel.cs

@@ -0,0 +1,93 @@
+using System.ComponentModel.DataAnnotations;
+using Bit.Core.Auth.Models.Api.Request.Accounts;
+using Bit.Core.Entities;
+using Bit.Core.Enums;
+using Bit.Core.Utilities;
+
+namespace Bit.Core.KeyManagement.Models.Api.Request;
+
+public class SetKeyConnectorKeyRequestModel : IValidatableObject
+{
+    // TODO will be removed with https://bitwarden.atlassian.net/browse/PM-27328
+    [Obsolete("Use KeyConnectorKeyWrappedUserKey instead")]
+    public string? Key { get; set; }
+
+    [Obsolete("Use AccountKeys instead")]
+    public KeysRequestModel? Keys { get; set; }
+    [Obsolete("Not used anymore")]
+    public KdfType? Kdf { get; set; }
+    [Obsolete("Not used anymore")]
+    public int? KdfIterations { get; set; }
+    [Obsolete("Not used anymore")]
+    public int? KdfMemory { get; set; }
+    [Obsolete("Not used anymore")]
+    public int? KdfParallelism { get; set; }
+
+    [EncryptedString]
+    public string? KeyConnectorKeyWrappedUserKey { get; set; }
+    public AccountKeysRequestModel? AccountKeys { get; set; }
+
+    [Required]
+    public required string OrgIdentifier { get; init; }
+
+    public IEnumerable<ValidationResult> Validate(ValidationContext validationContext)
+    {
+        if (IsV2Request())
+        {
+            // V2 registration
+            yield break;
+        }
+
+        // V1 registration
+        // TODO removed with https://bitwarden.atlassian.net/browse/PM-27328
+        if (string.IsNullOrEmpty(Key))
+        {
+            yield return new ValidationResult("Key must be supplied.");
+        }
+
+        if (Keys == null)
+        {
+            yield return new ValidationResult("Keys must be supplied.");
+        }
+
+        if (Kdf == null)
+        {
+            yield return new ValidationResult("Kdf must be supplied.");
+        }
+
+        if (KdfIterations == null)
+        {
+            yield return new ValidationResult("KdfIterations must be supplied.");
+        }
+
+        if (Kdf == KdfType.Argon2id)
+        {
+            if (KdfMemory == null)
+            {
+                yield return new ValidationResult("KdfMemory must be supplied when Kdf is Argon2id.");
+            }
+
+            if (KdfParallelism == null)
+            {
+                yield return new ValidationResult("KdfParallelism must be supplied when Kdf is Argon2id.");
+            }
+        }
+    }
+
+    public bool IsV2Request()
+    {
+        return !string.IsNullOrEmpty(KeyConnectorKeyWrappedUserKey) && AccountKeys != null;
+    }
+
+    // TODO removed with https://bitwarden.atlassian.net/browse/PM-27328
+    public User ToUser(User existingUser)
+    {
+        existingUser.Kdf = Kdf!.Value;
+        existingUser.KdfIterations = KdfIterations!.Value;
+        existingUser.KdfMemory = KdfMemory;
+        existingUser.KdfParallelism = KdfParallelism;
+        existingUser.Key = Key;
+        Keys!.ToUser(existingUser);
+        return existingUser;
+    }
+}

src/Core/KeyManagement/Queries/CanUseKeyConnectorQuery.cs

@@ -0,0 +1,31 @@
+using Bit.Core.Context;
+using Bit.Core.Entities;
+using Bit.Core.Enums;
+using Bit.Core.Exceptions;
+using Bit.Core.KeyManagement.Queries.Interfaces;
+
+namespace Bit.Core.KeyManagement.Queries;
+
+public class CanUseKeyConnectorQuery : ICanUseKeyConnectorQuery
+{
+    private readonly ICurrentContext _currentContext;
+
+    public CanUseKeyConnectorQuery(ICurrentContext currentContext)
+    {
+        _currentContext = currentContext;
+    }
+
+    public void VerifyCanUseKeyConnector(User user)
+    {
+        if (user.UsesKeyConnector)
+        {
+            throw new BadRequestException("Already uses Key Connector.");
+        }
+
+        if (_currentContext.Organizations.Any(u =>
+                u.Type is OrganizationUserType.Owner or OrganizationUserType.Admin))
+        {
+            throw new BadRequestException("Cannot use Key Connector when admin or owner of an organization.");
+        }
+    }
+}

src/Core/KeyManagement/Queries/Interfaces/ICanUseKeyConnectorQuery.cs

@@ -0,0 +1,15 @@
+using Bit.Core.Entities;
+
+namespace Bit.Core.KeyManagement.Queries.Interfaces;
+
+/// <summary>
+/// Query to verify if the user can use the key connector
+/// </summary>
+public interface ICanUseKeyConnectorQuery
+{
+    /// <summary>
+    /// Throws an exception if the user cannot use the key connector
+    /// </summary>
+    /// <param name="user">User to validate</param>
+    void VerifyCanUseKeyConnector(User user);
+}

src/Core/Repositories/IUserRepository.cs

@@ -55,6 +55,8 @@ Task SetV2AccountCryptographicStateAsync(
         UserAccountKeysData accountKeysData,
         IEnumerable<UpdateUserData>? updateUserDataActions = null);
     Task DeleteManyAsync(IEnumerable<User> users);
+
+    UpdateUserData SetKeyConnectorUserKey(Guid userId, string keyConnectorWrappedUserKey);
 }
 
 public delegate Task UpdateUserData(Microsoft.Data.SqlClient.SqlConnection? connection = null,

src/Core/Services/IUserService.cs

@@ -33,6 +33,8 @@ public interface IUserService
     Task<IdentityResult> ChangeEmailAsync(User user, string masterPassword, string newEmail, string newMasterPassword,
         string token, string key);
     Task<IdentityResult> ChangePasswordAsync(User user, string masterPassword, string newMasterPassword, string passwordHint, string key);
+    // TODO removed with https://bitwarden.atlassian.net/browse/PM-27328
+    [Obsolete("Use ISetKeyConnectorKeyCommand instead. This method will be removed in a future version.")]
     Task<IdentityResult> SetKeyConnectorKeyAsync(User user, string key, string orgIdentifier);
     Task<IdentityResult> ConvertToKeyConnectorAsync(User user);
     Task<IdentityResult> AdminResetPasswordAsync(OrganizationUserType type, Guid orgId, Guid id, string newMasterPassword, string key);

src/Core/Services/Implementations/UserService.cs

@@ -174,7 +174,7 @@
             return null;
         }
 
         _currentContext.User = await _userRepository.GetByIdAsync(userIdGuid);
         return _currentContext.User;
     }
 
@@ -185,7 +185,7 @@
             return _currentContext.User;
         }
 
         _currentContext.User = await _userRepository.GetByIdAsync(userId);
         return _currentContext.User;
     }
 
@@ -617,6 +617,7 @@
         return IdentityResult.Failed(_identityErrorDescriber.PasswordMismatch());
     }
 
+    // TODO removed with https://bitwarden.atlassian.net/browse/PM-27328
     public async Task<IdentityResult> SetKeyConnectorKeyAsync(User user, string key, string orgIdentifier)
     {
         var identityResult = CheckCanUseKeyConnector(user);
@@ -914,7 +915,7 @@
         }
         catch when (!_globalSettings.SelfHosted)
         {
             await paymentService.CancelAndRecoverChargesAsync(user);
             throw;
         }