[PM-27279] Implement TDE Registration with V2 Keys (#6671)

* Implement TDE v2 signup

* Clean up fallback logic for account keys

* Fix broken v2 logic

* Add comment

* Update comment

src/Api/Auth/Controllers/AccountsController.cs

@@ -18,6 +18,7 @@ using Bit.Core.Auth.UserFeatures.UserMasterPassword.Interfaces;
 using Bit.Core.Enums;
 using Bit.Core.Exceptions;
 using Bit.Core.KeyManagement.Kdf;
+using Bit.Core.KeyManagement.Models.Data;
 using Bit.Core.KeyManagement.Queries.Interfaces;
 using Bit.Core.Models.Api.Response;
 using Bit.Core.Repositories;
@@ -44,6 +45,7 @@ public class AccountsController : Controller
     private readonly IUserAccountKeysQuery _userAccountKeysQuery;
     private readonly ITwoFactorEmailService _twoFactorEmailService;
     private readonly IChangeKdfCommand _changeKdfCommand;
+    private readonly IUserRepository _userRepository;
 
     public AccountsController(
         IOrganizationService organizationService,
@@ -57,7 +59,8 @@ public class AccountsController : Controller
         IFeatureService featureService,
         IUserAccountKeysQuery userAccountKeysQuery,
         ITwoFactorEmailService twoFactorEmailService,
-        IChangeKdfCommand changeKdfCommand
+        IChangeKdfCommand changeKdfCommand,
+        IUserRepository userRepository
         )
     {
         _organizationService = organizationService;
@@ -72,6 +75,7 @@ public class AccountsController : Controller
         _userAccountKeysQuery = userAccountKeysQuery;
         _twoFactorEmailService = twoFactorEmailService;
         _changeKdfCommand = changeKdfCommand;
+        _userRepository = userRepository;
     }
 
 
@@ -440,8 +444,31 @@ public class AccountsController : Controller
             }
         }
 
-        await _userService.SaveUserAsync(model.ToUser(user));
-        return new KeysResponseModel(user);
+        if (model.AccountKeys != null)
+        {
+            var accountKeysData = model.AccountKeys.ToAccountKeysData();
+            if (!accountKeysData.IsV2Encryption())
+            {
+                throw new BadRequestException("AccountKeys are only supported for V2 encryption.");
+            }
+            await _userRepository.SetV2AccountCryptographicStateAsync(user.Id, accountKeysData);
+            return new KeysResponseModel(accountKeysData, user.Key);
+        }
+        else
+        {
+            // Todo: Drop this after a transition period. This will drop no-account-keys requests.
+            // The V1 check in the other branch should persist
+            // https://bitwarden.atlassian.net/browse/PM-27329
+            await _userService.SaveUserAsync(model.ToUser(user));
+            return new KeysResponseModel(new UserAccountKeysData
+            {
+                PublicKeyEncryptionKeyPairData = new PublicKeyEncryptionKeyPairData(
+                    user.PrivateKey,
+                    user.PublicKey
+                )
+            }, user.Key);
+        }
+
     }
 
     [HttpGet("keys")]
@@ -453,7 +480,8 @@ public class AccountsController : Controller
             throw new UnauthorizedAccessException();
         }
 
-        return new KeysResponseModel(user);
+        var accountKeys = await _userAccountKeysQuery.Run(user);
+        return new KeysResponseModel(accountKeys, user.Key);
     }
 
     [HttpDelete]

src/Api/Models/Response/KeysResponseModel.cs

