[PM-17210] Prevent unintentionally corrupting private keys (#5285)

* Prevent unintentionally corrupting private keys

* Deny key update only when replacing existing keys

* Fix incorrect use of existing user public/encrypted private key

* Fix test

* Fix tests

* Re-add test

* Pass through error for set-password

* Fix test

* Increase test coverage and simplify checks

src/Api/Auth/Controllers/AccountsController.cs

@@ -266,8 +266,18 @@ public class AccountsController : Controller
             throw new UnauthorizedAccessException();
         }
 
+        try
+        {
+            user = model.ToUser(user);
+        }
+        catch (Exception e)
+        {
+            ModelState.AddModelError(string.Empty, e.Message);
+            throw new BadRequestException(ModelState);
+        }
+
         var result = await _setInitialMasterPasswordCommand.SetInitialMasterPasswordAsync(
-            model.ToUser(user),
+            user,
             model.MasterPasswordHash,
             model.Key,
             model.OrgIdentifier);

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

@@ -1,26 +1,36 @@
 using System.ComponentModel.DataAnnotations;
 using Bit.Core.Entities;
+using Bit.Core.Utilities;
 
 namespace Bit.Core.Auth.Models.Api.Request.Accounts;
 
 public class KeysRequestModel
 {
+    [Required]
     public string PublicKey { get; set; }
     [Required]
     public string EncryptedPrivateKey { get; set; }
 
     public User ToUser(User existingUser)
     {
-        if (string.IsNullOrWhiteSpace(existingUser.PublicKey) && !string.IsNullOrWhiteSpace(PublicKey))
+        if (string.IsNullOrWhiteSpace(PublicKey) || string.IsNullOrWhiteSpace(EncryptedPrivateKey))
+        {
+            throw new InvalidOperationException("Public and private keys are required.");
+        }
+
+        if (string.IsNullOrWhiteSpace(existingUser.PublicKey) && string.IsNullOrWhiteSpace(existingUser.PrivateKey))
         {
             existingUser.PublicKey = PublicKey;
-        }
-
-        if (string.IsNullOrWhiteSpace(existingUser.PrivateKey))
-        {
             existingUser.PrivateKey = EncryptedPrivateKey;
+            return existingUser;
+        }
+        else if (PublicKey == existingUser.PublicKey && CoreHelpers.FixedTimeEquals(EncryptedPrivateKey, existingUser.PrivateKey))
+        {
+            return existingUser;
+        }
+        else
+        {
+            throw new InvalidOperationException("Cannot replace existing key(s) with new key(s).");
         }
-
-        return existingUser;
     }
 }

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

@@ -419,22 +419,32 @@ public class AccountsControllerTests : IDisposable
 
 
     [Theory]
