[PM-27280] Support v2 encryption on key-connector signups (#6712)

* account v2 registration for key connector

* use new user repository functions

* test coverage

* integration test coverage

* documentation

* code review

* missing test coverage

* fix failing test

* failing test

* incorrect ticket number

* moved back request model to Api, created dedicated data class in Core

* sql stored procedure type mismatch, simplification

* key connector authorization handler

src/Core/Constants.cs

@@ -212,6 +212,7 @@ public static class FeatureFlagKeys
     public const string ConsolidatedSessionTimeoutComponent = "pm-26056-consolidated-session-timeout-component";
     public const string V2RegistrationTDEJIT = "pm-27279-v2-registration-tde-jit";
     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/Authorization/KeyConnectorAuthorizationHandler.cs

@@ -0,0 +1,52 @@
+using Bit.Core.Context;
+using Bit.Core.Entities;
+using Bit.Core.Enums;
+using Microsoft.AspNetCore.Authorization;
+
+namespace Bit.Core.KeyManagement.Authorization;
+
+public class KeyConnectorAuthorizationHandler : AuthorizationHandler<KeyConnectorOperationsRequirement, User>
+{
+    private readonly ICurrentContext _currentContext;
+
+    public KeyConnectorAuthorizationHandler(ICurrentContext currentContext)
+    {
+        _currentContext = currentContext;
+    }
+
+    protected override Task HandleRequirementAsync(AuthorizationHandlerContext context,
+        KeyConnectorOperationsRequirement requirement,
+        User user)
+    {
+        var authorized = requirement switch
+        {
+            not null when requirement == KeyConnectorOperations.Use => CanUse(user),
+            _ => throw new ArgumentException("Unsupported operation requirement type provided.", nameof(requirement))
+        };
+
+        if (authorized)
+        {
+            context.Succeed(requirement);
+        }
+
+        return Task.CompletedTask;
+    }
+
+    private bool CanUse(User user)
+    {
+        // User cannot use Key Connector if they already use it
+        if (user.UsesKeyConnector)
+        {
+            return false;
+        }
+
+        // User cannot use Key Connector if they are an owner or admin of any organization
+        if (_currentContext.Organizations.Any(u =>
+                u.Type is OrganizationUserType.Owner or OrganizationUserType.Admin))
+        {
+            return false;
+        }
+
+        return true;
+    }
+}

src/Core/KeyManagement/Authorization/KeyConnectorOperations.cs

@@ -0,0 +1,16 @@
+using Microsoft.AspNetCore.Authorization.Infrastructure;
+
+namespace Bit.Core.KeyManagement.Authorization;
+
+public class KeyConnectorOperationsRequirement : OperationAuthorizationRequirement
+{
+    public KeyConnectorOperationsRequirement(string name)
+    {
+        Name = name;
+    }
+}
+
+public static class KeyConnectorOperations
+{
+    public static readonly KeyConnectorOperationsRequirement Use = new(nameof(Use));
+}

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

@@ -0,0 +1,13 @@
+using Bit.Core.Entities;
+using Bit.Core.KeyManagement.Models.Data;
+
+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, KeyConnectorKeysData keyConnectorKeysData);
+}

src/Core/KeyManagement/Commands/SetKeyConnectorKeyCommand.cs