@@ -1,27 +1,32 @@
-// FIXME: Update this file to be null safe and then delete the line below
-#nullable disable
-
-using Bit.Core.Entities;
+using Bit.Core.KeyManagement.Models.Api.Response;
+using Bit.Core.KeyManagement.Models.Data;
 using Bit.Core.Models.Api;
 
 namespace Bit.Api.Models.Response;
 
 public class KeysResponseModel : ResponseModel
 {
-    public KeysResponseModel(User user)
+    public KeysResponseModel(UserAccountKeysData accountKeys, string? masterKeyWrappedUserKey)
         : base("keys")
     {
-        if (user == null)
+        if (masterKeyWrappedUserKey != null)
         {
-            throw new ArgumentNullException(nameof(user));
+            Key = masterKeyWrappedUserKey;
         }
 
-        Key = user.Key;
-        PublicKey = user.PublicKey;
-        PrivateKey = user.PrivateKey;
+        PublicKey = accountKeys.PublicKeyEncryptionKeyPairData.PublicKey;
+        PrivateKey = accountKeys.PublicKeyEncryptionKeyPairData.WrappedPrivateKey;
+        AccountKeys = new PrivateKeysResponseModel(accountKeys);
     }
 
-    public string Key { get; set; }
+    /// <summary>
+    /// The master key wrapped user key. The master key can either be a master-password master key or a
+    /// key-connector master key.
+    /// </summary>
+    public string? Key { get; set; }
+    [Obsolete("Use AccountKeys.PublicKeyEncryptionKeyPair.PublicKey instead")]
     public string PublicKey { get; set; }
+    [Obsolete("Use AccountKeys.PublicKeyEncryptionKeyPair.WrappedPrivateKey instead")]
     public string PrivateKey { get; set; }
+    public PrivateKeysResponseModel AccountKeys { get; set; }
 }

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

@@ -3,17 +3,22 @@
 
 using System.ComponentModel.DataAnnotations;
 using Bit.Core.Entities;
