[PM-3797 Part 1] Layout new key rotation methods (#3425)

* layout new key rotation methods
- add endpoint with request model
- add command with data model
- add repository method

* layout new key rotation methods
- add endpoint with request model
- add command with data model
- add repository method

* formatting

* rename account recovery to reset password

* fix tests

* remove extra endpoint

* rename account recovery to reset password

* fix tests and formatting

* register db calls in command, removing list from user repo

* formatting

src/Api/Auth/Models/Request/Accounts/UpdateKeyRequestModel.cs

@@ -1,4 +1,5 @@
 using System.ComponentModel.DataAnnotations;
+using Bit.Api.AdminConsole.Models.Request.Organizations;
 using Bit.Api.Tools.Models.Request;
 using Bit.Api.Vault.Models.Request;
 
@@ -10,12 +11,13 @@ public class UpdateKeyRequestModel
     [StringLength(300)]
     public string MasterPasswordHash { get; set; }
     [Required]
-    public IEnumerable<CipherWithIdRequestModel> Ciphers { get; set; }
+    public string Key { get; set; }
     [Required]
+    public string PrivateKey { get; set; }
+    public IEnumerable<CipherWithIdRequestModel> Ciphers { get; set; }
     public IEnumerable<FolderWithIdRequestModel> Folders { get; set; }
     public IEnumerable<SendWithIdRequestModel> Sends { get; set; }
-    [Required]
+    public IEnumerable<EmergencyAccessUpdateRequestModel> EmergencyAccessKeys { get; set; }
-    public string PrivateKey { get; set; }
+    public IEnumerable<OrganizationUserUpdateRequestModel> ResetPasswordKeys { get; set; }
-    [Required]
+
-    public string Key { get; set; }
 }

src/Api/Controllers/AccountsController.cs

@@ -9,9 +9,12 @@ using Bit.Core.AdminConsole.Enums.Provider;
 using Bit.Core.AdminConsole.Repositories;
 using Bit.Core.Auth.Models.Api.Request.Accounts;
 using Bit.Core.Auth.Models.Api.Response.Accounts;
+using Bit.Core.Auth.Models.Data;
 using Bit.Core.Auth.Services;
+using Bit.Core.Auth.UserFeatures.UserKey;
 using Bit.Core.Auth.UserFeatures.UserMasterPassword.Interfaces;
 using Bit.Core.Auth.Utilities;
+using Bit.Core.Context;
 using Bit.Core.Enums;
 using Bit.Core.Exceptions;
 using Bit.Core.Models.Api.Response;
@@ -27,6 +30,7 @@ using Bit.Core.Utilities;
 using Bit.Core.Vault.Entities;
 using Bit.Core.Vault.Repositories;
 using Microsoft.AspNetCore.Authorization;
+using Microsoft.AspNetCore.Identity;
 using Microsoft.AspNetCore.Mvc;
 
 namespace Bit.Api.Controllers;
@@ -49,6 +53,9 @@ public class AccountsController : Controller
     private readonly ICaptchaValidationService _captchaValidationService;
     private readonly IPolicyService _policyService;
     private readonly ISetInitialMasterPasswordCommand _setInitialMasterPasswordCommand;
+    private readonly IRotateUserKeyCommand _rotateUserKeyCommand;
+    private readonly IFeatureService _featureService;
+    private readonly ICurrentContext _currentContext;
 
 
     public AccountsController(
@@ -65,7 +72,10 @@ public class AccountsController : Controller
         ISendService sendService,
         ICaptchaValidationService captchaValidationService,
         IPolicyService policyService,
-        ISetInitialMasterPasswordCommand setInitialMasterPasswordCommand
+        ISetInitialMasterPasswordCommand setInitialMasterPasswordCommand,
+        IRotateUserKeyCommand rotateUserKeyCommand,
+        IFeatureService featureService,
+        ICurrentContext currentContext
         )
     {
         _cipherRepository = cipherRepository;
@@ -82,6 +92,9 @@ public class AccountsController : Controller
         _captchaValidationService = captchaValidationService;
         _policyService = policyService;
         _setInitialMasterPasswordCommand = setInitialMasterPasswordCommand;
+        _rotateUserKeyCommand = rotateUserKeyCommand;
+        _featureService = featureService;
+        _currentContext = currentContext;
     }
 
     #region DEPRECATED (Moved to Identity Service)
