[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 */