+using Bit.Core.KeyManagement.Models.Api.Request;
 using Bit.Core.Utilities;
 
 namespace Bit.Core.Auth.Models.Api.Request.Accounts;
 
 public class KeysRequestModel
 {
+    [Obsolete("Use AccountKeys.AccountPublicKey instead")]
     [Required]
     public string PublicKey { get; set; }
+    [Obsolete("Use AccountKeys.UserKeyEncryptedAccountPrivateKey instead")]
     [Required]
     public string EncryptedPrivateKey { get; set; }
+    public AccountKeysRequestModel AccountKeys { get; set; }
 
+    [Obsolete("Use SetAccountKeysForUserCommand instead")]
     public User ToUser(User existingUser)
     {
         if (string.IsNullOrWhiteSpace(PublicKey) || string.IsNullOrWhiteSpace(EncryptedPrivateKey))

src/Core/Constants.cs

@@ -211,6 +211,7 @@ public static class FeatureFlagKeys
     public const string NoLogoutOnKdfChange = "pm-23995-no-logout-on-kdf-change";
     public const string DisableType0Decryption = "pm-25174-disable-type-0-decryption";
     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";
 
     /* Mobile Team */

test/Api.Test/Auth/Controllers/AccountsControllerTests.cs

@@ -11,6 +11,7 @@ using Bit.Core.Auth.UserFeatures.UserMasterPassword.Interfaces;
 using Bit.Core.Entities;
 using Bit.Core.Exceptions;
 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 +39,7 @@ public class AccountsControllerTests : IDisposable
     private readonly IUserAccountKeysQuery _userAccountKeysQuery;
     private readonly ITwoFactorEmailService _twoFactorEmailService;
     private readonly IChangeKdfCommand _changeKdfCommand;
+    private readonly IUserRepository _userRepository;
 
     public AccountsControllerTests()
     {
@@ -53,6 +55,7 @@ public class AccountsControllerTests : IDisposable
         _userAccountKeysQuery = Substitute.For<IUserAccountKeysQuery>();
         _twoFactorEmailService = Substitute.For<ITwoFactorEmailService>();
         _changeKdfCommand = Substitute.For<IChangeKdfCommand>();
+        _userRepository = Substitute.For<IUserRepository>();
 
         _sut = new AccountsController(
             _organizationService,
@@ -66,7 +69,8 @@ public class AccountsControllerTests : IDisposable
             _featureService,
             _userAccountKeysQuery,
             _twoFactorEmailService,
-            _changeKdfCommand
+            _changeKdfCommand,
+            _userRepository
         );
     }
 
@@ -738,5 +742,79 @@ public class AccountsControllerTests : IDisposable
         _userService.GetUserByIdAsync(Arg.Any<Guid>())
                     .Returns(Task.FromResult((User)null));
     }
+
+    [Theory, BitAutoData]
+    public async Task PostKeys_WithAccountKeys_CallsSetV2AccountCryptographicState(
+        User user,
+        KeysRequestModel model)
+    {
+        // Arrange
+        user.PublicKey = "public-key";
+        user.PrivateKey = "encrypted-private-key";
+        model.AccountKeys = new AccountKeysRequestModel
+        {
+            UserKeyEncryptedAccountPrivateKey = "wrapped-private-key",
+            AccountPublicKey = "public-key",
+            PublicKeyEncryptionKeyPair = new PublicKeyEncryptionKeyPairRequestModel
+            {
+                PublicKey = "public-key",
+                WrappedPrivateKey = "wrapped-private-key",
+                SignedPublicKey = "signed-public-key"
+            },
+            SignatureKeyPair = new SignatureKeyPairRequestModel
+            {
+                VerifyingKey = "verifying-key",
+                SignatureAlgorithm = "ed25519",
+                WrappedSigningKey = "wrapped-signing-key"
+            },
+            SecurityState = new SecurityStateModel
+            {
+                SecurityState = "security-state",
+                SecurityVersion = 2
+            }
+        };
+
+        _userService.GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>()).Returns(user);
+        _featureService.IsEnabled(Bit.Core.FeatureFlagKeys.ReturnErrorOnExistingKeypair).Returns(false);
+
+        // Act
+        var result = await _sut.PostKeys(model);
+
+        // Assert
+        await _userRepository.Received(1).SetV2AccountCryptographicStateAsync(
+            user.Id,
+            Arg.Any<UserAccountKeysData>());
+        await _userService.DidNotReceiveWithAnyArgs().SaveUserAsync(Arg.Any<User>());
+        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<ClaimsPrincipal>()).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<User>(u =>
+            u.PublicKey == model.PublicKey &&
+            u.PrivateKey == model.EncryptedPrivateKey));
+        await _userRepository.DidNotReceiveWithAnyArgs()
+            .SetV2AccountCryptographicStateAsync(Arg.Any<Guid>(), Arg.Any<UserAccountKeysData>());
+        Assert.NotNull(result);
+        Assert.Equal("keys", result.Object);
+    }
 }
 

test/Identity.IntegrationTest/Controllers/AccountsControllerTests.cs

