[PM-35599] Update set-password endpoint to handle V1 encryption when given Auth and Unlock data

src/Api/Auth/Controllers/AccountsController.cs

@@ -250,48 +250,58 @@ public async Task PostSetPasswordAsync([FromBody] SetInitialPasswordRequestModel
             throw new UnauthorizedAccessException();
         }
 
-        if (model.IsV2Request())
+        // V2 encryption - TDE user with "manage account recovery" permission
+        if (
+            model.HasAuthAndUnlockData() &&
+            model.IsTdeSetPasswordRequest() &&
+            _featureService.IsEnabled(FeatureFlagKeys.V2RegistrationTDEJIT))
         {
-            if (model.IsTdeSetPasswordRequest())
-            {
-                await _tdeSetPasswordCommand.SetMasterPasswordAsync(user, model.ToData());
-            }
-            else
-            {
-                await _finishSsoJitProvisionMasterPasswordCommand.FinishProvisionAsync(user, model.ToData());
-            }
+            await _tdeSetPasswordCommand.SetMasterPasswordAsync(user, model.ToData());
+            return;
         }
-        else
+
+        // V2 encryption - MP JIT
+        if (
+            model.HasAuthAndUnlockData() &&
+            model.IsJitMpSetPasswordRequest() &&
+            _featureService.IsEnabled(FeatureFlagKeys.EnableAccountEncryptionV2JitPasswordRegistration))
         {
-            // TODO removed with https://bitwarden.atlassian.net/browse/PM-27327
-            try
-            {
-                user = model.ToUser(user);
-            }
-            catch (Exception e)
-            {
-                ModelState.AddModelError(string.Empty, e.Message);
-                throw new BadRequestException(ModelState);
-            }
+            await _finishSsoJitProvisionMasterPasswordCommand.FinishProvisionAsync(user, model.ToData());
+            return;
+        }
 
-            var result = await _setInitialMasterPasswordCommandV1.SetInitialMasterPasswordAsync(
-                user,
-                model.MasterPasswordHash,
-                model.Key,
-                model.OrgIdentifier);
+        // V1 encryption — handles both modern clients (sending MPAD/MPUD + legacy Keys) and
+        // legacy clients (sending only legacy fields). The model's ToUser() handles the fallback.
+        // TODO: removal requires that BOTH flags have been removed:
+        //  - https://bitwarden.atlassian.net/browse/PM-27327 (MP)
+        //  - https://bitwarden.atlassian.net/browse/PM-27329 (TDE)
+        try
+        {
+            user = model.ToUser(user);
+        }
+        catch (Exception e)
+        {
+            ModelState.AddModelError(string.Empty, e.Message);
+            throw new BadRequestException(ModelState);
+        }
 
-            if (result.Succeeded)
-            {
-                return;
-            }
+        var result = await _setInitialMasterPasswordCommandV1.SetInitialMasterPasswordAsync(
+            user,
+            model.MasterPasswordAuthentication?.MasterPasswordAuthenticationHash ?? model.MasterPasswordHash,
+            model.MasterPasswordUnlock?.MasterKeyWrappedUserKey ?? model.Key,
+            model.OrgIdentifier);
 
-            foreach (var error in result.Errors)
-            {
-                ModelState.AddModelError(string.Empty, error.Description);
-            }
+        if (result.Succeeded)
+        {
+            return;
+        }
 
-            throw new BadRequestException(ModelState);
+        foreach (var error in result.Errors)
+        {
+            ModelState.AddModelError(string.Empty, error.Description);
         }
+
+        throw new BadRequestException(ModelState);
     }
 
     [HttpPost("verify-password")]

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

