Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. Weโ€™ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
quexten wants to merge 1 commit into
base: main
Choose a base branch
from
+165 โˆ’16
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 41 additions & 4 deletions src/Api/Auth/Controllers/AccountsController.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
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;
Expand All @@ -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,
Expand All @@ -57,7 +59,8 @@ public AccountsController(
IFeatureService featureService,
IUserAccountKeysQuery userAccountKeysQuery,
ITwoFactorEmailService twoFactorEmailService,
IChangeKdfCommand changeKdfCommand
IChangeKdfCommand changeKdfCommand,
IUserRepository userRepository
)
{
_organizationService = organizationService;
Expand All @@ -72,6 +75,7 @@ IChangeKdfCommand changeKdfCommand
_userAccountKeysQuery = userAccountKeysQuery;
_twoFactorEmailService = twoFactorEmailService;
_changeKdfCommand = changeKdfCommand;
_userRepository = userRepository;
}


Expand Down Expand Up @@ -440,8 +444,40 @@ public async Task<KeysResponseModel> PostKeys([FromBody] KeysRequestModel model)
}
}

await _userService.SaveUserAsync(model.ToUser(user));
return new KeysResponseModel(user);
if (model.AccountKeys != null)
{
var accountKeysData = model.AccountKeys.ToAccountKeysData();
if (accountKeysData.IsV2Encryption())
{
await _userRepository.SetV2AccountCryptographicStateAsync(user.Id, accountKeysData);
return new KeysResponseModel(accountKeysData, user.Key);
Comment on lines +450 to +453
Copy link
Contributor

@claude claude bot Dec 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

โœ… GOOD: Properly validating V2 encryption state with IsV2Encryption() before calling SetV2AccountCryptographicStateAsync. This prevents invalid cryptographic states from being persisted.

The repository method also performs this validation, providing defense-in-depth.

}
else
{
// Todo: Drop this after a transition period
await _userService.SaveUserAsync(model.ToUser(user));
return new KeysResponseModel(new UserAccountKeysData
{
PublicKeyEncryptionKeyPairData = new PublicKeyEncryptionKeyPairData(
user.PrivateKey,
user.PublicKey
)
}, user.Key);
Comment on lines +458 to +465
Copy link
Contributor

@claude claude bot Dec 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

โš ๏ธ DATA INCONSISTENCY RISK: After calling SaveUserAsync, the user object in memory may not reflect what was actually persisted to the database.

The SaveUserAsync method calls model.ToUser(user) which modifies the user object, but there's no guarantee that the returned user.PrivateKey and user.PublicKey match what was saved. This could happen if:

  1. Database triggers modify the values
  2. The save operation fails partially
  3. There are concurrent updates

Recommendation: Reload the user from the database after saving:

var updatedUser = model.ToUser(user);
await _userService.SaveUserAsync(updatedUser);
// Reload to ensure we return what was actually saved
user = await _userService.GetUserByIdAsync(updatedUser.Id);
return new KeysResponseModel(new UserAccountKeysData
{
    PublicKeyEncryptionKeyPairData = new PublicKeyEncryptionKeyPairData(
        user.PrivateKey,
        user.PublicKey
    )
}, user.Key);

Alternatively, use the updatedUser reference:

var updatedUser = model.ToUser(user);
await _userService.SaveUserAsync(updatedUser);
return new KeysResponseModel(new UserAccountKeysData
{
    PublicKeyEncryptionKeyPairData = new PublicKeyEncryptionKeyPairData(
        updatedUser.PrivateKey,
        updatedUser.PublicKey
    )
}, updatedUser.Key);

}
}
else
{
// Todo: Drop this after a transition period
await _userService.SaveUserAsync(model.ToUser(user));
return new KeysResponseModel(new UserAccountKeysData
{
PublicKeyEncryptionKeyPairData = new PublicKeyEncryptionKeyPairData(
user.PrivateKey,
user.PublicKey
)
}, user.Key);
}
Comment on lines +455 to +479
Copy link
Contributor

@claude claude bot Dec 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

โ™ป๏ธ CODE DUPLICATION: The v1 fallback logic (lines 458-466 and 471-479) is duplicated. This makes the code harder to maintain and increases the risk of bugs.

Recommendation: Extract the common v1 handling logic:

if (model.AccountKeys != null)
{
    var accountKeysData = model.AccountKeys.ToAccountKeysData();
    if (accountKeysData.IsV2Encryption())
    {
        await _userRepository.SetV2AccountCryptographicStateAsync(user.Id, accountKeysData);
        return new KeysResponseModel(accountKeysData, user.Key);
    }
}

// V1 fallback path (handles both null AccountKeys and non-V2 AccountKeys)
var updatedUser = model.ToUser(user);
await _userService.SaveUserAsync(updatedUser);
return new KeysResponseModel(new UserAccountKeysData
{
    PublicKeyEncryptionKeyPairData = new PublicKeyEncryptionKeyPairData(
        updatedUser.PrivateKey,
        updatedUser.PublicKey
    )
}, updatedUser.Key);

This simplifies the logic and makes it clear that there are two paths: V2 and V1 (fallback).


}

[HttpGet("keys")]
Expand All @@ -453,7 +489,8 @@ public async Task<KeysResponseModel> GetKeys()
throw new UnauthorizedAccessException();
}

return new KeysResponseModel(user);
var accountKeys = await _userAccountKeysQuery.Run(user);
return new KeysResponseModel(accountKeys, user.Key);
}

[HttpDelete]
Expand Down
27 changes: 16 additions & 11 deletions src/Api/Models/Response/KeysResponseModel.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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);
Comment on lines +9 to +19
Copy link
Contributor

@claude claude bot Dec 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