@@ -139,6 +139,7 @@ public class AccountsControllerTests : IClassFixture<IdentityApplicationFactory>
          [StringLength(1000), Required] string masterPasswordHash, [StringLength(50)] string masterPasswordHint, [Required] string userSymmetricKey,
          [Required] KeysRequestModel userAsymmetricKeys, int kdfMemory, int kdfParallelism)
     {
+        userAsymmetricKeys.AccountKeys = null;
         // Localize substitutions to this test.
         var localFactory = new IdentityApplicationFactory();
 
@@ -202,6 +203,7 @@ public class AccountsControllerTests : IClassFixture<IdentityApplicationFactory>
        [StringLength(1000), Required] string masterPasswordHash, [StringLength(50)] string masterPasswordHint, [Required] string userSymmetricKey,
        [Required] KeysRequestModel userAsymmetricKeys, int kdfMemory, int kdfParallelism)
     {
+        userAsymmetricKeys.AccountKeys = null;
         // Localize substitutions to this test.
         var localFactory = new IdentityApplicationFactory();
         localFactory.UpdateConfiguration("globalSettings:disableUserRegistration", "true");
@@ -233,6 +235,7 @@ public class AccountsControllerTests : IClassFixture<IdentityApplicationFactory>
          [StringLength(1000)] string masterPasswordHash, [StringLength(50)] string masterPasswordHint, string userSymmetricKey,
         KeysRequestModel userAsymmetricKeys, int kdfMemory, int kdfParallelism)
     {
+        userAsymmetricKeys.AccountKeys = null;
 
         // Localize factory to just this test.
         var localFactory = new IdentityApplicationFactory();
@@ -310,6 +313,7 @@ public class AccountsControllerTests : IClassFixture<IdentityApplicationFactory>
      [StringLength(1000)] string masterPasswordHash, [StringLength(50)] string masterPasswordHint, string userSymmetricKey,
     KeysRequestModel userAsymmetricKeys, int kdfMemory, int kdfParallelism, Guid orgSponsorshipId)
     {
+        userAsymmetricKeys.AccountKeys = null;
 
         // Localize factory to just this test.
         var localFactory = new IdentityApplicationFactory();
@@ -386,6 +390,7 @@ public class AccountsControllerTests : IClassFixture<IdentityApplicationFactory>
      [StringLength(1000)] string masterPasswordHash, [StringLength(50)] string masterPasswordHint, string userSymmetricKey,
     KeysRequestModel userAsymmetricKeys, int kdfMemory, int kdfParallelism, EmergencyAccess emergencyAccess)
     {
+        userAsymmetricKeys.AccountKeys = null;
 
         // Localize factory to just this test.
         var localFactory = new IdentityApplicationFactory();
@@ -455,6 +460,7 @@ public class AccountsControllerTests : IClassFixture<IdentityApplicationFactory>
      [StringLength(1000)] string masterPasswordHash, [StringLength(50)] string masterPasswordHint, string userSymmetricKey,
     KeysRequestModel userAsymmetricKeys, int kdfMemory, int kdfParallelism)
     {
+        userAsymmetricKeys.AccountKeys = null;
 
         // Localize factory to just this test.
         var localFactory = new IdentityApplicationFactory();

test/Identity.IntegrationTest/Endpoints/IdentityServerTests.cs

@@ -21,6 +21,13 @@ namespace Bit.Identity.IntegrationTest.Endpoints;
 [SutProviderCustomize]
 public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
 {
+    private static readonly KeysRequestModel TEST_ACCOUNT_KEYS = new KeysRequestModel
+    {
+        AccountKeys = null,
+        PublicKey = "public-key",
+        EncryptedPrivateKey = "encrypted-private-key",
+    };
+
     private const int SecondsInMinute = 60;
     private const int MinutesInHour = 60;
     private const int SecondsInHour = SecondsInMinute * MinutesInHour;
@@ -53,6 +60,7 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
     [Theory, BitAutoData, RegisterFinishRequestModelCustomize]
     public async Task TokenEndpoint_GrantTypePassword_Success(RegisterFinishRequestModel requestModel)
     {
+        requestModel.UserAsymmetricKeys = TEST_ACCOUNT_KEYS;
         var localFactory = new IdentityApplicationFactory();
         var user = await localFactory.RegisterNewIdentityFactoryUserAsync(requestModel);
 
@@ -78,6 +86,7 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
     public async Task TokenEndpoint_GrantTypePassword_WithAllUserTypes_WithSsoPolicyDisabled_WithEnforceSsoPolicyForAllUsersTrue_Success(
        OrganizationUserType organizationUserType, RegisterFinishRequestModel requestModel, Guid organizationId, int generatedUsername)
     {
+        requestModel.UserAsymmetricKeys = TEST_ACCOUNT_KEYS;
         requestModel.Email = $"{generatedUsername}@example.com";
 
         var localFactory = new IdentityApplicationFactory();
@@ -103,6 +112,7 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
     public async Task TokenEndpoint_GrantTypePassword_WithAllUserTypes_WithSsoPolicyDisabled_WithEnforceSsoPolicyForAllUsersFalse_Success(
         OrganizationUserType organizationUserType, RegisterFinishRequestModel requestModel, Guid organizationId, int generatedUsername)
     {
+        requestModel.UserAsymmetricKeys = TEST_ACCOUNT_KEYS;
         requestModel.Email = $"{generatedUsername}@example.com";
 
         var localFactory = new IdentityApplicationFactory();
@@ -129,6 +139,7 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
     public async Task TokenEndpoint_GrantTypePassword_WithAllUserTypes_WithSsoPolicyEnabled_WithEnforceSsoPolicyForAllUsersTrue_Throw(
         OrganizationUserType organizationUserType, RegisterFinishRequestModel requestModel, Guid organizationId, int generatedUsername)
     {
+        requestModel.UserAsymmetricKeys = TEST_ACCOUNT_KEYS;
         requestModel.Email = $"{generatedUsername}@example.com";
 
         var localFactory = new IdentityApplicationFactory();
@@ -152,6 +163,7 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
     public async Task TokenEndpoint_GrantTypePassword_WithOwnerOrAdmin_WithSsoPolicyEnabled_WithEnforceSsoPolicyForAllUsersFalse_Success(
         OrganizationUserType organizationUserType, RegisterFinishRequestModel requestModel, Guid organizationId, int generatedUsername)
     {
+        requestModel.UserAsymmetricKeys = TEST_ACCOUNT_KEYS;
         requestModel.Email = $"{generatedUsername}@example.com";
 
         var localFactory = new IdentityApplicationFactory();
@@ -175,6 +187,7 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
     public async Task TokenEndpoint_GrantTypePassword_WithNonOwnerOrAdmin_WithSsoPolicyEnabled_WithEnforceSsoPolicyForAllUsersFalse_Throws(
         OrganizationUserType organizationUserType, RegisterFinishRequestModel requestModel, Guid organizationId, int generatedUsername)
     {
+        requestModel.UserAsymmetricKeys = TEST_ACCOUNT_KEYS;
         requestModel.Email = $"{generatedUsername}@example.com";
 
         var localFactory = new IdentityApplicationFactory();
@@ -196,6 +209,7 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
     [Theory, BitAutoData, RegisterFinishRequestModelCustomize]
     public async Task TokenEndpoint_GrantTypeRefreshToken_Success(RegisterFinishRequestModel requestModel)
     {
+        requestModel.UserAsymmetricKeys = TEST_ACCOUNT_KEYS;
         var localFactory = new IdentityApplicationFactory();
 
         var user = await localFactory.RegisterNewIdentityFactoryUserAsync(requestModel);
@@ -218,6 +232,7 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
     [Theory, BitAutoData, RegisterFinishRequestModelCustomize]
     public async Task TokenEndpoint_GrantTypeClientCredentials_Success(RegisterFinishRequestModel model)
     {
+        model.UserAsymmetricKeys = TEST_ACCOUNT_KEYS;
         var localFactory = new IdentityApplicationFactory();
         var user = await localFactory.RegisterNewIdentityFactoryUserAsync(model);
 
@@ -242,6 +257,7 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
         RegisterFinishRequestModel model,
         string deviceId)
     {
+        model.UserAsymmetricKeys.AccountKeys = null;
         var localFactory = new IdentityApplicationFactory();
         var server = localFactory.WithWebHostBuilder(builder =>
         {
@@ -445,6 +461,7 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
     public async Task TokenEndpoint_TooQuickInOneSecond_BlockRequest(
         RegisterFinishRequestModel requestModel)
     {
+        requestModel.UserAsymmetricKeys = TEST_ACCOUNT_KEYS;
         const int AmountInOneSecondAllowed = 10;
 
         // The rule we are testing is 10 requests in 1 second