@@ -379,38 +392,59 @@ public class AccountsController : Controller
             throw new UnauthorizedAccessException();
         }
 
-        var ciphers = new List<Cipher>();
+        IdentityResult result;
-        if (model.Ciphers.Any())
+        if (_featureService.IsEnabled(FeatureFlagKeys.KeyRotationImprovements, _currentContext))
         {
-            var existingCiphers = await _cipherRepository.GetManyByUserIdAsync(user.Id);
+            var dataModel = new RotateUserKeyData
-            ciphers.AddRange(existingCiphers
+            {
-                .Join(model.Ciphers, c => c.Id, c => c.Id, (existing, c) => c.ToCipher(existing)));
+                MasterPasswordHash = model.MasterPasswordHash,
-        }
+                Key = model.Key,
+                PrivateKey = model.PrivateKey,
+                // Ciphers = await _cipherValidator.ValidateAsync(user, model.Ciphers),
+                // Folders = await _folderValidator.ValidateAsync(user, model.Folders),
+                // Sends = await _sendValidator.ValidateAsync(user, model.Sends),
+                // EmergencyAccessKeys = await _emergencyAccessValidator.ValidateAsync(user, model.EmergencyAccessKeys),
+                // ResetPasswordKeys = await _accountRecoveryValidator.ValidateAsync(user, model.ResetPasswordKeys),
+            };
 
-        var folders = new List<Folder>();
+            result = await _rotateUserKeyCommand.RotateUserKeyAsync(dataModel);
-        if (model.Folders.Any())
-        {
-            var existingFolders = await _folderRepository.GetManyByUserIdAsync(user.Id);
-            folders.AddRange(existingFolders
-                .Join(model.Folders, f => f.Id, f => f.Id, (existing, f) => f.ToFolder(existing)));
         }
-
+        else
-        var sends = new List<Send>();
-        if (model.Sends?.Any() == true)
         {
-            var existingSends = await _sendRepository.GetManyByUserIdAsync(user.Id);
+            var ciphers = new List<Cipher>();
-            sends.AddRange(existingSends
+            if (model.Ciphers.Any())
-                .Join(model.Sends, s => s.Id, s => s.Id, (existing, s) => s.ToSend(existing, _sendService)));
+            {
+                var existingCiphers = await _cipherRepository.GetManyByUserIdAsync(user.Id);
+                ciphers.AddRange(existingCiphers
+                    .Join(model.Ciphers, c => c.Id, c => c.Id, (existing, c) => c.ToCipher(existing)));
+            }
+
+            var folders = new List<Folder>();
+            if (model.Folders.Any())
+            {
+                var existingFolders = await _folderRepository.GetManyByUserIdAsync(user.Id);
+                folders.AddRange(existingFolders
+                    .Join(model.Folders, f => f.Id, f => f.Id, (existing, f) => f.ToFolder(existing)));
+            }
+
+            var sends = new List<Send>();
+            if (model.Sends?.Any() == true)
+            {
+                var existingSends = await _sendRepository.GetManyByUserIdAsync(user.Id);
+                sends.AddRange(existingSends
+                    .Join(model.Sends, s => s.Id, s => s.Id, (existing, s) => s.ToSend(existing, _sendService)));
+            }
+
+            result = await _userService.UpdateKeyAsync(
+                user,
+                model.MasterPasswordHash,
+                model.Key,
+                model.PrivateKey,
+                ciphers,
+                folders,
+                sends);
         }
 