-    [BitAutoData(true, false)]  // User has PublicKey and PrivateKey, and Keys in request are NOT null
+    [BitAutoData(true, "existingPrivateKey", "existingPublicKey", true)] // allow providing existing keys in the request
-    [BitAutoData(true, true)]   // User has PublicKey and PrivateKey, and Keys in request are null
+    [BitAutoData(true, null, null, true)] // allow not setting the public key when the user already has a key
-    [BitAutoData(false, false)] // User has neither PublicKey nor PrivateKey, and Keys in request are NOT null
+    [BitAutoData(false, "newPrivateKey", "newPublicKey", true)] // allow setting new keys when the user has no keys
-    [BitAutoData(false, true)]  // User has neither PublicKey nor PrivateKey, and Keys in request are null
+    [BitAutoData(false, null, null, true)] // allow not setting the public key when the user has no keys
+    // do not allow single key
+    [BitAutoData(false, "existingPrivateKey", null, false)]
+    [BitAutoData(false, null, "existingPublicKey", false)]
+    [BitAutoData(false, "newPrivateKey", null, false)]
+    [BitAutoData(false, null, "newPublicKey", false)]
+    [BitAutoData(true, "existingPrivateKey", null, false)]
+    [BitAutoData(true, null, "existingPublicKey", false)]
+    [BitAutoData(true, "newPrivateKey", null, false)]
+    [BitAutoData(true, null, "newPublicKey", false)]
+    // reject overwriting existing keys
+    [BitAutoData(true, "newPrivateKey", "newPublicKey", false)]
     public async Task PostSetPasswordAsync_WhenUserExistsAndSettingPasswordSucceeds_ShouldHandleKeysCorrectlyAndReturn(
-    bool hasExistingKeys,
+        bool hasExistingKeys,
-    bool shouldSetKeysToNull,
+        string requestPrivateKey,
-    User user,
+        string requestPublicKey,
-    SetPasswordRequestModel setPasswordRequestModel)
+        bool shouldSucceed,
+        User user,
+        SetPasswordRequestModel setPasswordRequestModel)
     {
         // Arrange
         const string existingPublicKey = "existingPublicKey";
-        const string existingEncryptedPrivateKey = "existingEncryptedPrivateKey";
+        const string existingEncryptedPrivateKey = "existingPrivateKey";
-
-        const string newPublicKey = "newPublicKey";
-        const string newEncryptedPrivateKey = "newEncryptedPrivateKey";
 
         if (hasExistingKeys)
         {
@@ -447,16 +457,16 @@ public class AccountsControllerTests : IDisposable
             user.PrivateKey = null;
         }
 
-        if (shouldSetKeysToNull)
+        if (requestPrivateKey == null && requestPublicKey == null)
         {
             setPasswordRequestModel.Keys = null;
         }
         else
         {
-            setPasswordRequestModel.Keys = new KeysRequestModel()
+            setPasswordRequestModel.Keys = new KeysRequestModel
             {
-                PublicKey = newPublicKey,
+                EncryptedPrivateKey = requestPrivateKey,
-                EncryptedPrivateKey = newEncryptedPrivateKey
+                PublicKey = requestPublicKey
             };
         }
 
@@ -469,44 +479,66 @@ public class AccountsControllerTests : IDisposable
             .Returns(Task.FromResult(IdentityResult.Success));
 
         // Act
-        await _sut.PostSetPasswordAsync(setPasswordRequestModel);
+        if (shouldSucceed)
-
-        // Assert
-        await _setInitialMasterPasswordCommand.Received(1)
-            .SetInitialMasterPasswordAsync(
-                Arg.Is<User>(u => u == user),
-                Arg.Is<string>(s => s == setPasswordRequestModel.MasterPasswordHash),
-                Arg.Is<string>(s => s == setPasswordRequestModel.Key),
-                Arg.Is<string>(s => s == setPasswordRequestModel.OrgIdentifier));
-
-        // Additional Assertions for User object modifications
-        Assert.Equal(setPasswordRequestModel.MasterPasswordHint, user.MasterPasswordHint);
-        Assert.Equal(setPasswordRequestModel.Kdf, user.Kdf);
-        Assert.Equal(setPasswordRequestModel.KdfIterations, user.KdfIterations);
-        Assert.Equal(setPasswordRequestModel.KdfMemory, user.KdfMemory);
-        Assert.Equal(setPasswordRequestModel.KdfParallelism, user.KdfParallelism);
-        Assert.Equal(setPasswordRequestModel.Key, user.Key);
-
-        if (hasExistingKeys)
         {
-            // User Keys should not be modified
+            await _sut.PostSetPasswordAsync(setPasswordRequestModel);
-            Assert.Equal(existingPublicKey, user.PublicKey);
+            // Assert
-            Assert.Equal(existingEncryptedPrivateKey, user.PrivateKey);
+            await _setInitialMasterPasswordCommand.Received(1)
-        }
+                .SetInitialMasterPasswordAsync(
-        else if (!shouldSetKeysToNull)
+                    Arg.Is<User>(u => u == user),
-        {
+                    Arg.Is<string>(s => s == setPasswordRequestModel.MasterPasswordHash),
-            // User had no keys so they should be set to the request model's keys
+                    Arg.Is<string>(s => s == setPasswordRequestModel.Key),
-            Assert.Equal(setPasswordRequestModel.Keys.PublicKey, user.PublicKey);
+                    Arg.Is<string>(s => s == setPasswordRequestModel.OrgIdentifier));
-            Assert.Equal(setPasswordRequestModel.Keys.EncryptedPrivateKey, user.PrivateKey);
+
+            // Additional Assertions for User object modifications
+            Assert.Equal(setPasswordRequestModel.MasterPasswordHint, user.MasterPasswordHint);
+            Assert.Equal(setPasswordRequestModel.Kdf, user.Kdf);
+            Assert.Equal(setPasswordRequestModel.KdfIterations, user.KdfIterations);
+            Assert.Equal(setPasswordRequestModel.KdfMemory, user.KdfMemory);
+            Assert.Equal(setPasswordRequestModel.KdfParallelism, user.KdfParallelism);
+            Assert.Equal(setPasswordRequestModel.Key, user.Key);
         }
         else
         {
-            // User had no keys and the request model's keys were null, so they should be set to null
+            await Assert.ThrowsAsync<BadRequestException>(() => _sut.PostSetPasswordAsync(setPasswordRequestModel));
-            Assert.Null(user.PublicKey);
-            Assert.Null(user.PrivateKey);
         }
     }
 
+    [Theory]
+    [BitAutoData]
+    public async Task PostSetPasswordAsync_WhenUserExistsAndHasKeysAndKeysAreUpdated_ShouldThrowAsync(
+    User user,
+    SetPasswordRequestModel setPasswordRequestModel)
+    {
+        // Arrange
+        const string existingPublicKey = "existingPublicKey";
+        const string existingEncryptedPrivateKey = "existingEncryptedPrivateKey";
+
+        const string newPublicKey = "newPublicKey";
+        const string newEncryptedPrivateKey = "newEncryptedPrivateKey";
+
+        user.PublicKey = existingPublicKey;
+        user.PrivateKey = existingEncryptedPrivateKey;
+
+        setPasswordRequestModel.Keys = new KeysRequestModel()
+        {
+            PublicKey = newPublicKey,
+            EncryptedPrivateKey = newEncryptedPrivateKey
+        };
+
+        _userService.GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>()).Returns(Task.FromResult(user));
+        _setInitialMasterPasswordCommand.SetInitialMasterPasswordAsync(
+                user,
+                setPasswordRequestModel.MasterPasswordHash,
+                setPasswordRequestModel.Key,
+                setPasswordRequestModel.OrgIdentifier)
+            .Returns(Task.FromResult(IdentityResult.Success));
+
+        // Act & Assert
+        await Assert.ThrowsAsync<BadRequestException>(() => _sut.PostSetPasswordAsync(setPasswordRequestModel));
+    }
+
+
     [Theory]
     [BitAutoData]
     public async Task PostSetPasswordAsync_WhenUserDoesNotExist_ShouldThrowUnauthorizedAccessException(
@@ -525,6 +557,7 @@ public class AccountsControllerTests : IDisposable
         User user,
         SetPasswordRequestModel model)
     {
+        model.Keys = null;
         // Arrange
         _userService.GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>()).Returns(Task.FromResult(user));
         _setInitialMasterPasswordCommand.SetInitialMasterPasswordAsync(Arg.Any<User>(), Arg.Any<string>(), Arg.Any<string>(), Arg.Any<string>())