๐Ÿ‘ GOOD REFACTORING: The new constructor properly:

  1. Accepts UserAccountKeysData instead of the full User entity, following separation of concerns
  2. Makes masterKeyWrappedUserKey nullable and optional
  3. Populates both legacy properties (for backward compatibility) and new AccountKeys property
  4. Removes the #nullable disable directive

This is a clean migration path that maintains backward compatibility while supporting V2 encryption.

}

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; }
}
5 changes: 5 additions & 0 deletions src/Core/Auth/Models/Api/Request/Accounts/KeysRequestModel.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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; }
Copy link
Contributor

@claude claude bot Dec 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

โš ๏ธ MISSING VALIDATION: The AccountKeys property has no validation attributes, while the obsolete properties have [Required].

This creates an ambiguous state where:

  • Old clients send only PublicKey and EncryptedPrivateKey (both marked [Required])
  • New clients send only AccountKeys (no validation)
  • But a request could potentially provide neither, or both

The current #nullable disable at the top of the file masks this issue.

Recommendation: Add validation to ensure at least one set of keys is provided:

public class KeysRequestModel : IValidatableObject
{
    [Obsolete("Use AccountKeys.AccountPublicKey instead")]
    public string? PublicKey { get; set; }
    
    [Obsolete("Use AccountKeys.UserKeyEncryptedAccountPrivateKey instead")]
    public string? EncryptedPrivateKey { get; set; }
    
    public AccountKeysRequestModel? AccountKeys { get; set; }

    public IEnumerable<ValidationResult> Validate(ValidationContext validationContext)
    {
        // Either new AccountKeys or both legacy fields must be provided
        if (AccountKeys == null && 
            (string.IsNullOrWhiteSpace(PublicKey) || string.IsNullOrWhiteSpace(EncryptedPrivateKey)))
        {
            yield return new ValidationResult(
                "Either AccountKeys or both PublicKey and EncryptedPrivateKey must be provided.",
                new[] { nameof(AccountKeys), nameof(PublicKey), nameof(EncryptedPrivateKey) });
        }
    }
}

This would also allow removing the #nullable disable directive and the now-incorrect [Required] attributes.


[Obsolete("Use SetAccountKeysForUserCommand instead")]
public User ToUser(User existingUser)
{
if (string.IsNullOrWhiteSpace(PublicKey) || string.IsNullOrWhiteSpace(EncryptedPrivateKey))
Expand Down
1 change: 1 addition & 0 deletions src/Core/Constants.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,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 */
Expand Down
80 changes: 79 additions & 1 deletion test/Api.Test/Auth/Controllers/AccountsControllerTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
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;
Expand Down Expand Up @@ -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()
{
Expand All @@ -53,6 +55,7 @@ public AccountsControllerTests()
_userAccountKeysQuery = Substitute.For<IUserAccountKeysQuery>();
_twoFactorEmailService = Substitute.For<ITwoFactorEmailService>();
_changeKdfCommand = Substitute.For<IChangeKdfCommand>();
_userRepository = Substitute.For<IUserRepository>();

_sut = new AccountsController(
_organizationService,
Expand All @@ -66,7 +69,8 @@ public AccountsControllerTests()
_featureService,
_userAccountKeysQuery,
_twoFactorEmailService,
_changeKdfCommand
_changeKdfCommand,
_userRepository
);
}

Expand Down Expand Up @@ -738,5 +742,79 @@ private void ConfigureUserServiceToReturnNullUserId()
_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);
}
Comment on lines +788 to +790
Copy link
Contributor

@claude claude bot Dec 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

๐Ÿ“ TEST COVERAGE GAP: This test only validates that the result is not null and has the correct object type. It doesn't verify:

  1. The response contains the correct key data
  2. The AccountKeys property is properly populated
  3. The Key property matches user.Key

Recommendation: Add comprehensive response validation:

// 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);
Assert.Equal(user.Key, result.Key);
Assert.NotNull(result.AccountKeys);
Assert.NotNull(result.PublicKey);
Assert.NotNull(result.PrivateKey);

Similar validation should be added to the PostKeys_WithoutAccountKeys_CallsSaveUser test as well.


[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);
}
Comment on lines +816 to +818
Copy link
Contributor

@claude claude bot Dec 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

๐Ÿ“ TEST COVERAGE GAP: Similar to the V2 test, this test needs better response validation to catch potential bugs in response construction.

Recommendation: Add assertions to verify the response data:

// 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);
Assert.Equal(user.Key, result.Key);
Assert.NotNull(result.AccountKeys);
Assert.Equal(model.PublicKey, result.PublicKey);
Assert.Equal(model.EncryptedPrivateKey, result.PrivateKey);

Note: This assumes the controller is fixed to use the updated user reference (see my comment on AccountsController.cs).

}

6 changes: 6 additions & 0 deletions test/Identity.IntegrationTest/Controllers/AccountsControllerTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,7 @@ public async Task RegistrationWithEmailVerification_WithEmailVerificationToken_S
[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();

Expand Down Expand Up @@ -202,6 +203,7 @@ public async Task RegistrationWithEmailVerification_OpenRegistrationDisabled_Thr
[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");
Expand Down Expand Up @@ -233,6 +235,7 @@ public async Task RegistrationWithEmailVerification_WithOrgInviteToken_Succeeds(
[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();
Expand Down Expand Up @@ -310,6 +313,7 @@ public async Task RegistrationWithEmailVerification_WithOrgSponsoredFreeFamilyPl
[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();
Expand Down Expand Up @@ -386,6 +390,7 @@ public async Task RegistrationWithEmailVerification_WithAcceptEmergencyAccessInv
[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();
Expand Down Expand Up @@ -455,6 +460,7 @@ public async Task RegistrationWithEmailVerification_WithProviderInviteToken_Succ
[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();
Expand Down
Loading
Loading