-        var result = await _userService.UpdateKeyAsync(
-            user,
-            model.MasterPasswordHash,
-            model.Key,
-            model.PrivateKey,
-            ciphers,
-            folders,
-            sends);
 
         if (result.Succeeded)
         {

src/Api/Startup.cs

@@ -15,6 +15,8 @@ using Bit.SharedWeb.Utilities;
 using Microsoft.AspNetCore.Diagnostics.HealthChecks;
 using Microsoft.Extensions.DependencyInjection.Extensions;
 using Bit.Core.Auth.Identity;
+using Bit.Core.Auth.UserFeatures.UserKey;
+using Bit.Core.Auth.UserFeatures.UserKey.Implementations;
 using Bit.Core.OrganizationFeatures.OrganizationSubscriptions;
 
 #if !OSS
@@ -131,6 +133,9 @@ public class Startup
 
         services.AddScoped<AuthenticatorTokenProvider>();
 
+        // Key Rotation
+        services.AddScoped<IRotateUserKeyCommand, RotateUserKeyCommand>();
+
         // Services
         services.AddBaseServices(globalSettings);
         services.AddDefaultServices(globalSettings);

src/Core/Auth/Models/Data/RotateUserKeyData.cs

@@ -0,0 +1,19 @@
+using Bit.Core.Auth.Entities;
+using Bit.Core.Entities;
+using Bit.Core.Tools.Entities;
+using Bit.Core.Vault.Entities;
+
+namespace Bit.Core.Auth.Models.Data;
+
+public class RotateUserKeyData
+{
+    public User User { get; set; }
+    public string MasterPasswordHash { get; set; }
+    public string Key { get; set; }
+    public string PrivateKey { get; set; }
+    public IEnumerable<Cipher> Ciphers { get; set; }
+    public IEnumerable<Folder> Folders { get; set; }
+    public IEnumerable<Send> Sends { get; set; }
+    public IEnumerable<EmergencyAccess> EmergencyAccessKeys { get; set; }
+    public IEnumerable<OrganizationUser> ResetPasswordKeys { get; set; }
+}

src/Core/Auth/UserFeatures/UserKey/IRotateUserKeyCommand.cs

@@ -0,0 +1,18 @@
+using Bit.Core.Auth.Models.Data;
+using Microsoft.AspNetCore.Identity;
+using Microsoft.Data.SqlClient;
+
+namespace Bit.Core.Auth.UserFeatures.UserKey;
+
+public interface IRotateUserKeyCommand
+{
+    /// <summary>
+    /// Sets a new user key and updates all encrypted data.
+    /// </summary>
+    /// <param name="model">All necessary information for rotation. Warning: Any encrypted data not included will be lost.</param>
+    /// <returns>An IdentityResult for verification of the master password hash</returns>
+    /// <exception cref="ArgumentNullException">User must be provided.</exception>
+    Task<IdentityResult> RotateUserKeyAsync(RotateUserKeyData model);
+}
+
+public delegate Task UpdateEncryptedDataForKeyRotation(SqlTransaction transaction = null);

src/Core/Auth/UserFeatures/UserKey/Implementations/RotateUserKeyCommand.cs

@@ -0,0 +1,61 @@
+using Bit.Core.Auth.Models.Data;
+using Bit.Core.Repositories;
+using Bit.Core.Services;
+using Microsoft.AspNetCore.Identity;
+
+namespace Bit.Core.Auth.UserFeatures.UserKey.Implementations;
+
+public class RotateUserKeyCommand : IRotateUserKeyCommand
+{
+    private readonly IUserService _userService;
+    private readonly IUserRepository _userRepository;
+    private readonly IPushNotificationService _pushService;
+    private readonly IdentityErrorDescriber _identityErrorDescriber;
+
+    public RotateUserKeyCommand(IUserService userService, IUserRepository userRepository,
+        IPushNotificationService pushService, IdentityErrorDescriber errors)
+    {
+        _userService = userService;
+        _userRepository = userRepository;
+        _pushService = pushService;
+        _identityErrorDescriber = errors;
+    }
+
+    /// <inheritdoc />
+    public async Task<IdentityResult> RotateUserKeyAsync(RotateUserKeyData model)
+    {
+        if (model.User == null)
+        {
+            throw new ArgumentNullException(nameof(model.User));
+        }
+
+        if (!await _userService.CheckPasswordAsync(model.User, model.MasterPasswordHash))
+        {
+            return IdentityResult.Failed(_identityErrorDescriber.PasswordMismatch());
+        }
+
+        var now = DateTime.UtcNow;
+        model.User.RevisionDate = model.User.AccountRevisionDate = now;
+        model.User.LastKeyRotationDate = now;
+        model.User.SecurityStamp = Guid.NewGuid().ToString();
+        model.User.Key = model.Key;
+        model.User.PrivateKey = model.PrivateKey;
+        if (model.Ciphers.Any() || model.Folders.Any() || model.Sends.Any() || model.EmergencyAccessKeys.Any() ||
+            model.ResetPasswordKeys.Any())
+        {
+            List<UpdateEncryptedDataForKeyRotation> saveEncryptedDataActions = new();
+            // if (model.Ciphers.Any())
+            // {
+            //      saveEncryptedDataActions.Add(_cipherRepository.SaveRotatedData);
+            // }
+            await _userRepository.UpdateUserKeyAndEncryptedDataAsync(model.User, saveEncryptedDataActions);
+        }
+        else
+        {
+            await _userRepository.ReplaceAsync(model.User);
+        }
+
+        await _pushService.PushLogOutAsync(model.User.Id, excludeCurrentContextFromPush: true);
+        return IdentityResult.Success;
+    }
+}

src/Core/Constants.cs

@@ -65,6 +65,7 @@ public static class FeatureFlagKeys
     public const string ItemShare = "item-share";
     public const string BillingPlansUpgrade = "billing-plans-upgrade";
     public const string BillingStarterPlan = "billing-starter-plan";
+    public const string KeyRotationImprovements = "key-rotation-improvements";
 
     public static List<string> GetAllKeys()
     {

src/Core/Repositories/IUserRepository.cs

@@ -1,4 +1,5 @@
-using Bit.Core.Entities;
+using Bit.Core.Auth.UserFeatures.UserKey;
+using Bit.Core.Entities;
 using Bit.Core.Models.Data;
 
 namespace Bit.Core.Repositories;
@@ -15,4 +16,13 @@ public interface IUserRepository : IRepository<User, Guid>
     Task UpdateStorageAsync(Guid id);
     Task UpdateRenewalReminderDateAsync(Guid id, DateTime renewalReminderDate);
     Task<IEnumerable<User>> GetManyAsync(IEnumerable<Guid> ids);
+    /// <summary>
+    /// Sets a new user key and updates all encrypted data.
+    /// <para>Warning: Any user key encrypted data not included will be lost.</para>
+    /// </summary>
+    /// <param name="user">The user to update</param>
+    /// <param name="updateDataActions">Registered database calls to update re-encrypted data.</param>
+    [Obsolete("Intended for future improvements to key rotation. Do not use.")]
+    Task UpdateUserKeyAndEncryptedDataAsync(User user,
+        IEnumerable<UpdateEncryptedDataForKeyRotation> updateDataActions);
 }

src/Infrastructure.Dapper/Repositories/UserRepository.cs

@@ -1,5 +1,6 @@
 using System.Data;
 using Bit.Core;
+using Bit.Core.Auth.UserFeatures.UserKey;
 using Bit.Core.Entities;
 using Bit.Core.Models.Data;
 using Bit.Core.Repositories;
@@ -175,6 +176,52 @@ public class UserRepository : Repository<User, Guid>, IUserRepository
         }
     }
 
