From 7a490701cd3d144515f64783bfbcfc60dc34a18b Mon Sep 17 00:00:00 2001 From: Maciej Zieniuk Date: Fri, 7 Aug 2026 12:41:53 +0100 Subject: [PATCH 1/3] Reuse the connection and transaction during user key rotation Replace UpdateEncryptedDataForKeyRotation with the generic DatabaseTransactionAction, which takes a DbConnection and DbTransaction rather than the SQL Server specific SqlConnection and SqlTransaction. This matches the delegate introduced in #7548. Every EF repository that takes part in a key rotation now binds its DatabaseContext to the caller's connection and transaction through GetTransactionalDatabaseContext, instead of opening a second connection. Opening a second connection deadlocks on single-writer providers such as SQLite, and left each write outside the rotation transaction on the others, which breaks integration tests. --- .../IOrganizationUserRepository.cs | 3 +- .../IEmergencyAccessRepository.cs | 3 +- .../IWebAuthnCredentialRepository.cs | 3 +- .../IUserSignatureKeyPairRepository.cs | 5 +- .../UserKey/IRotateUserAccountKeysCommand.cs | 10 ---- .../RotateUserAccountKeysCommand.cs | 20 ++++---- .../Repositories/DatabaseTransactionAction.cs | 9 ++++ src/Core/Repositories/IDeviceRepository.cs | 3 +- src/Core/Repositories/IUserRepository.cs | 5 +- .../Tools/Repositories/ISendRepository.cs | 3 +- .../Vault/Repositories/ICipherRepository.cs | 3 +- .../Vault/Repositories/IFolderRepository.cs | 5 +- .../OrganizationUserRepository.cs | 3 +- .../Repositories/EmergencyAccessRepository.cs | 8 ++-- .../WebAuthnCredentialRepository.cs | 6 +-- .../UserSignatureKeyPairRepository.cs | 10 ++-- .../Repositories/DeviceRepository.cs | 5 +- .../Repositories/UserRepository.cs | 5 +- .../Tools/Repositories/SendRepository.cs | 9 ++-- .../Vault/Repositories/CipherRepository.cs | 9 ++-- .../Vault/Repositories/FolderRepository.cs | 9 ++-- .../OrganizationUserRepository.cs | 7 ++- .../Repositories/EmergencyAccessRepository.cs | 8 ++-- .../WebAuthnCredentialRepository.cs | 8 ++-- .../UserSignatureKeyPairRepository.cs | 14 +++--- .../BaseEntityFrameworkRepository.cs | 12 ++++- .../Repositories/DeviceRepository.cs | 7 ++- .../Repositories/UserRepository.cs | 21 ++++---- .../Tools/Repositories/SendRepository.cs | 8 ++-- .../Vault/Repositories/CipherRepository.cs | 9 ++-- .../Vault/Repositories/FolderRepository.cs | 9 ++-- .../RotateUserAccountKeysCommandTests.cs | 47 +++++++++--------- .../PushServiceCollectionExtensionsTests.cs | 3 +- .../DatabaseTransactionActionTestHelper.cs | 48 +++++++++++++++++++ .../Repositories/UserRepositoryTests.cs | 25 +++++++++- 35 files changed, 217 insertions(+), 145 deletions(-) create mode 100644 src/Core/Repositories/DatabaseTransactionAction.cs create mode 100644 test/Infrastructure.IntegrationTest/DatabaseTransactionActionTestHelper.cs diff --git a/src/Core/AdminConsole/Repositories/IOrganizationUserRepository.cs b/src/Core/AdminConsole/Repositories/IOrganizationUserRepository.cs index 5b021831097a..52f55c06b905 100644 --- a/src/Core/AdminConsole/Repositories/IOrganizationUserRepository.cs +++ b/src/Core/AdminConsole/Repositories/IOrganizationUserRepository.cs @@ -3,7 +3,6 @@ using Bit.Core.AdminConsole.OrganizationFeatures.OrganizationUsers.InviteUsers.Models; using Bit.Core.Entities; using Bit.Core.Enums; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Models.Data; using Bit.Core.Models.Data.Organizations.OrganizationUsers; @@ -88,7 +87,7 @@ Task> GetManyDetailsByUserAsync /// /// The user that initiated the key rotation /// A list of organization users with updated reset password keys - UpdateEncryptedDataForKeyRotation UpdateForKeyRotation(Guid userId, + DatabaseTransactionAction UpdateForKeyRotation(Guid userId, IEnumerable resetPasswordKeys); /// diff --git a/src/Core/Auth/Repositories/IEmergencyAccessRepository.cs b/src/Core/Auth/Repositories/IEmergencyAccessRepository.cs index e46a0520e5fe..02dafc8ecf9c 100644 --- a/src/Core/Auth/Repositories/IEmergencyAccessRepository.cs +++ b/src/Core/Auth/Repositories/IEmergencyAccessRepository.cs @@ -1,6 +1,5 @@ using Bit.Core.Auth.Entities; using Bit.Core.Auth.Models.Data; -using Bit.Core.KeyManagement.UserKey; namespace Bit.Core.Repositories; @@ -40,7 +39,7 @@ public interface IEmergencyAccessRepository : IRepository /// /// The grantor that initiated the key rotation /// A list of emergency access with updated keys - UpdateEncryptedDataForKeyRotation UpdateForKeyRotation(Guid grantorId, + DatabaseTransactionAction UpdateForKeyRotation(Guid grantorId, IEnumerable emergencyAccessKeys); /// diff --git a/src/Core/Auth/Repositories/IWebAuthnCredentialRepository.cs b/src/Core/Auth/Repositories/IWebAuthnCredentialRepository.cs index 29ed9d2210bf..859a6f8ed6b1 100644 --- a/src/Core/Auth/Repositories/IWebAuthnCredentialRepository.cs +++ b/src/Core/Auth/Repositories/IWebAuthnCredentialRepository.cs @@ -1,6 +1,5 @@ using Bit.Core.Auth.Entities; using Bit.Core.Auth.Models.Data; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Repositories; #nullable enable @@ -12,5 +11,5 @@ public interface IWebAuthnCredentialRepository : IRepository GetByIdAsync(Guid id, Guid userId); Task> GetManyByUserIdAsync(Guid userId); Task UpdateAsync(WebAuthnCredential credential); - UpdateEncryptedDataForKeyRotation UpdateKeysForRotationAsync(Guid userId, IEnumerable credentials); + DatabaseTransactionAction UpdateKeysForRotationAsync(Guid userId, IEnumerable credentials); } diff --git a/src/Core/KeyManagement/Repositories/IUserSignatureKeyPairRepository.cs b/src/Core/KeyManagement/Repositories/IUserSignatureKeyPairRepository.cs index ce8979620f8e..bfb20332e24b 100644 --- a/src/Core/KeyManagement/Repositories/IUserSignatureKeyPairRepository.cs +++ b/src/Core/KeyManagement/Repositories/IUserSignatureKeyPairRepository.cs @@ -1,7 +1,6 @@  using Bit.Core.KeyManagement.Entities; using Bit.Core.KeyManagement.Models.Data; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Repositories; namespace Bit.Core.KeyManagement.Repositories; @@ -9,6 +8,6 @@ namespace Bit.Core.KeyManagement.Repositories; public interface IUserSignatureKeyPairRepository : IRepository { public Task GetByUserIdAsync(Guid userId); - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation(Guid grantorId, SignatureKeyPairData signatureKeyPair); - public UpdateEncryptedDataForKeyRotation SetUserSignatureKeyPair(Guid userId, SignatureKeyPairData signatureKeyPair); + public DatabaseTransactionAction UpdateForKeyRotation(Guid grantorId, SignatureKeyPairData signatureKeyPair); + public DatabaseTransactionAction SetUserSignatureKeyPair(Guid userId, SignatureKeyPairData signatureKeyPair); } diff --git a/src/Core/KeyManagement/UserKey/IRotateUserAccountKeysCommand.cs b/src/Core/KeyManagement/UserKey/IRotateUserAccountKeysCommand.cs index 14c9974f0083..3b80076e6dc4 100644 --- a/src/Core/KeyManagement/UserKey/IRotateUserAccountKeysCommand.cs +++ b/src/Core/KeyManagement/UserKey/IRotateUserAccountKeysCommand.cs @@ -5,7 +5,6 @@ using Bit.Core.Exceptions; using Bit.Core.KeyManagement.UserKey.Models.Data; using Microsoft.AspNetCore.Identity; -using Microsoft.Data.SqlClient; namespace Bit.Core.KeyManagement.UserKey; @@ -51,12 +50,3 @@ public interface IRotateUserAccountKeysCommand /// Thrown when is not a key connector user. Task KeyConnectorRotateUserAccountKeysAsync(User user, KeyConnectorRotateUserAccountKeysData model); } - -/// -/// A type used to implement updates to the database for key rotations. Each domain that requires an update of encrypted -/// data during a key rotation should use this to implement its own database call. The user repository loops through -/// these during a key rotation. -/// Note: connection and transaction are only used for Dapper. They won't be available in EF -/// -public delegate Task UpdateEncryptedDataForKeyRotation(SqlConnection connection = null, - SqlTransaction transaction = null); diff --git a/src/Core/KeyManagement/UserKey/Implementations/RotateUserAccountKeysCommand.cs b/src/Core/KeyManagement/UserKey/Implementations/RotateUserAccountKeysCommand.cs index 3b3cb4351832..b06614937878 100644 --- a/src/Core/KeyManagement/UserKey/Implementations/RotateUserAccountKeysCommand.cs +++ b/src/Core/KeyManagement/UserKey/Implementations/RotateUserAccountKeysCommand.cs @@ -89,7 +89,7 @@ public async Task PasswordChangeAndRotateUserAccountKeysAsync(Us model.ValidateForUser(user); - List saveEncryptedDataActions = []; + List saveEncryptedDataActions = []; var shouldPersistV2UpgradeToken = await BaseRotateUserAccountKeysAsync(model.BaseData, user, saveEncryptedDataActions); // Delegate the master password mutation (hash, wrapped user key, hint, time markers) to @@ -124,7 +124,7 @@ public async Task MasterPasswordRotateUserAccountKeysAsync(User user, MasterPass model.ValidateForUser(user); - List saveEncryptedDataActions = []; + List saveEncryptedDataActions = []; var shouldPersistV2UpgradeToken = await BaseRotateUserAccountKeysAsync(model.BaseData, user, saveEncryptedDataActions); user.Key = model.MasterPasswordUnlockData.MasterKeyWrappedUserKey; @@ -141,7 +141,7 @@ public async Task TdeRotateUserAccountKeysAsync(User user, TdeRotateUserAccountK model.ValidateForUser(user); - List saveEncryptedDataActions = []; + List saveEncryptedDataActions = []; var shouldPersistV2UpgradeToken = await BaseRotateUserAccountKeysAsync(model.BaseData, user, saveEncryptedDataActions); await _userRepository.UpdateUserKeyAndEncryptedDataV2Async(user, saveEncryptedDataActions); @@ -156,7 +156,7 @@ public async Task KeyConnectorRotateUserAccountKeysAsync(User user, KeyConnector model.ValidateForUser(user); - List saveEncryptedDataActions = []; + List saveEncryptedDataActions = []; var shouldPersistV2UpgradeToken = await BaseRotateUserAccountKeysAsync(model.BaseData, user, saveEncryptedDataActions); user.Key = model.KeyConnectorKeyWrappedUserKey; @@ -165,7 +165,7 @@ public async Task KeyConnectorRotateUserAccountKeysAsync(User user, KeyConnector await HandlePushNotificationAsync(shouldPersistV2UpgradeToken, user); } - private async Task RotateV2AccountKeysAsync(BaseRotateUserAccountKeysData model, User user, List saveEncryptedDataActions) + private async Task RotateV2AccountKeysAsync(BaseRotateUserAccountKeysData model, User user, List saveEncryptedDataActions) { ValidateV2Encryption(model); await ValidateVerifyingKeyUnchangedAsync(model, user); @@ -176,7 +176,7 @@ private async Task RotateV2AccountKeysAsync(BaseRotateUserAccountKeysData model, user.SecurityVersion = model.AccountKeys.SecurityStateData.SecurityVersion; } - private void UpgradeV1ToV2Keys(BaseRotateUserAccountKeysData model, User user, List saveEncryptedDataActions) + private void UpgradeV1ToV2Keys(BaseRotateUserAccountKeysData model, User user, List saveEncryptedDataActions) { ValidateV2Encryption(model); saveEncryptedDataActions.Add(_userSignatureKeyPairRepository.SetUserSignatureKeyPair(user.Id, model.AccountKeys.SignatureKeyPairData)); @@ -185,7 +185,7 @@ private void UpgradeV1ToV2Keys(BaseRotateUserAccountKeysData model, User user, L user.SecurityVersion = model.AccountKeys.SecurityStateData.SecurityVersion; } - internal async Task UpdateAccountKeysAsync(BaseRotateUserAccountKeysData model, User user, List saveEncryptedDataActions) + internal async Task UpdateAccountKeysAsync(BaseRotateUserAccountKeysData model, User user, List saveEncryptedDataActions) { ValidatePublicKeyEncryptionKeyPairUnchanged(model, user); @@ -210,7 +210,7 @@ internal async Task UpdateAccountKeysAsync(BaseRotateUserAccountKeysData model, user.PrivateKey = model.AccountKeys.PublicKeyEncryptionKeyPairData.WrappedPrivateKey; } - internal void UpdateUserData(BaseRotateUserAccountKeysData model, User user, List saveEncryptedDataActions) + internal void UpdateUserData(BaseRotateUserAccountKeysData model, User user, List saveEncryptedDataActions) { // The revision date has to be updated so that de-synced clients don't accidentally post over the re-encrypted data // with an old-user key-encrypted copy @@ -290,7 +290,7 @@ private static void ValidateV2Encryption(BaseRotateUserAccountKeysData model) } } - private void UpdateBaseUnlockMethods(BaseRotateUserAccountKeysData model, User user, List saveEncryptedDataActions) + private void UpdateBaseUnlockMethods(BaseRotateUserAccountKeysData model, User user, List saveEncryptedDataActions) { if (model.EmergencyAccesses.Any()) { @@ -314,7 +314,7 @@ private void UpdateBaseUnlockMethods(BaseRotateUserAccountKeysData model, User u } private async Task BaseRotateUserAccountKeysAsync(BaseRotateUserAccountKeysData baseModel, User user, - List saveEncryptedDataActions) + List saveEncryptedDataActions) { var now = DateTime.UtcNow; user.RevisionDate = user.AccountRevisionDate = now; diff --git a/src/Core/Repositories/DatabaseTransactionAction.cs b/src/Core/Repositories/DatabaseTransactionAction.cs new file mode 100644 index 000000000000..ae4aa92a19f4 --- /dev/null +++ b/src/Core/Repositories/DatabaseTransactionAction.cs @@ -0,0 +1,9 @@ +using System.Data.Common; + +namespace Bit.Core.Repositories; + +/// +/// A database operation that participates in an existing database connection and transaction. +/// Used to compose multiple repository operations into a single atomic transaction. +/// +public delegate Task DatabaseTransactionAction(DbConnection connection, DbTransaction transaction); diff --git a/src/Core/Repositories/IDeviceRepository.cs b/src/Core/Repositories/IDeviceRepository.cs index d4ce7a299a72..235b0438ecf2 100644 --- a/src/Core/Repositories/IDeviceRepository.cs +++ b/src/Core/Repositories/IDeviceRepository.cs @@ -1,6 +1,5 @@ using Bit.Core.Auth.Models.Data; using Bit.Core.Entities; -using Bit.Core.KeyManagement.UserKey; #nullable enable @@ -17,7 +16,7 @@ public interface IDeviceRepository : IRepository // other requests. Task> GetManyByUserIdWithDeviceAuth(Guid userId); Task ClearPushTokenAsync(Guid id); - UpdateEncryptedDataForKeyRotation UpdateKeysForRotationAsync(Guid userId, IEnumerable devices); + DatabaseTransactionAction UpdateKeysForRotationAsync(Guid userId, IEnumerable devices); /// /// Updates the device's last-activity state: moves LastActivityDate to today (if not /// already today) and writes to ClientVersion (if diff --git a/src/Core/Repositories/IUserRepository.cs b/src/Core/Repositories/IUserRepository.cs index c7f5bc1668f0..92e3a5748b24 100644 --- a/src/Core/Repositories/IUserRepository.cs +++ b/src/Core/Repositories/IUserRepository.cs @@ -2,7 +2,6 @@ using Bit.Core.Billing.Premium.Models; using Bit.Core.Entities; using Bit.Core.KeyManagement.Models.Data; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Models.Data; namespace Bit.Core.Repositories; @@ -60,9 +59,9 @@ public interface IUserRepository : IRepository /// The user to update /// Registered database calls to update re-encrypted data. Task UpdateUserKeyAndEncryptedDataAsync(User user, - IEnumerable updateDataActions); + IEnumerable updateDataActions); Task UpdateUserKeyAndEncryptedDataV2Async(User user, - IEnumerable updateDataActions); + IEnumerable updateDataActions); /// /// Sets the account cryptographic state to a user in a single transaction. The provided /// MUST be a V2 encryption state. Passing in a V1 encryption state will throw. diff --git a/src/Core/Tools/Repositories/ISendRepository.cs b/src/Core/Tools/Repositories/ISendRepository.cs index 4f7ced15df5e..55c40bda9d6a 100644 --- a/src/Core/Tools/Repositories/ISendRepository.cs +++ b/src/Core/Tools/Repositories/ISendRepository.cs @@ -1,6 +1,5 @@ #nullable enable -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Repositories; using Bit.Core.Tools.Entities; @@ -76,7 +75,7 @@ public interface ISendRepository : IRepository /// /// The user that initiated the key rotation /// A list of sends with updated data - UpdateEncryptedDataForKeyRotation UpdateForKeyRotation(Guid userId, + DatabaseTransactionAction UpdateForKeyRotation(Guid userId, IEnumerable sends); /// diff --git a/src/Core/Vault/Repositories/ICipherRepository.cs b/src/Core/Vault/Repositories/ICipherRepository.cs index bc4bd64cfa7f..4c15e97a387c 100644 --- a/src/Core/Vault/Repositories/ICipherRepository.cs +++ b/src/Core/Vault/Repositories/ICipherRepository.cs @@ -1,5 +1,4 @@ using Bit.Core.Entities; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Repositories; using Bit.Core.Vault.Entities; using Bit.Core.Vault.Models.Data; @@ -68,7 +67,7 @@ Task> GetCipherPermissionsForOrganizat /// /// The user that initiated the key rotation /// A list of ciphers with updated data - UpdateEncryptedDataForKeyRotation UpdateForKeyRotation(Guid userId, + DatabaseTransactionAction UpdateForKeyRotation(Guid userId, IEnumerable ciphers); /// diff --git a/src/Core/Vault/Repositories/IFolderRepository.cs b/src/Core/Vault/Repositories/IFolderRepository.cs index c4693b2a1343..9073706e5c74 100644 --- a/src/Core/Vault/Repositories/IFolderRepository.cs +++ b/src/Core/Vault/Repositories/IFolderRepository.cs @@ -1,5 +1,4 @@ -using Bit.Core.KeyManagement.UserKey; -using Bit.Core.Repositories; +using Bit.Core.Repositories; using Bit.Core.Vault.Entities; namespace Bit.Core.Vault.Repositories; @@ -14,6 +13,6 @@ public interface IFolderRepository : IRepository /// /// The user that initiated the key rotation /// A list of folders with updated data - UpdateEncryptedDataForKeyRotation UpdateForKeyRotation(Guid userId, + DatabaseTransactionAction UpdateForKeyRotation(Guid userId, IEnumerable folders); } diff --git a/src/Infrastructure.Dapper/AdminConsole/Repositories/OrganizationUserRepository.cs b/src/Infrastructure.Dapper/AdminConsole/Repositories/OrganizationUserRepository.cs index 418adf51325e..2dbf2644093c 100644 --- a/src/Infrastructure.Dapper/AdminConsole/Repositories/OrganizationUserRepository.cs +++ b/src/Infrastructure.Dapper/AdminConsole/Repositories/OrganizationUserRepository.cs @@ -7,7 +7,6 @@ using Bit.Core.AdminConsole.Utilities.DebuggingInstruments; using Bit.Core.Entities; using Bit.Core.Enums; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Models.Data; using Bit.Core.Models.Data.Organizations.OrganizationUsers; using Bit.Core.Repositories; @@ -596,7 +595,7 @@ public async Task> GetManyAcco } /// - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation( + public DatabaseTransactionAction UpdateForKeyRotation( Guid userId, IEnumerable resetPasswordKeys) { return async (connection, transaction) => diff --git a/src/Infrastructure.Dapper/Auth/Repositories/EmergencyAccessRepository.cs b/src/Infrastructure.Dapper/Auth/Repositories/EmergencyAccessRepository.cs index c76dfb525d17..39da66a1f2d3 100644 --- a/src/Infrastructure.Dapper/Auth/Repositories/EmergencyAccessRepository.cs +++ b/src/Infrastructure.Dapper/Auth/Repositories/EmergencyAccessRepository.cs @@ -1,7 +1,6 @@ using System.Data; using Bit.Core.Auth.Entities; using Bit.Core.Auth.Models.Data; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Repositories; using Bit.Core.Settings; using Bit.Infrastructure.Dapper.Auth.Helpers; @@ -124,11 +123,14 @@ public async Task> GetExpiredRecoveriesAsync } /// - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation( + public DatabaseTransactionAction UpdateForKeyRotation( Guid grantorId, IEnumerable emergencyAccessKeys) { - return async (SqlConnection connection, SqlTransaction transaction) => + return async (dbConnection, dbTransaction) => { + var connection = (SqlConnection)dbConnection; + var transaction = (SqlTransaction)dbTransaction; + // Create temp table var sqlCreateTemp = @" SELECT TOP 0 * diff --git a/src/Infrastructure.Dapper/Auth/Repositories/WebAuthnCredentialRepository.cs b/src/Infrastructure.Dapper/Auth/Repositories/WebAuthnCredentialRepository.cs index 7dfcd15d4921..3c6290238921 100644 --- a/src/Infrastructure.Dapper/Auth/Repositories/WebAuthnCredentialRepository.cs +++ b/src/Infrastructure.Dapper/Auth/Repositories/WebAuthnCredentialRepository.cs @@ -2,7 +2,7 @@ using Bit.Core.Auth.Entities; using Bit.Core.Auth.Models.Data; using Bit.Core.Auth.Repositories; -using Bit.Core.KeyManagement.UserKey; +using Bit.Core.Repositories; using Bit.Core.Settings; using Bit.Core.Utilities; using Bit.Infrastructure.Dapper.Repositories; @@ -61,9 +61,9 @@ public async Task UpdateAsync(WebAuthnCredential credential) return affectedRows > 0; } - public UpdateEncryptedDataForKeyRotation UpdateKeysForRotationAsync(Guid userId, IEnumerable credentials) + public DatabaseTransactionAction UpdateKeysForRotationAsync(Guid userId, IEnumerable credentials) { - return async (SqlConnection connection, SqlTransaction transaction) => + return async (connection, transaction) => { const string sql = @" UPDATE WC diff --git a/src/Infrastructure.Dapper/KeyManagement/Repositories/UserSignatureKeyPairRepository.cs b/src/Infrastructure.Dapper/KeyManagement/Repositories/UserSignatureKeyPairRepository.cs index 5dcc2943b8f7..9f501d9092aa 100644 --- a/src/Infrastructure.Dapper/KeyManagement/Repositories/UserSignatureKeyPairRepository.cs +++ b/src/Infrastructure.Dapper/KeyManagement/Repositories/UserSignatureKeyPairRepository.cs @@ -2,7 +2,7 @@ using Bit.Core.KeyManagement.Entities; using Bit.Core.KeyManagement.Models.Data; using Bit.Core.KeyManagement.Repositories; -using Bit.Core.KeyManagement.UserKey; +using Bit.Core.Repositories; using Bit.Core.Settings; using Bit.Core.Utilities; using Bit.Infrastructure.Dapper.Repositories; @@ -37,9 +37,9 @@ public UserSignatureKeyPairRepository(string connectionString, string readOnlyCo } } - public UpdateEncryptedDataForKeyRotation SetUserSignatureKeyPair(Guid userId, SignatureKeyPairData signingKeys) + public DatabaseTransactionAction SetUserSignatureKeyPair(Guid userId, SignatureKeyPairData signingKeys) { - return async (SqlConnection connection, SqlTransaction transaction) => + return async (connection, transaction) => { await connection.QueryAsync( "[dbo].[UserSignatureKeyPair_SetForRotation]", @@ -58,9 +58,9 @@ await connection.QueryAsync( }; } - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation(Guid grantorId, SignatureKeyPairData signingKeys) + public DatabaseTransactionAction UpdateForKeyRotation(Guid grantorId, SignatureKeyPairData signingKeys) { - return async (SqlConnection connection, SqlTransaction transaction) => + return async (connection, transaction) => { await connection.QueryAsync( "[dbo].[UserSignatureKeyPair_UpdateForRotation]", diff --git a/src/Infrastructure.Dapper/Repositories/DeviceRepository.cs b/src/Infrastructure.Dapper/Repositories/DeviceRepository.cs index f165d60955f5..8198a9374453 100644 --- a/src/Infrastructure.Dapper/Repositories/DeviceRepository.cs +++ b/src/Infrastructure.Dapper/Repositories/DeviceRepository.cs @@ -1,7 +1,6 @@ using System.Data; using Bit.Core.Auth.Models.Data; using Bit.Core.Entities; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Repositories; using Bit.Core.Settings; using Bit.Core.Utilities; @@ -130,9 +129,9 @@ await connection.ExecuteAsync( } } - public UpdateEncryptedDataForKeyRotation UpdateKeysForRotationAsync(Guid userId, IEnumerable devices) + public DatabaseTransactionAction UpdateKeysForRotationAsync(Guid userId, IEnumerable devices) { - return async (SqlConnection connection, SqlTransaction transaction) => + return async (connection, transaction) => { const string sql = @" UPDATE D diff --git a/src/Infrastructure.Dapper/Repositories/UserRepository.cs b/src/Infrastructure.Dapper/Repositories/UserRepository.cs index 0ea07d4bd40c..04cdb7e21c9b 100644 --- a/src/Infrastructure.Dapper/Repositories/UserRepository.cs +++ b/src/Infrastructure.Dapper/Repositories/UserRepository.cs @@ -6,7 +6,6 @@ using Bit.Core.Enums; using Bit.Core.KeyManagement.Kdf; using Bit.Core.KeyManagement.Models.Data; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Models.Data; using Bit.Core.Repositories; using Bit.Core.Settings; @@ -242,7 +241,7 @@ await connection.ExecuteAsync( /// public async Task UpdateUserKeyAndEncryptedDataAsync( User user, - IEnumerable updateDataActions) + IEnumerable updateDataActions) { await using var connection = new SqlConnection(ConnectionString); connection.Open(); @@ -292,7 +291,7 @@ public async Task UpdateUserKeyAndEncryptedDataAsync( public async Task UpdateUserKeyAndEncryptedDataV2Async( User user, - IEnumerable updateDataActions) + IEnumerable updateDataActions) { await using var connection = new SqlConnection(ConnectionString); connection.Open(); diff --git a/src/Infrastructure.Dapper/Tools/Repositories/SendRepository.cs b/src/Infrastructure.Dapper/Tools/Repositories/SendRepository.cs index a62d649c9dfc..e123e918f2b7 100644 --- a/src/Infrastructure.Dapper/Tools/Repositories/SendRepository.cs +++ b/src/Infrastructure.Dapper/Tools/Repositories/SendRepository.cs @@ -3,7 +3,7 @@ using System.Data; using System.Security.Cryptography; using Bit.Core; -using Bit.Core.KeyManagement.UserKey; +using Bit.Core.Repositories; using Bit.Core.Settings; using Bit.Core.Tools.Entities; using Bit.Core.Tools.Repositories; @@ -132,10 +132,13 @@ public override async Task ReplaceAsync(Send send) } /// - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation(Guid userId, IEnumerable sends) + public DatabaseTransactionAction UpdateForKeyRotation(Guid userId, IEnumerable sends) { - return async (connection, transaction) => + return async (dbConnection, dbTransaction) => { + var connection = (SqlConnection)dbConnection; + var transaction = (SqlTransaction)dbTransaction; + // Protect all sends before bulk update var sendsList = sends.ToList(); foreach (var send in sendsList) diff --git a/src/Infrastructure.Dapper/Vault/Repositories/CipherRepository.cs b/src/Infrastructure.Dapper/Vault/Repositories/CipherRepository.cs index fb43858902f4..aa58e365d244 100644 --- a/src/Infrastructure.Dapper/Vault/Repositories/CipherRepository.cs +++ b/src/Infrastructure.Dapper/Vault/Repositories/CipherRepository.cs @@ -4,7 +4,7 @@ using System.Data; using System.Text.Json; using Bit.Core.Entities; -using Bit.Core.KeyManagement.UserKey; +using Bit.Core.Repositories; using Bit.Core.Settings; using Bit.Core.Tools.Entities; using Bit.Core.Utilities; @@ -366,11 +366,14 @@ public async Task> GetUserSecurityTasksByCip } /// - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation( + public DatabaseTransactionAction UpdateForKeyRotation( Guid userId, IEnumerable ciphers) { - return async (SqlConnection connection, SqlTransaction transaction) => + return async (dbConnection, dbTransaction) => { + var connection = (SqlConnection)dbConnection; + var transaction = (SqlTransaction)dbTransaction; + // Create temp table var sqlCreateTemp = @" SELECT TOP 0 * diff --git a/src/Infrastructure.Dapper/Vault/Repositories/FolderRepository.cs b/src/Infrastructure.Dapper/Vault/Repositories/FolderRepository.cs index 63da064f8808..c6313de347a0 100644 --- a/src/Infrastructure.Dapper/Vault/Repositories/FolderRepository.cs +++ b/src/Infrastructure.Dapper/Vault/Repositories/FolderRepository.cs @@ -2,7 +2,7 @@ #nullable disable using System.Data; -using Bit.Core.KeyManagement.UserKey; +using Bit.Core.Repositories; using Bit.Core.Settings; using Bit.Core.Vault.Entities; using Bit.Core.Vault.Repositories; @@ -48,11 +48,14 @@ public async Task> GetManyByUserIdAsync(Guid userId) } /// - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation( + public DatabaseTransactionAction UpdateForKeyRotation( Guid userId, IEnumerable folders) { - return async (SqlConnection connection, SqlTransaction transaction) => + return async (dbConnection, dbTransaction) => { + var connection = (SqlConnection)dbConnection; + var transaction = (SqlTransaction)dbTransaction; + // Create temp table var sqlCreateTemp = @" SELECT TOP 0 * diff --git a/src/Infrastructure.EntityFramework/AdminConsole/Repositories/OrganizationUserRepository.cs b/src/Infrastructure.EntityFramework/AdminConsole/Repositories/OrganizationUserRepository.cs index 81f00c8f571f..bb5f8a690b50 100644 --- a/src/Infrastructure.EntityFramework/AdminConsole/Repositories/OrganizationUserRepository.cs +++ b/src/Infrastructure.EntityFramework/AdminConsole/Repositories/OrganizationUserRepository.cs @@ -8,7 +8,6 @@ using Bit.Core.AdminConsole.OrganizationFeatures.OrganizationUsers.InviteUsers.Models; using Bit.Core.Enums; using Bit.Core.Exceptions; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Models.Data; using Bit.Core.Models.Data.Organizations.OrganizationUsers; using Bit.Core.Repositories; @@ -857,14 +856,14 @@ on ou.OrganizationId equals o.Id } /// - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation( + public DatabaseTransactionAction UpdateForKeyRotation( Guid userId, IEnumerable resetPasswordKeys) { - return async (_, _) => + return async (connection, transaction) => { var newOrganizationUsers = resetPasswordKeys.ToList(); using var scope = ServiceScopeFactory.CreateScope(); - var dbContext = GetDatabaseContext(scope); + var dbContext = GetTransactionalDatabaseContext(scope, connection, transaction); // Get user organization users var userOrganizationUsers = await GetDbSet(dbContext) diff --git a/src/Infrastructure.EntityFramework/Auth/Repositories/EmergencyAccessRepository.cs b/src/Infrastructure.EntityFramework/Auth/Repositories/EmergencyAccessRepository.cs index ed811e9e0983..e13b2a4e6a8d 100644 --- a/src/Infrastructure.EntityFramework/Auth/Repositories/EmergencyAccessRepository.cs +++ b/src/Infrastructure.EntityFramework/Auth/Repositories/EmergencyAccessRepository.cs @@ -1,12 +1,10 @@ using AutoMapper; using Bit.Core.Auth.Enums; using Bit.Core.Auth.Models.Data; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Repositories; using Bit.Infrastructure.EntityFramework.Auth.Models; using Bit.Infrastructure.EntityFramework.Auth.Repositories.Queries; using Bit.Infrastructure.EntityFramework.Repositories; -using Microsoft.Data.SqlClient; using Microsoft.EntityFrameworkCore; using Microsoft.Extensions.DependencyInjection; @@ -153,14 +151,14 @@ public async Task> GetManyToNotifyAsync() } /// - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation( + public DatabaseTransactionAction UpdateForKeyRotation( Guid grantorId, IEnumerable emergencyAccessKeys) { - return async (SqlConnection connection, SqlTransaction transaction) => + return async (connection, transaction) => { var newKeys = emergencyAccessKeys.ToList(); using var scope = ServiceScopeFactory.CreateScope(); - var dbContext = GetDatabaseContext(scope); + var dbContext = GetTransactionalDatabaseContext(scope, connection, transaction); var userEmergencyAccess = await GetDbSet(dbContext) .Where(ea => ea.GrantorId == grantorId) .ToListAsync(); diff --git a/src/Infrastructure.EntityFramework/Auth/Repositories/WebAuthnCredentialRepository.cs b/src/Infrastructure.EntityFramework/Auth/Repositories/WebAuthnCredentialRepository.cs index ca32c44211bf..182f316a5c34 100644 --- a/src/Infrastructure.EntityFramework/Auth/Repositories/WebAuthnCredentialRepository.cs +++ b/src/Infrastructure.EntityFramework/Auth/Repositories/WebAuthnCredentialRepository.cs @@ -1,7 +1,7 @@ using AutoMapper; using Bit.Core.Auth.Models.Data; using Bit.Core.Auth.Repositories; -using Bit.Core.KeyManagement.UserKey; +using Bit.Core.Repositories; using Bit.Infrastructure.EntityFramework.Auth.Models; using Bit.Infrastructure.EntityFramework.Repositories; using Microsoft.EntityFrameworkCore; @@ -61,13 +61,13 @@ public async Task UpdateAsync(Core.Auth.Entities.WebAuthnCredential creden } } - public UpdateEncryptedDataForKeyRotation UpdateKeysForRotationAsync(Guid userId, IEnumerable credentials) + public DatabaseTransactionAction UpdateKeysForRotationAsync(Guid userId, IEnumerable credentials) { - return async (_, _) => + return async (connection, transaction) => { var newCreds = credentials.ToList(); using var scope = ServiceScopeFactory.CreateScope(); - var dbContext = GetDatabaseContext(scope); + var dbContext = GetTransactionalDatabaseContext(scope, connection, transaction); var newCredIds = newCreds.Select(nwc => nwc.Id).ToList(); var validUserWebauthnCredentials = await GetDbSet(dbContext) diff --git a/src/Infrastructure.EntityFramework/KeyManagement/Repositories/UserSignatureKeyPairRepository.cs b/src/Infrastructure.EntityFramework/KeyManagement/Repositories/UserSignatureKeyPairRepository.cs index 04f055501d09..11c645c0c9a0 100644 --- a/src/Infrastructure.EntityFramework/KeyManagement/Repositories/UserSignatureKeyPairRepository.cs +++ b/src/Infrastructure.EntityFramework/KeyManagement/Repositories/UserSignatureKeyPairRepository.cs @@ -2,7 +2,7 @@ using AutoMapper; using Bit.Core.KeyManagement.Models.Data; using Bit.Core.KeyManagement.Repositories; -using Bit.Core.KeyManagement.UserKey; +using Bit.Core.Repositories; using Bit.Core.Utilities; using Bit.Infrastructure.EntityFramework.Repositories; using Microsoft.EntityFrameworkCore; @@ -25,12 +25,12 @@ public class UserSignatureKeyPairRepository(IServiceScopeFactory serviceScopeFac return signingKeys.ToSignatureKeyPairData(); } - public UpdateEncryptedDataForKeyRotation SetUserSignatureKeyPair(Guid userId, SignatureKeyPairData signingKeys) + public DatabaseTransactionAction SetUserSignatureKeyPair(Guid userId, SignatureKeyPairData signingKeys) { - return async (_, _) => + return async (connection, transaction) => { await using var scope = ServiceScopeFactory.CreateAsyncScope(); - var dbContext = GetDatabaseContext(scope); + var dbContext = GetTransactionalDatabaseContext(scope, connection, transaction); var entity = new Models.UserSignatureKeyPair { Id = CoreHelpers.GenerateComb(), @@ -46,12 +46,12 @@ public UpdateEncryptedDataForKeyRotation SetUserSignatureKeyPair(Guid userId, Si }; } - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation(Guid grantorId, SignatureKeyPairData signingKeys) + public DatabaseTransactionAction UpdateForKeyRotation(Guid grantorId, SignatureKeyPairData signingKeys) { - return async (_, _) => + return async (connection, transaction) => { await using var scope = ServiceScopeFactory.CreateAsyncScope(); - var dbContext = GetDatabaseContext(scope); + var dbContext = GetTransactionalDatabaseContext(scope, connection, transaction); var entity = await dbContext.UserSignatureKeyPairs.FirstOrDefaultAsync(x => x.UserId == grantorId); if (entity != null) { diff --git a/src/Infrastructure.EntityFramework/Repositories/BaseEntityFrameworkRepository.cs b/src/Infrastructure.EntityFramework/Repositories/BaseEntityFrameworkRepository.cs index 6cf7cbb46efc..c98c8ff085f1 100644 --- a/src/Infrastructure.EntityFramework/Repositories/BaseEntityFrameworkRepository.cs +++ b/src/Infrastructure.EntityFramework/Repositories/BaseEntityFrameworkRepository.cs @@ -1,4 +1,5 @@ -using System.Text.Json; +using System.Data.Common; +using System.Text.Json; using AutoMapper; using Bit.Infrastructure.EntityFramework.AdminConsole.Models; using Bit.Infrastructure.EntityFramework.Repositories.Queries; @@ -31,6 +32,15 @@ public DatabaseContext GetDatabaseContext(IServiceScope serviceScope) return serviceScope.ServiceProvider.GetRequiredService(); } + protected DatabaseContext GetTransactionalDatabaseContext( + IServiceScope scope, DbConnection connection, DbTransaction transaction) + { + var dbContext = GetDatabaseContext(scope); + dbContext.Database.SetDbConnection(connection); + dbContext.Database.UseTransaction(transaction); + return dbContext; + } + public void ClearChangeTracking() { using (var scope = ServiceScopeFactory.CreateScope()) diff --git a/src/Infrastructure.EntityFramework/Repositories/DeviceRepository.cs b/src/Infrastructure.EntityFramework/Repositories/DeviceRepository.cs index cbc7709d204c..4e2c70fe5c53 100644 --- a/src/Infrastructure.EntityFramework/Repositories/DeviceRepository.cs +++ b/src/Infrastructure.EntityFramework/Repositories/DeviceRepository.cs @@ -1,6 +1,5 @@ using AutoMapper; using Bit.Core.Auth.Models.Data; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Repositories; using Bit.Core.Settings; using Bit.Infrastructure.EntityFramework.Auth.Repositories.Queries; @@ -200,13 +199,13 @@ await dbContext.Devices : d.ClientVersion)); } - public UpdateEncryptedDataForKeyRotation UpdateKeysForRotationAsync(Guid userId, IEnumerable devices) + public DatabaseTransactionAction UpdateKeysForRotationAsync(Guid userId, IEnumerable devices) { - return async (_, _) => + return async (connection, transaction) => { var deviceUpdates = devices.ToList(); using var scope = ServiceScopeFactory.CreateScope(); - var dbContext = GetDatabaseContext(scope); + var dbContext = GetTransactionalDatabaseContext(scope, connection, transaction); var userDevices = await GetDbSet(dbContext) .Where(device => device.UserId == userId) .ToListAsync(); diff --git a/src/Infrastructure.EntityFramework/Repositories/UserRepository.cs b/src/Infrastructure.EntityFramework/Repositories/UserRepository.cs index d7b10aa62b2f..790148757dea 100644 --- a/src/Infrastructure.EntityFramework/Repositories/UserRepository.cs +++ b/src/Infrastructure.EntityFramework/Repositories/UserRepository.cs @@ -4,7 +4,6 @@ using Bit.Core.Enums; using Bit.Core.KeyManagement.Kdf; using Bit.Core.KeyManagement.Models.Data; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Models.Data; using Bit.Core.Repositories; using Bit.Infrastructure.EntityFramework.Models; @@ -195,12 +194,15 @@ public async Task UpdateRenewalReminderDateAsync(Guid id, DateTime renewalRemind /// public async Task UpdateUserKeyAndEncryptedDataAsync(Core.Entities.User user, - IEnumerable updateDataActions) + IEnumerable updateDataActions) { using var scope = ServiceScopeFactory.CreateScope(); var dbContext = GetDatabaseContext(scope); - await using var transaction = await dbContext.Database.BeginTransactionAsync(); + var connection = dbContext.Database.GetDbConnection(); + await connection.OpenAsync(); + await using var transaction = await connection.BeginTransactionAsync(); + await dbContext.Database.UseTransactionAsync(transaction); try { @@ -225,8 +227,7 @@ public async Task UpdateUserKeyAndEncryptedDataAsync(Core.Entities.User user, // Update re-encrypted data foreach (var action in updateDataActions) { - // connection and transaction aren't used in EF - await action(); + await action(connection, transaction); } await transaction.CommitAsync(); @@ -241,12 +242,15 @@ public async Task UpdateUserKeyAndEncryptedDataAsync(Core.Entities.User user, public async Task UpdateUserKeyAndEncryptedDataV2Async(Core.Entities.User user, - IEnumerable updateDataActions) + IEnumerable updateDataActions) { using var scope = ServiceScopeFactory.CreateScope(); var dbContext = GetDatabaseContext(scope); - await using var transaction = await dbContext.Database.BeginTransactionAsync(); + var connection = dbContext.Database.GetDbConnection(); + await connection.OpenAsync(); + await using var transaction = await connection.BeginTransactionAsync(); + await dbContext.Database.UseTransactionAsync(transaction); // Update user var userEntity = await dbContext.Users.FindAsync(user.Id); @@ -287,8 +291,7 @@ public async Task UpdateUserKeyAndEncryptedDataV2Async(Core.Entities.User user, // Update re-encrypted data foreach (var action in updateDataActions) { - // connection and transaction aren't used in EF - await action(); + await action(connection, transaction); } await transaction.CommitAsync(); diff --git a/src/Infrastructure.EntityFramework/Tools/Repositories/SendRepository.cs b/src/Infrastructure.EntityFramework/Tools/Repositories/SendRepository.cs index 0ab18e3d255d..13882c7f8937 100644 --- a/src/Infrastructure.EntityFramework/Tools/Repositories/SendRepository.cs +++ b/src/Infrastructure.EntityFramework/Tools/Repositories/SendRepository.cs @@ -3,7 +3,7 @@ using System.Security.Cryptography; using AutoMapper; using Bit.Core; -using Bit.Core.KeyManagement.UserKey; +using Bit.Core.Repositories; using Bit.Core.Tools.Enums; using Bit.Core.Tools.Repositories; using Bit.Infrastructure.EntityFramework.Models; @@ -150,17 +150,17 @@ public override async Task ReplaceAsync(Core.Tools.Entities.Send send) } /// - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation(Guid userId, + public DatabaseTransactionAction UpdateForKeyRotation(Guid userId, IEnumerable sends) { - return async (_, _) => + return async (connection, transaction) => { // No Emails protect/unprotect needed here: this only mutates Key on tracked entities, and EF // writes only the changed column, so the already-protected Emails at rest is untouched. (The // Dapper implementation protects because it bulk-copies whole rows.) var newSends = sends.ToDictionary(s => s.Id); using var scope = ServiceScopeFactory.CreateScope(); - var dbContext = GetDatabaseContext(scope); + var dbContext = GetTransactionalDatabaseContext(scope, connection, transaction); var userSends = await GetDbSet(dbContext) .Where(s => s.UserId == userId) .ToListAsync(); diff --git a/src/Infrastructure.EntityFramework/Vault/Repositories/CipherRepository.cs b/src/Infrastructure.EntityFramework/Vault/Repositories/CipherRepository.cs index e74c05b2293e..dcc119664e7d 100644 --- a/src/Infrastructure.EntityFramework/Vault/Repositories/CipherRepository.cs +++ b/src/Infrastructure.EntityFramework/Vault/Repositories/CipherRepository.cs @@ -5,7 +5,7 @@ using System.Text.Json.Nodes; using AutoMapper; using Bit.Core.Enums; -using Bit.Core.KeyManagement.UserKey; +using Bit.Core.Repositories; using Bit.Core.Utilities; using Bit.Core.Vault.Enums; using Bit.Core.Vault.Models.Data; @@ -19,7 +19,6 @@ using Bit.Infrastructure.EntityFramework.Vault.Models; using Bit.Infrastructure.EntityFramework.Vault.Repositories.Queries; using LinqToDB.EntityFrameworkCore; -using Microsoft.Data.SqlClient; using Microsoft.EntityFrameworkCore; using Microsoft.Extensions.DependencyInjection; using NS = Newtonsoft.Json; @@ -1035,14 +1034,14 @@ public async Task UpdatePartialAsync(Guid id, Guid userId, Guid? folderId, bool } /// - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation( + public DatabaseTransactionAction UpdateForKeyRotation( Guid userId, IEnumerable ciphers) { - return async (SqlConnection _, SqlTransaction _) => + return async (connection, transaction) => { var newCiphers = ciphers.ToList(); using var scope = ServiceScopeFactory.CreateScope(); - var dbContext = GetDatabaseContext(scope); + var dbContext = GetTransactionalDatabaseContext(scope, connection, transaction); var userCiphers = await GetDbSet(dbContext) .Where(c => c.UserId == userId) .ToListAsync(); diff --git a/src/Infrastructure.EntityFramework/Vault/Repositories/FolderRepository.cs b/src/Infrastructure.EntityFramework/Vault/Repositories/FolderRepository.cs index 83fa442eb4f9..406a4e81c4d1 100644 --- a/src/Infrastructure.EntityFramework/Vault/Repositories/FolderRepository.cs +++ b/src/Infrastructure.EntityFramework/Vault/Repositories/FolderRepository.cs @@ -2,11 +2,10 @@ #nullable disable using AutoMapper; -using Bit.Core.KeyManagement.UserKey; +using Bit.Core.Repositories; using Bit.Core.Vault.Repositories; using Bit.Infrastructure.EntityFramework.Repositories; using Bit.Infrastructure.EntityFramework.Vault.Models; -using Microsoft.Data.SqlClient; using Microsoft.EntityFrameworkCore; using Microsoft.Extensions.DependencyInjection; @@ -43,14 +42,14 @@ public FolderRepository(IServiceScopeFactory serviceScopeFactory, IMapper mapper } /// - public UpdateEncryptedDataForKeyRotation UpdateForKeyRotation( + public DatabaseTransactionAction UpdateForKeyRotation( Guid userId, IEnumerable folders) { - return async (SqlConnection _, SqlTransaction _) => + return async (connection, transaction) => { var newFolders = folders.ToList(); using var scope = ServiceScopeFactory.CreateScope(); - var dbContext = GetDatabaseContext(scope); + var dbContext = GetTransactionalDatabaseContext(scope, connection, transaction); var userFolders = await GetDbSet(dbContext) .Where(f => f.UserId == userId) .ToListAsync(); diff --git a/test/Core.Test/KeyManagement/UserKey/RotateUserAccountKeysCommandTests.cs b/test/Core.Test/KeyManagement/UserKey/RotateUserAccountKeysCommandTests.cs index 5d04dda75998..c225336303ca 100644 --- a/test/Core.Test/KeyManagement/UserKey/RotateUserAccountKeysCommandTests.cs +++ b/test/Core.Test/KeyManagement/UserKey/RotateUserAccountKeysCommandTests.cs @@ -6,7 +6,6 @@ using Bit.Core.KeyManagement.Enums; using Bit.Core.KeyManagement.Models.Data; using Bit.Core.KeyManagement.Repositories; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.KeyManagement.UserKey.Implementations; using Bit.Core.KeyManagement.UserKey.Models.Data; using Bit.Core.Platform.Push; @@ -164,7 +163,7 @@ public async Task UpdateAccountKeysAsync_PublicKeyChange_Rejects(SutProvider(); + var saveEncryptedDataActions = new List(); await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); } @@ -176,7 +175,7 @@ public async Task UpdateAccountKeysAsync_V2User_PrivateKeyNotXChaCha20_Rejects(S SetV2ModelUser(model); model.AccountKeys.PublicKeyEncryptionKeyPairData.WrappedPrivateKey = _mockEncryptedType2String; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); } @@ -188,7 +187,7 @@ public async Task UpdateAccountKeysAsync_V1User_PrivateKeyNotAesCbcHmac_Rejects( SetV1ModelUser(model); model.AccountKeys.PublicKeyEncryptionKeyPairData.WrappedPrivateKey = _mockEncryptedType7String; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); var ex = await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); Assert.Equal("The provided account private key was not wrapped with AES-256-CBC-HMAC", ex.Message); } @@ -200,7 +199,7 @@ public async Task UpdateAccountKeysAsync_V1_Success(SutProvider(); + var saveEncryptedDataActions = new List(); await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions); Assert.Empty(saveEncryptedDataActions); } @@ -212,7 +211,7 @@ public async Task UpdateAccountKeysAsync_V2_Success(SutProvider(); + var saveEncryptedDataActions = new List(); await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions); Assert.NotEmpty(saveEncryptedDataActions); Assert.Equal(user.SecurityState, model.AccountKeys.SecurityStateData!.SecurityState); @@ -226,7 +225,7 @@ public async Task UpdateAccountKeysAsync_V2User_VerifyingKeyMismatch_Rejects(Sut SetV2ModelUser(model); model.AccountKeys.SignatureKeyPairData.VerifyingKey = "different-verifying-key"; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); var ex = await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); Assert.Equal("The provided verifying key does not match the user's current verifying key.", ex.Message); } @@ -239,7 +238,7 @@ public async Task UpdateAccountKeysAsync_V2User_SignedPublicKeyNullOrEmpty_Rejec SetV2ModelUser(model); model.AccountKeys.PublicKeyEncryptionKeyPairData.SignedPublicKey = null; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); var ex = await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); Assert.Equal("No signed public key provided, but the user already has a signature key pair.", ex.Message); } @@ -252,7 +251,7 @@ public async Task UpdateAccountKeysAsync_V2User_WrappedSigningKeyNotXChaCha20_Re SetV2ModelUser(model); model.AccountKeys.SignatureKeyPairData.WrappedSigningKey = _mockEncryptedType2String; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); var ex = await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); Assert.Equal("The provided signing key data is not wrapped with XChaCha20-Poly1305.", ex.Message); } @@ -265,7 +264,7 @@ public async Task UpdateAccountKeys_UpgradeToV2_InvalidVerifyingKey_Rejects(SutP SetV2ModelUser(model); model.AccountKeys.SignatureKeyPairData.VerifyingKey = ""; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); var ex = await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); Assert.Equal("The provided signature key pair data does not contain a valid verifying key.", ex.Message); } @@ -278,7 +277,7 @@ public async Task UpdateAccountKeysAsync_UpgradeToV2_IncorrectlyWrappedPrivateKe SetV2ModelUser(model); model.AccountKeys.PublicKeyEncryptionKeyPairData.WrappedPrivateKey = _mockEncryptedType2String; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); var ex = await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); Assert.Equal("The provided private key encryption key is not wrapped with XChaCha20-Poly1305.", ex.Message); } @@ -291,7 +290,7 @@ public async Task UpdateAccountKeysAsync_UpgradeToV2_NoSignedPublicKey_Rejects(S SetV2ModelUser(model); model.AccountKeys.PublicKeyEncryptionKeyPairData.SignedPublicKey = null; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); var ex = await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); Assert.Equal("No signed public key provided, but the user already has a signature key pair.", ex.Message); } @@ -304,7 +303,7 @@ public async Task UpdateAccountKeysAsync_UpgradeToV2_NoSecurityState_Rejects(Sut SetV2ModelUser(model); model.AccountKeys.SecurityStateData = null; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); var ex = await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); Assert.Equal("No signed security state provider for V2 user", ex.Message); } @@ -317,7 +316,7 @@ public async Task UpdateAccountKeysAsync_RotateV2_NoSignatureKeyPair_Rejects(Sut SetV2ModelUser(model); model.AccountKeys.SignatureKeyPairData = null; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); var ex = await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); Assert.Equal("Signature key pair data is required for V2 encryption.", ex.Message); } @@ -330,7 +329,7 @@ public async Task UpdateAccountKeysAsync_GetEncryptionType_EmptyString_Rejects(S SetV1ModelUser(model); model.AccountKeys.PublicKeyEncryptionKeyPairData.WrappedPrivateKey = ""; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); var ex = await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); Assert.Equal("Invalid encryption type string.", ex.Message); } @@ -343,7 +342,7 @@ public async Task UpdateAccountKeysAsync_GetEncryptionType_InvalidString_Rejects SetV1ModelUser(model); model.AccountKeys.PublicKeyEncryptionKeyPairData.WrappedPrivateKey = "9.xxx"; - var saveEncryptedDataActions = new List(); + var saveEncryptedDataActions = new List(); var ex = await Assert.ThrowsAsync(async () => await sutProvider.Sut.UpdateAccountKeysAsync(model, user, saveEncryptedDataActions)); Assert.Equal("Invalid encryption type string.", ex.Message); } @@ -365,12 +364,12 @@ public async Task UpdateUserData_RevisionDateChanged_Success(SutProvider(); + var saveEncryptedDataActions = new List(); sutProvider.Sut.UpdateUserData(model, user, saveEncryptedDataActions); foreach (var dataAction in saveEncryptedDataActions) { - await dataAction.Invoke(); + await dataAction.Invoke(null!, null!); } var updatedCiphers = sutProvider.GetDependency() @@ -653,7 +652,7 @@ public async Task PasswordChangeAndRotateUserAccountKeysAsync_MasterPasswordServ Assert.False(result.Succeeded); Assert.Contains(result.Errors, e => e.Code == "SomeError"); await sutProvider.GetDependency().DidNotReceive() - .UpdateUserKeyAndEncryptedDataV2Async(Arg.Any(), Arg.Any>()); + .UpdateUserKeyAndEncryptedDataV2Async(Arg.Any(), Arg.Any>()); await sutProvider.GetDependency().DidNotReceive() .PushLogOutAsync(Arg.Any(), Arg.Any(), Arg.Any()); } @@ -748,7 +747,7 @@ public async Task MasterPasswordRotateUserAccountKeysAsync_V2User_Success( Assert.Equal(model.MasterPasswordUnlockData.MasterKeyWrappedUserKey, user.Key); await sutProvider.GetDependency().Received(1) - .UpdateUserKeyAndEncryptedDataV2Async(user, Arg.Any>()); + .UpdateUserKeyAndEncryptedDataV2Async(user, Arg.Any>()); Assert.NotEqual(originalSecurityStamp, user.SecurityStamp); await sutProvider.GetDependency().Received(1) .PushLogOutAsync(user.Id); @@ -846,7 +845,7 @@ public async Task TdeRotateUserAccountKeysAsync_V1User_Success( Assert.Null(user.Key); Assert.NotEqual(originalSecurityStamp, user.SecurityStamp); await sutProvider.GetDependency().Received(1) - .UpdateUserKeyAndEncryptedDataV2Async(user, Arg.Any>()); + .UpdateUserKeyAndEncryptedDataV2Async(user, Arg.Any>()); await sutProvider.GetDependency().Received(1) .PushLogOutAsync(user.Id); } @@ -867,7 +866,7 @@ public async Task TdeRotateUserAccountKeysAsync_V2User_Success( Assert.Null(user.Key); Assert.NotEqual(originalSecurityStamp, user.SecurityStamp); await sutProvider.GetDependency().Received(1) - .UpdateUserKeyAndEncryptedDataV2Async(user, Arg.Any>()); + .UpdateUserKeyAndEncryptedDataV2Async(user, Arg.Any>()); await sutProvider.GetDependency().Received(1) .PushLogOutAsync(user.Id); } @@ -988,7 +987,7 @@ public async Task KeyConnectorRotateUserAccountKeysAsync_V1User_Success( Assert.Equal(model.KeyConnectorKeyWrappedUserKey, user.Key); Assert.NotEqual(originalSecurityStamp, user.SecurityStamp); await sutProvider.GetDependency().Received(1) - .UpdateUserKeyAndEncryptedDataV2Async(user, Arg.Any>()); + .UpdateUserKeyAndEncryptedDataV2Async(user, Arg.Any>()); await sutProvider.GetDependency().Received(1) .PushLogOutAsync(user.Id); } @@ -1009,7 +1008,7 @@ public async Task KeyConnectorRotateUserAccountKeysAsync_V2User_Success( Assert.Equal(model.KeyConnectorKeyWrappedUserKey, user.Key); Assert.NotEqual(originalSecurityStamp, user.SecurityStamp); await sutProvider.GetDependency().Received(1) - .UpdateUserKeyAndEncryptedDataV2Async(user, Arg.Any>()); + .UpdateUserKeyAndEncryptedDataV2Async(user, Arg.Any>()); await sutProvider.GetDependency().Received(1) .PushLogOutAsync(user.Id); } diff --git a/test/Core.Test/Platform/Push/PushServiceCollectionExtensionsTests.cs b/test/Core.Test/Platform/Push/PushServiceCollectionExtensionsTests.cs index dabddc605851..086e82fc5ebc 100644 --- a/test/Core.Test/Platform/Push/PushServiceCollectionExtensionsTests.cs +++ b/test/Core.Test/Platform/Push/PushServiceCollectionExtensionsTests.cs @@ -1,6 +1,5 @@ using Bit.Core.Auth.Models.Data; using Bit.Core.Entities; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Platform.Push; using Bit.Core.Platform.Push.Internal; using Bit.Core.Repositories; @@ -192,7 +191,7 @@ private class StubDeviceRepository : IDeviceRepository public Task> GetManyByUserIdAsync(Guid userId) => throw new NotImplementedException(); public Task> GetManyByUserIdWithDeviceAuth(Guid userId) => throw new NotImplementedException(); public Task ReplaceAsync(Device obj) => throw new NotImplementedException(); - public UpdateEncryptedDataForKeyRotation UpdateKeysForRotationAsync(Guid userId, IEnumerable devices) => throw new NotImplementedException(); + public DatabaseTransactionAction UpdateKeysForRotationAsync(Guid userId, IEnumerable devices) => throw new NotImplementedException(); public Task UpsertAsync(Device obj) => throw new NotImplementedException(); public Task UpdateLastActivityByIdAsync(Guid deviceId, string? clientVersion) => throw new NotImplementedException(); public Task UpdateLastActivityByIdentifierAndUserIdAsync(string identifier, Guid userId, string? clientVersion) => throw new NotImplementedException(); diff --git a/test/Infrastructure.IntegrationTest/DatabaseTransactionActionTestHelper.cs b/test/Infrastructure.IntegrationTest/DatabaseTransactionActionTestHelper.cs new file mode 100644 index 000000000000..7a308a100437 --- /dev/null +++ b/test/Infrastructure.IntegrationTest/DatabaseTransactionActionTestHelper.cs @@ -0,0 +1,48 @@ +using Bit.Core.Enums; +using Bit.Core.Repositories; +using Bit.Infrastructure.EntityFramework.Repositories; +using Microsoft.Data.SqlClient; +using Microsoft.EntityFrameworkCore; +using Microsoft.Extensions.DependencyInjection; + +namespace Bit.Infrastructure.IntegrationTest; + +/// +/// Executes delegates in integration tests. +/// Opens a connection and transaction appropriate for the database provider, executes the actions, and commits. +/// +public static class DatabaseTransactionActionTestHelper +{ + public static Task ExecuteAsync(Database database, DatabaseTransactionAction action, + IServiceProvider serviceProvider) + => ExecuteAsync(database, [action], serviceProvider); + + public static async Task ExecuteAsync(Database database, IEnumerable actions, + IServiceProvider serviceProvider) + { + var isDapper = database.Type == SupportedDatabaseProviders.SqlServer && !database.UseEf; + var connection = isDapper + ? new SqlConnection(database.ConnectionString) + : serviceProvider.GetRequiredService().Database.GetDbConnection(); + + try + { + await connection.OpenAsync(); + await using var transaction = await connection.BeginTransactionAsync(); + + foreach (var action in actions) + { + await action(connection, transaction); + } + + await transaction.CommitAsync(); + } + finally + { + if (isDapper) + { + await connection.DisposeAsync(); + } + } + } +} diff --git a/test/Infrastructure.IntegrationTest/Repositories/UserRepositoryTests.cs b/test/Infrastructure.IntegrationTest/Repositories/UserRepositoryTests.cs index 444c5c525f80..21b534eae820 100644 --- a/test/Infrastructure.IntegrationTest/Repositories/UserRepositoryTests.cs +++ b/test/Infrastructure.IntegrationTest/Repositories/UserRepositoryTests.cs @@ -5,9 +5,10 @@ using Bit.Core.KeyManagement.Enums; using Bit.Core.KeyManagement.Kdf; using Bit.Core.KeyManagement.Models.Data; -using Bit.Core.KeyManagement.UserKey; using Bit.Core.Models.Data; using Bit.Core.Repositories; +using Bit.Core.Vault.Entities; +using Bit.Core.Vault.Repositories; using Bit.Infrastructure.IntegrationTest.AdminConsole; using Microsoft.Data.SqlClient; using Xunit; @@ -771,7 +772,7 @@ public async Task UpdateUserKeyAndEncryptedDataV2Async_InvokesUpdateDataActions( user.RevisionDate = DateTime.UtcNow; var actionWasInvoked = false; - UpdateEncryptedDataForKeyRotation action = (_, _) => + DatabaseTransactionAction action = (_, _) => { actionWasInvoked = true; return Task.CompletedTask; @@ -784,6 +785,26 @@ public async Task UpdateUserKeyAndEncryptedDataV2Async_InvokesUpdateDataActions( Assert.True(actionWasInvoked); } + [Theory, DatabaseData] + public async Task UpdateUserKeyAndEncryptedDataV2Async_WithRepositoryUpdateDataAction_PersistsRotatedData( + IUserRepository userRepository, IFolderRepository folderRepository) + { + // Arrange + var user = await userRepository.CreateTestUserAsync(); + var folder = await folderRepository.CreateAsync(new Folder { UserId = user.Id, Name = "2.original-name" }); + folder.Name = "2.rotated-name"; + user.RevisionDate = DateTime.UtcNow; + + // Act + await userRepository.UpdateUserKeyAndEncryptedDataV2Async( + user, [folderRepository.UpdateForKeyRotation(user.Id, [folder])]); + + // Assert + var rotatedFolder = await folderRepository.GetByIdAsync(folder.Id); + Assert.NotNull(rotatedFolder); + Assert.Equal("2.rotated-name", rotatedFolder.Name); + } + [Theory, DatabaseData] public async Task GetKdfInformationByEmailAsync_WithPbkdf2User_ReturnsKdfInformation( IUserRepository userRepository) From 8e9bdbcc923fd3416776a1d0a4c6982163b8d193 Mon Sep 17 00:00:00 2001 From: Maciej Zieniuk Date: Fri, 7 Aug 2026 14:36:14 +0100 Subject: [PATCH 2/3] remove redundant test --- .../Repositories/UserRepositoryTests.cs | 22 ------------------- 1 file changed, 22 deletions(-) diff --git a/test/Infrastructure.IntegrationTest/Repositories/UserRepositoryTests.cs b/test/Infrastructure.IntegrationTest/Repositories/UserRepositoryTests.cs index 21b534eae820..f7720a793ec9 100644 --- a/test/Infrastructure.IntegrationTest/Repositories/UserRepositoryTests.cs +++ b/test/Infrastructure.IntegrationTest/Repositories/UserRepositoryTests.cs @@ -7,8 +7,6 @@ using Bit.Core.KeyManagement.Models.Data; using Bit.Core.Models.Data; using Bit.Core.Repositories; -using Bit.Core.Vault.Entities; -using Bit.Core.Vault.Repositories; using Bit.Infrastructure.IntegrationTest.AdminConsole; using Microsoft.Data.SqlClient; using Xunit; @@ -785,26 +783,6 @@ public async Task UpdateUserKeyAndEncryptedDataV2Async_InvokesUpdateDataActions( Assert.True(actionWasInvoked); } - [Theory, DatabaseData] - public async Task UpdateUserKeyAndEncryptedDataV2Async_WithRepositoryUpdateDataAction_PersistsRotatedData( - IUserRepository userRepository, IFolderRepository folderRepository) - { - // Arrange - var user = await userRepository.CreateTestUserAsync(); - var folder = await folderRepository.CreateAsync(new Folder { UserId = user.Id, Name = "2.original-name" }); - folder.Name = "2.rotated-name"; - user.RevisionDate = DateTime.UtcNow; - - // Act - await userRepository.UpdateUserKeyAndEncryptedDataV2Async( - user, [folderRepository.UpdateForKeyRotation(user.Id, [folder])]); - - // Assert - var rotatedFolder = await folderRepository.GetByIdAsync(folder.Id); - Assert.NotNull(rotatedFolder); - Assert.Equal("2.rotated-name", rotatedFolder.Name); - } - [Theory, DatabaseData] public async Task GetKdfInformationByEmailAsync_WithPbkdf2User_ReturnsKdfInformation( IUserRepository userRepository) From 09674868683ffc30259b847d41fe29193f94e4cb Mon Sep 17 00:00:00 2001 From: Maciej Zieniuk Date: Fri, 7 Aug 2026 12:46:24 +0100 Subject: [PATCH 3/3] [PM-39455] Add OrganizationUser V2 upgrade token rotation tests Cover OrganizationUserRepository.UpdateForKeyRotation across all four providers: it persists the reset password key and the V2 upgrade token together, clears a token left over from an earlier upgrade, and leaves another member's membership untouched. These need both the V2UpgradeToken column and the shared rotation transaction, so they are kept out of the branches that add each. --- .../OrganizationUserRepositoryTests.cs | 88 +++++++++++++++++++ 1 file changed, 88 insertions(+) diff --git a/test/Infrastructure.IntegrationTest/AdminConsole/Repositories/OrganizationUserRepository/OrganizationUserRepositoryTests.cs b/test/Infrastructure.IntegrationTest/AdminConsole/Repositories/OrganizationUserRepository/OrganizationUserRepositoryTests.cs index c1036369a18f..5b65560091ea 100644 --- a/test/Infrastructure.IntegrationTest/AdminConsole/Repositories/OrganizationUserRepository/OrganizationUserRepositoryTests.cs +++ b/test/Infrastructure.IntegrationTest/AdminConsole/Repositories/OrganizationUserRepository/OrganizationUserRepositoryTests.cs @@ -17,6 +17,10 @@ namespace Bit.Infrastructure.IntegrationTest.AdminConsole.Repositories.Organizat public class OrganizationUserRepositoryTests { + private const string _resetPasswordKey = "4.reset-password-key"; + private const string _v2UpgradeToken = """{"WrappedUserKey1":"7.key-one","WrappedUserKey2":"2.key-two"}"""; + private const string _otherV2UpgradeToken = """{"WrappedUserKey1":"7.key-three","WrappedUserKey2":"2.key-four"}"""; + [Theory, DatabaseData] public async Task GetOccupiedSmSeatCountByOrganizationIdAsync_ExcludesRevokedAndStaged( IUserRepository userRepository, @@ -1793,4 +1797,88 @@ public async Task UpdateGroupsAsync_BumpsGroupRevisionDate( Assert.Equal(expectedRevisionDate, actualGroup1.RevisionDate, TimeSpan.FromMilliseconds(10)); Assert.Equal(expectedRevisionDate, actualGroup2.RevisionDate, TimeSpan.FromMilliseconds(10)); } + + [Theory, DatabaseData] + public async Task UpdateForKeyRotation_WithV2UpgradeToken_PersistsResetPasswordKeyAndToken( + IUserRepository userRepository, + IOrganizationRepository organizationRepository, + IOrganizationUserRepository organizationUserRepository) + { + // Arrange + var user = await userRepository.CreateTestUserAsync(); + var organization = await organizationRepository.CreateTestOrganizationAsync(); + var organizationUser = await organizationUserRepository.CreateTestOrganizationUserAsync(organization, user); + + organizationUser.ResetPasswordKey = _resetPasswordKey; + organizationUser.V2UpgradeToken = _v2UpgradeToken; + + // Act + await userRepository.UpdateUserKeyAndEncryptedDataV2Async(user, + [organizationUserRepository.UpdateForKeyRotation(user.Id, [organizationUser])]); + + // Assert + var updated = await organizationUserRepository.GetByIdAsync(organizationUser.Id); + Assert.NotNull(updated); + Assert.Equal(_resetPasswordKey, updated.ResetPasswordKey); + Assert.Equal(_v2UpgradeToken, updated.V2UpgradeToken); + } + + [Theory, DatabaseData] + public async Task UpdateForKeyRotation_WithoutV2UpgradeToken_ClearsStaleToken( + IUserRepository userRepository, + IOrganizationRepository organizationRepository, + IOrganizationUserRepository organizationUserRepository) + { + // Arrange + var user = await userRepository.CreateTestUserAsync(); + var organization = await organizationRepository.CreateTestOrganizationAsync(); + var organizationUser = await organizationUserRepository.CreateTestOrganizationUserAsync(organization, user); + + // A previous upgrade rotation left a token behind + organizationUser.ResetPasswordKey = _resetPasswordKey; + organizationUser.V2UpgradeToken = _v2UpgradeToken; + await userRepository.UpdateUserKeyAndEncryptedDataV2Async(user, + [organizationUserRepository.UpdateForKeyRotation(user.Id, [organizationUser])]); + + organizationUser.V2UpgradeToken = null; + + // Act + await userRepository.UpdateUserKeyAndEncryptedDataV2Async(user, + [organizationUserRepository.UpdateForKeyRotation(user.Id, [organizationUser])]); + + // Assert + var updated = await organizationUserRepository.GetByIdAsync(organizationUser.Id); + Assert.NotNull(updated); + Assert.Equal(_resetPasswordKey, updated.ResetPasswordKey); + Assert.Null(updated.V2UpgradeToken); + } + + [Theory, DatabaseData] + public async Task UpdateForKeyRotation_WithOtherUsersMembership_LeavesItUnchanged( + IUserRepository userRepository, + IOrganizationRepository organizationRepository, + IOrganizationUserRepository organizationUserRepository) + { + // Arrange + var user = await userRepository.CreateTestUserAsync("rotating"); + var otherUser = await userRepository.CreateTestUserAsync("other"); + var organization = await organizationRepository.CreateTestOrganizationAsync(); + var otherOrganizationUser = + await organizationUserRepository.CreateTestOrganizationUserAsync(organization, otherUser); + + otherOrganizationUser.ResetPasswordKey = _resetPasswordKey; + otherOrganizationUser.V2UpgradeToken = _v2UpgradeToken; + await userRepository.UpdateUserKeyAndEncryptedDataV2Async(otherUser, + [organizationUserRepository.UpdateForKeyRotation(otherUser.Id, [otherOrganizationUser])]); + + // Act - the rotating user submits another member's membership, carrying their own token + otherOrganizationUser.V2UpgradeToken = _otherV2UpgradeToken; + await userRepository.UpdateUserKeyAndEncryptedDataV2Async(user, + [organizationUserRepository.UpdateForKeyRotation(user.Id, [otherOrganizationUser])]); + + // Assert - the UserId filter keeps the caller from writing onto a membership they do not own + var updated = await organizationUserRepository.GetByIdAsync(otherOrganizationUser.Id); + Assert.NotNull(updated); + Assert.Equal(_v2UpgradeToken, updated.V2UpgradeToken); + } }