@@ -10,7 +10,9 @@ namespace Bit.Api.Auth.Models.Request.Accounts;
 
 public class SetInitialPasswordRequestModel : IValidatableObject
 {
-    // TODO will be removed with https://bitwarden.atlassian.net/browse/PM-27327
+    // TODO: removal requires that BOTH flags have been removed:
+    //  - https://bitwarden.atlassian.net/browse/PM-27327 (MP)
+    //  - https://bitwarden.atlassian.net/browse/PM-27329 (TDE)
     [Obsolete("Use MasterPasswordAuthentication instead")]
     [StringLength(300)]
     public string? MasterPasswordHash { get; set; }
@@ -43,22 +45,26 @@ public class SetInitialPasswordRequestModel : IValidatableObject
     [Required]
     public required string OrgIdentifier { get; set; }
 
-    // TODO removed with https://bitwarden.atlassian.net/browse/PM-27327
+    // Reads KDF/key from MasterPasswordAuthentication/MasterPasswordUnlock when present (modern clients),
+    // and falls back to the top-level legacy properties when not (clients ≤3 releases back).
+    // TODO: removal requires that BOTH flags have been removed:
+    //  - https://bitwarden.atlassian.net/browse/PM-27327 (MP)
+    //  - https://bitwarden.atlassian.net/browse/PM-27329 (TDE)
     public User ToUser(User existingUser)
     {
         existingUser.MasterPasswordHint = MasterPasswordHint;
-        existingUser.Kdf = Kdf!.Value;
-        existingUser.KdfIterations = KdfIterations!.Value;
-        existingUser.KdfMemory = KdfMemory;
-        existingUser.KdfParallelism = KdfParallelism;
-        existingUser.Key = Key;
+        existingUser.Kdf = MasterPasswordAuthentication?.Kdf.KdfType ?? Kdf!.Value;
+        existingUser.KdfIterations = MasterPasswordAuthentication?.Kdf.Iterations ?? KdfIterations!.Value;
+        existingUser.KdfMemory = MasterPasswordAuthentication?.Kdf.Memory ?? KdfMemory;
+        existingUser.KdfParallelism = MasterPasswordAuthentication?.Kdf.Parallelism ?? KdfParallelism;
+        existingUser.Key = MasterPasswordUnlock?.MasterKeyWrappedUserKey ?? Key;
         Keys?.ToUser(existingUser);
         return existingUser;
     }
 
     public IEnumerable<ValidationResult> Validate(ValidationContext validationContext)
     {
-        if (IsV2Request())
+        if (HasAuthAndUnlockData())
         {
             // V2 registration - validate KDF equality, salt equality, and KDF settings
             foreach (var validationResult in KdfSettingsValidator.ValidateAuthenticationAndUnlockData(
@@ -71,7 +77,9 @@ public IEnumerable<ValidationResult> Validate(ValidationContext validationContex
         }
 
         // V1 registration
-        // TODO removed with https://bitwarden.atlassian.net/browse/PM-27327
+        // TODO: removal requires that BOTH flags have been removed:
+        //  - https://bitwarden.atlassian.net/browse/PM-27327 (MP)
+        //  - https://bitwarden.atlassian.net/browse/PM-27329 (TDE)
         if (string.IsNullOrEmpty(MasterPasswordHash))
         {
             yield return new ValidationResult("MasterPasswordHash must be supplied.");
@@ -115,15 +123,30 @@ public IEnumerable<ValidationResult> Validate(ValidationContext validationContex
         }
     }
 
-    public bool IsV2Request()
+    /// <summary>
+    /// True when the request uses the new data shape (MasterPasswordAuthentication + MasterPasswordUnlock).
+    /// This is a shape check, NOT a guarantee that V2 encryption will run. It is possible for V1 encryption
+    /// to run even when the request contains these new data types (see `set-password` endpoint).
+    /// Feature flags and AccountKeys presence determine the actual flow (V1 or V2).
+    /// </summary>
+    public bool HasAuthAndUnlockData()
     {
         // AccountKeys can be null for TDE users, so we don't check that here
         return MasterPasswordAuthentication != null && MasterPasswordUnlock != null;
     }
 
+    // TDE users don't send any key material (their keypair already exists).
+    // Checks both AccountKeys (new) and Keys (legacy) so the predicate is correct for
+    // the transitional period where clients may send either key shape.
     public bool IsTdeSetPasswordRequest()
     {
-        return AccountKeys == null;
+        return AccountKeys == null && Keys == null;
+    }
+
+    // MP JIT users send new key material — either via the new AccountKeys shape or the legacy Keys shape.
+    public bool IsJitMpSetPasswordRequest()
+    {
+        return AccountKeys != null || Keys != null;
     }
 
     public SetInitialMasterPasswordDataModel ToData()

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

@@ -928,6 +928,7 @@ public async Task PostSetPasswordAsync_V2_WhenUserExistsAndSettingPasswordSuccee
     {
         // Arrange
         UpdateSetInitialPasswordRequestModelToV2(setInitialPasswordRequestModel);
+        _featureService.IsEnabled(FeatureFlagKeys.EnableAccountEncryptionV2JitPasswordRegistration).Returns(true);
         _userService.GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>()).Returns(Task.FromResult(user));
         _finishSsoJitProvisionMasterPasswordCommand.FinishProvisionAsync(user, Arg.Any<SetInitialMasterPasswordDataModel>())
             .Returns(Task.CompletedTask);
@@ -954,6 +955,7 @@ public async Task PostSetPasswordAsync_V2_WithTdeSetPassword_ShouldCallTdeSetPas
     {
         // Arrange
         UpdateSetInitialPasswordRequestModelToV2(setInitialPasswordRequestModel, includeTdeSetPassword: true);
+        _featureService.IsEnabled(FeatureFlagKeys.V2RegistrationTDEJIT).Returns(true);
         _userService.GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>()).Returns(Task.FromResult(user));
         _tdeSetPasswordCommand.SetMasterPasswordAsync(user, Arg.Any<SetInitialMasterPasswordDataModel>())
             .Returns(Task.CompletedTask);
@@ -993,6 +995,7 @@ public async Task PostSetPasswordAsync_V2_WhenSettingPasswordFails_ShouldThrowEx
     {
         // Arrange
         UpdateSetInitialPasswordRequestModelToV2(setInitialPasswordRequestModel);
+        _featureService.IsEnabled(FeatureFlagKeys.EnableAccountEncryptionV2JitPasswordRegistration).Returns(true);
         _userService.GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>()).Returns(Task.FromResult(user));
         _finishSsoJitProvisionMasterPasswordCommand.FinishProvisionAsync(user, Arg.Any<SetInitialMasterPasswordDataModel>())
             .Returns(Task.FromException(new Exception("Setting password failed")));
@@ -1001,6 +1004,105 @@ public async Task PostSetPasswordAsync_V2_WhenSettingPasswordFails_ShouldThrowEx
         await Assert.ThrowsAsync<Exception>(() => _sut.PostSetPasswordAsync(setInitialPasswordRequestModel));
     }
 
+    // V1 encryption with new data types (transitional path — V2 flags off, modern client carries MPAD/MPUD + legacy Keys)
+    // TODO: removal requires that BOTH flags have been removed:
+    //  - https://bitwarden.atlassian.net/browse/PM-27327 (MP)
+    //  - https://bitwarden.atlassian.net/browse/PM-27329 (TDE)
+    [Theory]
+    [BitAutoData]
+    public async Task PostSetPasswordAsync_V1_NewClientMpJit_UsesMpadMpudValues_ShouldCallV1CommandAsync(
+        User user,
+        SetInitialPasswordRequestModel setInitialPasswordRequestModel)
+    {
+        // Arrange — modern MP JIT client: sends MPAD + MPUD + legacy Keys (no AccountKeys, no V2 flag).
+        // ToUser() should map KDF from MPAD and the wrapped user key from MPUD; legacy Keys?.ToUser sets the keypair.
+        UpdateSetInitialPasswordRequestModelToV2(setInitialPasswordRequestModel);
+        setInitialPasswordRequestModel.AccountKeys = null;
+        setInitialPasswordRequestModel.Keys = new KeysRequestModel
+        {
+            PublicKey = "newPublicKey",
+            EncryptedPrivateKey = "newEncryptedPrivateKey"
+        };
+        user.PublicKey = null;
+        user.PrivateKey = null;
+
+        _userService.GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>()).Returns(Task.FromResult(user));
+        _setInitialMasterPasswordCommandV1.SetInitialMasterPasswordAsync(
+                user,
+                setInitialPasswordRequestModel.MasterPasswordAuthentication.MasterPasswordAuthenticationHash,
+                setInitialPasswordRequestModel.MasterPasswordUnlock.MasterKeyWrappedUserKey,
+                setInitialPasswordRequestModel.OrgIdentifier)
+            .Returns(Task.FromResult(IdentityResult.Success));
+
+        // Act
+        await _sut.PostSetPasswordAsync(setInitialPasswordRequestModel);
+
+        // Assert — V1 command called with MPAD hash + MPUD wrapped key (not legacy MasterPasswordHash/Key)
+        await _setInitialMasterPasswordCommandV1.Received(1)
+            .SetInitialMasterPasswordAsync(
+                Arg.Is<User>(u => u == user),
+                Arg.Is<string>(s => s == setInitialPasswordRequestModel.MasterPasswordAuthentication.MasterPasswordAuthenticationHash),
+                Arg.Is<string>(s => s == setInitialPasswordRequestModel.MasterPasswordUnlock.MasterKeyWrappedUserKey),
+                Arg.Is<string>(s => s == setInitialPasswordRequestModel.OrgIdentifier));
+
+        // KDF mapped from MPAD
+        Assert.Equal(setInitialPasswordRequestModel.MasterPasswordHint, user.MasterPasswordHint);
+        Assert.Equal(setInitialPasswordRequestModel.MasterPasswordAuthentication.Kdf.KdfType, user.Kdf);
+        Assert.Equal(setInitialPasswordRequestModel.MasterPasswordAuthentication.Kdf.Iterations, user.KdfIterations);
+
+        // Public/private keys mapped from legacy Keys
+        Assert.Equal("newPublicKey", user.PublicKey);
+        Assert.Equal("newEncryptedPrivateKey", user.PrivateKey);
+
+        // V2 commands not called
+        await _finishSsoJitProvisionMasterPasswordCommand.DidNotReceiveWithAnyArgs()
+            .FinishProvisionAsync(Arg.Any<User>(), Arg.Any<SetInitialMasterPasswordDataModel>());
+        await _tdeSetPasswordCommand.DidNotReceiveWithAnyArgs()
+            .SetMasterPasswordAsync(Arg.Any<User>(), Arg.Any<SetInitialMasterPasswordDataModel>());
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task PostSetPasswordAsync_V1_NewClientTde_UsesMpadMpudValues_DoesNotMutateExistingKeysAsync(
+        User user,
+        SetInitialPasswordRequestModel setInitialPasswordRequestModel)
+    {
+        // Arrange — modern TDE client: sends MPAD + MPUD with both AccountKeys and Keys null (V2 TDE flag off).
+        // TDE users already have a keypair; Keys?.ToUser is a no-op and the existing keys are left alone.
+        UpdateSetInitialPasswordRequestModelToV2(setInitialPasswordRequestModel, includeTdeSetPassword: true);
+
+        const string existingPublicKey = "tdeUserExistingPublicKey";
+        const string existingPrivateKey = "tdeUserExistingPrivateKey";
+        user.PublicKey = existingPublicKey;
+        user.PrivateKey = existingPrivateKey;
+
+        _userService.GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>()).Returns(Task.FromResult(user));
+        _setInitialMasterPasswordCommandV1.SetInitialMasterPasswordAsync(
+                Arg.Any<User>(), Arg.Any<string>(), Arg.Any<string>(), Arg.Any<string>())
+            .Returns(Task.FromResult(IdentityResult.Success));
+
+        // Act
+        await _sut.PostSetPasswordAsync(setInitialPasswordRequestModel);
+
+        // Assert — V1 command called with MPAD hash + MPUD wrapped key
+        await _setInitialMasterPasswordCommandV1.Received(1)
+            .SetInitialMasterPasswordAsync(
+                Arg.Is<User>(u => u == user),
+                Arg.Is<string>(s => s == setInitialPasswordRequestModel.MasterPasswordAuthentication.MasterPasswordAuthenticationHash),
+                Arg.Is<string>(s => s == setInitialPasswordRequestModel.MasterPasswordUnlock.MasterKeyWrappedUserKey),
+                Arg.Is<string>(s => s == setInitialPasswordRequestModel.OrgIdentifier));
+
+        // Existing keypair preserved
+        Assert.Equal(existingPublicKey, user.PublicKey);
+        Assert.Equal(existingPrivateKey, user.PrivateKey);
+
+        // V2 commands not called
+        await _finishSsoJitProvisionMasterPasswordCommand.DidNotReceiveWithAnyArgs()
+            .FinishProvisionAsync(Arg.Any<User>(), Arg.Any<SetInitialMasterPasswordDataModel>());
+        await _tdeSetPasswordCommand.DidNotReceiveWithAnyArgs()
+            .SetMasterPasswordAsync(Arg.Any<User>(), Arg.Any<SetInitialMasterPasswordDataModel>());
+    }
+
     private void UpdateSetInitialPasswordRequestModelToV1(SetInitialPasswordRequestModel model)
     {
         model.MasterPasswordAuthentication = null;