+    /// <inheritdoc />
+    public async Task UpdateUserKeyAndEncryptedDataAsync(
+        User user,
+        IEnumerable<UpdateEncryptedDataForKeyRotation> updateDataActions)
+    {
+        await using var connection = new SqlConnection(ConnectionString);
+        connection.Open();
+
+        await using var transaction = connection.BeginTransaction();
+        try
+        {
+            // Update user
+            await using (var cmd = new SqlCommand("[dbo].[User_UpdateKeys]", connection, transaction))
+            {
+                cmd.CommandType = CommandType.StoredProcedure;
+                cmd.Parameters.Add("@Id", SqlDbType.UniqueIdentifier).Value = user.Id;
+                cmd.Parameters.Add("@SecurityStamp", SqlDbType.NVarChar).Value = user.SecurityStamp;
+                cmd.Parameters.Add("@Key", SqlDbType.VarChar).Value = user.Key;
+
+                cmd.Parameters.Add("@PrivateKey", SqlDbType.VarChar).Value =
+                    string.IsNullOrWhiteSpace(user.PrivateKey) ? DBNull.Value : user.PrivateKey;
+
+                cmd.Parameters.Add("@RevisionDate", SqlDbType.DateTime2).Value = user.RevisionDate;
+                cmd.Parameters.Add("@AccountRevisionDate", SqlDbType.DateTime2).Value =
+                    user.AccountRevisionDate;
+                cmd.Parameters.Add("@LastKeyRotationDate", SqlDbType.DateTime2).Value =
+                    user.LastKeyRotationDate;
+                cmd.ExecuteNonQuery();
+            }
+
+            //  Update re-encrypted data
+            foreach (var action in updateDataActions)
+            {
+                await action(transaction);
+            }
+
+            transaction.Commit();
+        }
+        catch
+        {
+            transaction.Rollback();
+            throw;
+        }
+    }
+
+
     public async Task<IEnumerable<User>> GetManyAsync(IEnumerable<Guid> ids)
     {
         using (var connection = new SqlConnection(ReadOnlyConnectionString))

src/Infrastructure.EntityFramework/Repositories/UserRepository.cs

@@ -1,4 +1,5 @@
 using AutoMapper;
+using Bit.Core.Auth.UserFeatures.UserKey;
 using Bit.Core.Repositories;
 using Bit.Infrastructure.EntityFramework.Models;
 using Microsoft.EntityFrameworkCore;
@@ -135,6 +136,48 @@ public class UserRepository : Repository<Core.Entities.User, User, Guid>, IUserR
         }
     }
 