@@ -0,0 +1,60 @@
+using Bit.Core.Context;
+using Bit.Core.Entities;
+using Bit.Core.Enums;
+using Bit.Core.Exceptions;
+using Bit.Core.KeyManagement.Authorization;
+using Bit.Core.KeyManagement.Commands.Interfaces;
+using Bit.Core.KeyManagement.Models.Data;
+using Bit.Core.OrganizationFeatures.OrganizationUsers.Interfaces;
+using Bit.Core.Repositories;
+using Bit.Core.Services;
+using Microsoft.AspNetCore.Authorization;
+
+namespace Bit.Core.KeyManagement.Commands;
+
+public class SetKeyConnectorKeyCommand : ISetKeyConnectorKeyCommand
+{
+    private readonly IAuthorizationService _authorizationService;
+    private readonly ICurrentContext _currentContext;
+    private readonly IEventService _eventService;
+    private readonly IAcceptOrgUserCommand _acceptOrgUserCommand;
+    private readonly IUserService _userService;
+    private readonly IUserRepository _userRepository;
+
+    public SetKeyConnectorKeyCommand(
+        IAuthorizationService authorizationService,
+        ICurrentContext currentContext,
+        IEventService eventService,
+        IAcceptOrgUserCommand acceptOrgUserCommand,
+        IUserService userService,
+        IUserRepository userRepository)
+    {
+        _authorizationService = authorizationService;
+        _currentContext = currentContext;
+        _eventService = eventService;
+        _acceptOrgUserCommand = acceptOrgUserCommand;
+        _userService = userService;
+        _userRepository = userRepository;
+    }
+
+    public async Task SetKeyConnectorKeyForUserAsync(User user, KeyConnectorKeysData keyConnectorKeysData)
+    {
+        var authorizationResult = await _authorizationService.AuthorizeAsync(_currentContext.HttpContext.User, user,
+            KeyConnectorOperations.Use);
+        if (!authorizationResult.Succeeded)
+        {
+            throw new BadRequestException("Cannot use Key Connector");
+        }
+
+        var setKeyConnectorUserKeyTask =
+            _userRepository.SetKeyConnectorUserKey(user.Id, keyConnectorKeysData.KeyConnectorKeyWrappedUserKey);
+
+        await _userRepository.SetV2AccountCryptographicStateAsync(user.Id,
+            keyConnectorKeysData.AccountKeys.ToAccountKeysData(), [setKeyConnectorUserKeyTask]);
+
+        await _eventService.LogUserEventAsync(user.Id, EventType.User_MigratedKeyToKeyConnector);
+
+        await _acceptOrgUserCommand.AcceptOrgUserByOrgSsoIdAsync(keyConnectorKeysData.OrgIdentifier, user,
+            _userService);
+    }
+}

src/Core/KeyManagement/KeyManagementServiceCollectionExtensions.cs

@@ -1,9 +1,11 @@
-using Bit.Core.KeyManagement.Commands;
+using Bit.Core.KeyManagement.Authorization;
+using Bit.Core.KeyManagement.Commands;
 using Bit.Core.KeyManagement.Commands.Interfaces;
 using Bit.Core.KeyManagement.Kdf;
 using Bit.Core.KeyManagement.Kdf.Implementations;
 using Bit.Core.KeyManagement.Queries;
 using Bit.Core.KeyManagement.Queries.Interfaces;
+using Microsoft.AspNetCore.Authorization;
 using Microsoft.Extensions.DependencyInjection;
 
 namespace Bit.Core.KeyManagement;
@@ -12,15 +14,22 @@ public static class KeyManagementServiceCollectionExtensions
 {
     public static void AddKeyManagementServices(this IServiceCollection services)
     {
+        services.AddKeyManagementAuthorizationHandlers();
         services.AddKeyManagementCommands();
         services.AddKeyManagementQueries();
         services.AddSendPasswordServices();
     }
 
+    private static void AddKeyManagementAuthorizationHandlers(this IServiceCollection services)
+    {
+        services.AddScoped<IAuthorizationHandler, KeyConnectorAuthorizationHandler>();
+    }
+
     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)

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

@@ -0,0 +1,12 @@
+using Bit.Core.KeyManagement.Models.Api.Request;
+
+namespace Bit.Core.KeyManagement.Models.Data;
+
+public class KeyConnectorKeysData
+{
+    public required string KeyConnectorKeyWrappedUserKey { get; set; }
+
+    public required AccountKeysRequestModel AccountKeys { get; set; }
+
+    public required string OrgIdentifier { get; init; }
+}

src/Core/Repositories/IUserRepository.cs

@@ -72,6 +72,8 @@ public interface IUserRepository : IRepository<User, Guid>
         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

@@ -621,6 +621,7 @@ public class UserService : UserManager<User>, IUserService
         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);