diff --git a/test/Api.Test/Auth/Controllers/AccountsControllerTests.cs b/test/Api.Test/Auth/Controllers/AccountsControllerTests.cs index f1aa11d068..e0b7eedecc 100644 --- a/test/Api.Test/Auth/Controllers/AccountsControllerTests.cs +++ b/test/Api.Test/Auth/Controllers/AccountsControllerTests.cs @@ -10,7 +10,9 @@ using Bit.Core.Auth.UserFeatures.TwoFactorAuth.Interfaces; using Bit.Core.Auth.UserFeatures.UserMasterPassword.Interfaces; using Bit.Core.Entities; using Bit.Core.Exceptions; +using Bit.Core.KeyManagement.Commands; using Bit.Core.KeyManagement.Kdf; +using Bit.Core.KeyManagement.Models.Api.Request; using Bit.Core.KeyManagement.Models.Data; using Bit.Core.KeyManagement.Queries.Interfaces; using Bit.Core.Repositories; @@ -38,6 +40,7 @@ public class AccountsControllerTests : IDisposable private readonly IUserAccountKeysQuery _userAccountKeysQuery; private readonly ITwoFactorEmailService _twoFactorEmailService; private readonly IChangeKdfCommand _changeKdfCommand; + private readonly SetAccountKeysForUserCommand _setAccountKeysForUserCommand; public AccountsControllerTests() { @@ -53,6 +56,7 @@ public class AccountsControllerTests : IDisposable _userAccountKeysQuery = Substitute.For(); _twoFactorEmailService = Substitute.For(); _changeKdfCommand = Substitute.For(); + _setAccountKeysForUserCommand = Substitute.For(); _sut = new AccountsController( _organizationService, @@ -66,7 +70,8 @@ public class AccountsControllerTests : IDisposable _featureService, _userAccountKeysQuery, _twoFactorEmailService, - _changeKdfCommand + _changeKdfCommand, + _setAccountKeysForUserCommand ); } @@ -738,5 +743,62 @@ public class AccountsControllerTests : IDisposable _userService.GetUserByIdAsync(Arg.Any()) .Returns(Task.FromResult((User)null)); } + + [Theory, BitAutoData] + public async Task PostKeys_WithAccountKeys_CallsSetAccountKeysCommand( + User user, + KeysRequestModel model) + { + // Arrange + user.PublicKey = null; + user.PrivateKey = null; + model.AccountKeys = new AccountKeysRequestModel + { + UserKeyEncryptedAccountPrivateKey = "wrapped-private-key", + AccountPublicKey = "public-key" + }; + + _userService.GetUserByPrincipalAsync(Arg.Any()).Returns(user); + _featureService.IsEnabled(Bit.Core.FeatureFlagKeys.ReturnErrorOnExistingKeypair).Returns(false); + + // Act + var result = await _sut.PostKeys(model); + + // Assert + await _setAccountKeysForUserCommand.Received(1).SetAccountKeysForUserAsync( + user.Id, + model.AccountKeys); + await _userService.DidNotReceiveWithAnyArgs().SaveUserAsync(Arg.Any()); + Assert.NotNull(result); + Assert.Equal("keys", result.Object); + } + + [Theory, BitAutoData] + public async Task PostKeys_WithoutAccountKeys_CallsSaveUser( + User user, + KeysRequestModel model) + { + // Arrange + user.PublicKey = null; + user.PrivateKey = null; + model.AccountKeys = null; + model.PublicKey = "public-key"; + model.EncryptedPrivateKey = "encrypted-private-key"; + + _userService.GetUserByPrincipalAsync(Arg.Any()).Returns(user); + _featureService.IsEnabled(Bit.Core.FeatureFlagKeys.ReturnErrorOnExistingKeypair).Returns(false); + + // Act + var result = await _sut.PostKeys(model); + + // Assert + await _userService.Received(1).SaveUserAsync(Arg.Is(u => + u.PublicKey == model.PublicKey && + u.PrivateKey == model.EncryptedPrivateKey)); + await _setAccountKeysForUserCommand.DidNotReceiveWithAnyArgs() + .SetAccountKeysForUserAsync(Arg.Any(), Arg.Any()); + Assert.NotNull(result); + Assert.Equal("keys", result.Object); + } }