+    /// <inheritdoc />
+    public async Task UpdateUserKeyAndEncryptedDataAsync(Core.Entities.User user,
+        IEnumerable<UpdateEncryptedDataForKeyRotation> updateDataActions)
+    {
+        using var scope = ServiceScopeFactory.CreateScope();
+        var dbContext = GetDatabaseContext(scope);
+
+        await using var transaction = await dbContext.Database.BeginTransactionAsync();
+
+        try
+        {
+            // Update user
+            var entity = await dbContext.Users.FindAsync(user.Id);
+            if (entity == null)
+            {
+                throw new ArgumentException("User not found", nameof(user));
+            }
+
+            entity.SecurityStamp = user.SecurityStamp;
+            entity.Key = user.Key;
+            entity.PrivateKey = user.PrivateKey;
+            entity.LastKeyRotationDate = user.LastKeyRotationDate;
+            entity.AccountRevisionDate = user.AccountRevisionDate;
+            entity.RevisionDate = user.RevisionDate;
+
+            //  Update re-encrypted data
+            foreach (var action in updateDataActions)
+            {
+                // TODO (jlf0dev): Check if transaction captures these operations
+                await action();
+            }
+
+            await transaction.CommitAsync();
+        }
+        catch
+        {
+            await transaction.RollbackAsync();
+            throw;
+        }
+
+    }
+
     public async Task<IEnumerable<Core.Entities.User>> GetManyAsync(IEnumerable<Guid> ids)
     {
         using (var scope = ServiceScopeFactory.CreateScope())

test/Api.Test/Controllers/AccountsControllerTests.cs

@@ -4,7 +4,9 @@ using Bit.Api.Controllers;
 using Bit.Core.AdminConsole.Repositories;
 using Bit.Core.Auth.Models.Api.Request.Accounts;
 using Bit.Core.Auth.Services;
+using Bit.Core.Auth.UserFeatures.UserKey;
 using Bit.Core.Auth.UserFeatures.UserMasterPassword.Interfaces;
+using Bit.Core.Context;
 using Bit.Core.Entities;
 using Bit.Core.Enums;
 using Bit.Core.Exceptions;
@@ -40,6 +42,10 @@ public class AccountsControllerTests : IDisposable
     private readonly ICaptchaValidationService _captchaValidationService;
     private readonly IPolicyService _policyService;
     private readonly ISetInitialMasterPasswordCommand _setInitialMasterPasswordCommand;
+    private readonly IRotateUserKeyCommand _rotateUserKeyCommand;
+    private readonly IFeatureService _featureService;
+    private readonly ICurrentContext _currentContext;
+
 
     public AccountsControllerTests()
     {
@@ -57,6 +63,9 @@ public class AccountsControllerTests : IDisposable
         _captchaValidationService = Substitute.For<ICaptchaValidationService>();
         _policyService = Substitute.For<IPolicyService>();
         _setInitialMasterPasswordCommand = Substitute.For<ISetInitialMasterPasswordCommand>();
+        _rotateUserKeyCommand = Substitute.For<IRotateUserKeyCommand>();
+        _featureService = Substitute.For<IFeatureService>();
+        _currentContext = Substitute.For<ICurrentContext>();
 
         _sut = new AccountsController(
             _globalSettings,
@@ -72,7 +81,10 @@ public class AccountsControllerTests : IDisposable
             _sendService,
             _captchaValidationService,
             _policyService,
-            _setInitialMasterPasswordCommand
+            _setInitialMasterPasswordCommand,
+            _rotateUserKeyCommand,
+            _featureService,
+            _currentContext
         );
     }
 

test/Core.Test/Auth/UserFeatures/UserKey/RotateUserKeyCommandTests.cs

@@ -0,0 +1,50 @@
+using Bit.Core.Auth.Models.Data;
+using Bit.Core.Auth.UserFeatures.UserKey.Implementations;
+using Bit.Core.Services;
+using Bit.Test.Common.AutoFixture;
+using Bit.Test.Common.AutoFixture.Attributes;
+using Microsoft.AspNetCore.Identity;
+using NSubstitute;
+using Xunit;
+
+namespace Bit.Core.Test.Auth.UserFeatures.UserKey;
+
+[SutProviderCustomize]
+public class RotateUserKeyCommandTests
+{
+    [Theory, BitAutoData]
+    public async Task RotateUserKeyAsync_Success(SutProvider<RotateUserKeyCommand> sutProvider, RotateUserKeyData model)
+    {
+        sutProvider.GetDependency<IUserService>().CheckPasswordAsync(model.User, model.MasterPasswordHash)
+            .Returns(true);
+
+        var result = await sutProvider.Sut.RotateUserKeyAsync(model);
+
+        Assert.Equal(IdentityResult.Success, result);
+    }
+
+    [Theory, BitAutoData]
+    public async Task RotateUserKeyAsync_InvalidMasterPasswordHash_ReturnsFailedIdentityResult(
+        SutProvider<RotateUserKeyCommand> sutProvider, RotateUserKeyData model)
+    {
+        sutProvider.GetDependency<IUserService>().CheckPasswordAsync(model.User, model.MasterPasswordHash)
+            .Returns(false);
+
+        var result = await sutProvider.Sut.RotateUserKeyAsync(model);
+
+        Assert.False(result.Succeeded);
+    }
+
+    [Theory, BitAutoData]
+    public async Task RotateUserKeyAsync_LogsOutUser(
+        SutProvider<RotateUserKeyCommand> sutProvider, RotateUserKeyData model)
+    {
+        sutProvider.GetDependency<IUserService>().CheckPasswordAsync(model.User, model.MasterPasswordHash)
+            .Returns(true);
+
+        await sutProvider.Sut.RotateUserKeyAsync(model);
+
+        await sutProvider.GetDependency<IPushNotificationService>().ReceivedWithAnyArgs()
+            .PushLogOutAsync(default, default);
+    }
+}