Add UpdateAccountCryptographicState repository function (#6669)

* Add user repository update function for account cryptographic state

* Remove comment

* Remove transaction logic

* Fix security version

* Apply feedback

* Update tests

* Add support for external actions

src/Core/KeyManagement/Models/Data/UserAccountKeysData.cs

@@ -1,9 +1,34 @@
 namespace Bit.Core.KeyManagement.Models.Data;
 
-
+/// <summary>
+/// Represents an expanded account cryptographic state for a user. Expanded here means
+/// that it does not only contain the (wrapped) private / signing key, but also the public
+/// key / verifying key. The client side only needs a subset of this data to unlock
+/// their vault and the public parts can be derived.
+/// </summary>
 public class UserAccountKeysData
 {
     public required PublicKeyEncryptionKeyPairData PublicKeyEncryptionKeyPairData { get; set; }
     public SignatureKeyPairData? SignatureKeyPairData { get; set; }
     public SecurityStateData? SecurityStateData { get; set; }
+
+    /// <summary>
+    /// Checks whether the account cryptographic state is for a V1 encryption user or a V2 encryption user.
+    /// Throws if the state is invalid
+    /// </summary>
+    public bool IsV2Encryption()
+    {
+        if (PublicKeyEncryptionKeyPairData.SignedPublicKey != null && SignatureKeyPairData != null && SecurityStateData != null)
+        {
+            return true;
+        }
+        else if (PublicKeyEncryptionKeyPairData.SignedPublicKey == null && SignatureKeyPairData == null && SecurityStateData == null)
+        {
+            return false;
+        }
+        else
+        {
+            throw new InvalidOperationException("Invalid account cryptographic state: V2 encryption fields must be either all present or all absent.");
+        }
+    }
 }

src/Core/Repositories/IUserRepository.cs

@@ -1,4 +1,5 @@
 using Bit.Core.Entities;
+using Bit.Core.KeyManagement.Models.Data;
 using Bit.Core.KeyManagement.UserKey;
 using Bit.Core.Models.Data;
 
@@ -44,5 +45,17 @@ public interface IUserRepository : IRepository<User, Guid>
         IEnumerable<UpdateEncryptedDataForKeyRotation> updateDataActions);
     Task UpdateUserKeyAndEncryptedDataV2Async(User user,
         IEnumerable<UpdateEncryptedDataForKeyRotation> updateDataActions);
+    /// <summary>
+    /// 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.
+    /// Extra actions can be passed in case other user data needs to be updated in the same transaction.
+    /// </summary>
+    Task SetV2AccountCryptographicStateAsync(
+        Guid userId,
+        UserAccountKeysData accountKeysData,
+        IEnumerable<UpdateUserData>? updateUserDataActions = null);
     Task DeleteManyAsync(IEnumerable<User> users);
 }
+
+public delegate Task UpdateUserData(Microsoft.Data.SqlClient.SqlConnection? connection = null,
+    Microsoft.Data.SqlClient.SqlTransaction? transaction = null);

src/Infrastructure.Dapper/Repositories/UserRepository.cs

@@ -2,16 +2,16 @@
 using System.Text.Json;
 using Bit.Core;
 using Bit.Core.Entities;
+using Bit.Core.KeyManagement.Models.Data;
 using Bit.Core.KeyManagement.UserKey;
 using Bit.Core.Models.Data;
 using Bit.Core.Repositories;
 using Bit.Core.Settings;
+using Bit.Core.Utilities;
 using Dapper;
 using Microsoft.AspNetCore.DataProtection;
 using Microsoft.Data.SqlClient;
 
-#nullable enable
-
 namespace Bit.Infrastructure.Dapper.Repositories;
 
 public class UserRepository : Repository<User, Guid>, IUserRepository
@@ -288,6 +288,63 @@ public class UserRepository : Repository<User, Guid>, IUserRepository
         UnprotectData(user);
     }
 
+    public async Task SetV2AccountCryptographicStateAsync(
+        Guid userId,
+        UserAccountKeysData accountKeysData,
+        IEnumerable<UpdateUserData>? updateUserDataActions = null)
+    {
+        if (!accountKeysData.IsV2Encryption())
+        {
+            throw new ArgumentException("Provided account keys data is not valid V2 encryption data.", nameof(accountKeysData));
+        }
+
+        var timestamp = DateTime.UtcNow;
+        var signatureKeyPairId = CoreHelpers.GenerateComb();
+
+        await using var connection = new SqlConnection(ConnectionString);
+        await connection.OpenAsync();
+
+        await using var transaction = connection.BeginTransaction();
+        try
+        {
+            await connection.ExecuteAsync(
+                "[dbo].[User_UpdateAccountCryptographicState]",
+                new
+                {
+                    Id = userId,
+                    PublicKey = accountKeysData.PublicKeyEncryptionKeyPairData.PublicKey,
+                    PrivateKey = accountKeysData.PublicKeyEncryptionKeyPairData.WrappedPrivateKey,
+                    SignedPublicKey = accountKeysData.PublicKeyEncryptionKeyPairData.SignedPublicKey,
+                    SecurityState = accountKeysData.SecurityStateData!.SecurityState,
+                    SecurityVersion = accountKeysData.SecurityStateData!.SecurityVersion,
+                    SignatureKeyPairId = signatureKeyPairId,
+                    SignatureAlgorithm = accountKeysData.SignatureKeyPairData!.SignatureAlgorithm,
+                    SigningKey = accountKeysData.SignatureKeyPairData!.WrappedSigningKey,
+                    VerifyingKey = accountKeysData.SignatureKeyPairData!.VerifyingKey,
+                    RevisionDate = timestamp,
+                    AccountRevisionDate = timestamp
+                },
+                transaction: transaction,
+                commandType: CommandType.StoredProcedure);
+
+            //  Update user data that depends on cryptographic state
+            if (updateUserDataActions != null)
+            {
+                foreach (var action in updateUserDataActions)
+                {
+                    await action(connection, transaction);
+                }
+            }
+
+            await transaction.CommitAsync();
+        }
+        catch
+        {
+            await transaction.RollbackAsync();
+            throw;
+        }
+    }
+
     public async Task<IEnumerable<User>> GetManyAsync(IEnumerable<Guid> ids)
     {
         using (var connection = new SqlConnection(ReadOnlyConnectionString))

src/Infrastructure.EntityFramework/Repositories/UserRepository.cs

@@ -1,4 +1,5 @@
 using AutoMapper;
+using Bit.Core.KeyManagement.Models.Data;
 using Bit.Core.KeyManagement.UserKey;
 using Bit.Core.Models.Data;
 using Bit.Core.Repositories;
@@ -241,6 +242,80 @@ public class UserRepository : Repository<Core.Entities.User, User, Guid>, IUserR
         await transaction.CommitAsync();
     }
 
+    public async Task SetV2AccountCryptographicStateAsync(
+        Guid userId,
+        UserAccountKeysData accountKeysData,
+        IEnumerable<UpdateUserData>? updateUserDataActions = null)
+    {
+        if (!accountKeysData.IsV2Encryption())
+        {
+            throw new ArgumentException("Provided account keys data is not valid V2 encryption data.", nameof(accountKeysData));
+        }
+
+        using var scope = ServiceScopeFactory.CreateScope();
+        var dbContext = GetDatabaseContext(scope);
+
+        await using var transaction = await dbContext.Database.BeginTransactionAsync();
+
+        // Update user
+        var userEntity = await dbContext.Users.FindAsync(userId);
+        if (userEntity == null)
+        {
+            throw new ArgumentException("User not found", nameof(userId));
+        }
+
+        // Update public key encryption key pair
+        var timestamp = DateTime.UtcNow;
+
+        userEntity.RevisionDate = timestamp;
+        userEntity.AccountRevisionDate = timestamp;
+
+        // V1 + V2 user crypto changes
+        userEntity.PublicKey = accountKeysData.PublicKeyEncryptionKeyPairData.PublicKey;
+        userEntity.PrivateKey = accountKeysData.PublicKeyEncryptionKeyPairData.WrappedPrivateKey;
+
+        userEntity.SecurityState = accountKeysData.SecurityStateData!.SecurityState;
+        userEntity.SecurityVersion = accountKeysData.SecurityStateData.SecurityVersion;
+        userEntity.SignedPublicKey = accountKeysData.PublicKeyEncryptionKeyPairData.SignedPublicKey;
+
+        // Replace existing keypair if it exists
+        var existingKeyPair = await dbContext.UserSignatureKeyPairs
+            .FirstOrDefaultAsync(x => x.UserId == userId);
+        if (existingKeyPair != null)
+        {
+            existingKeyPair.SignatureAlgorithm = accountKeysData.SignatureKeyPairData!.SignatureAlgorithm;
+            existingKeyPair.SigningKey = accountKeysData.SignatureKeyPairData.WrappedSigningKey;
+            existingKeyPair.VerifyingKey = accountKeysData.SignatureKeyPairData.VerifyingKey;
+            existingKeyPair.RevisionDate = timestamp;
+        }
+        else
+        {
+            var newKeyPair = new UserSignatureKeyPair
+            {
+                UserId = userId,
+                SignatureAlgorithm = accountKeysData.SignatureKeyPairData!.SignatureAlgorithm,
+                SigningKey = accountKeysData.SignatureKeyPairData.WrappedSigningKey,
+                VerifyingKey = accountKeysData.SignatureKeyPairData.VerifyingKey,
+                CreationDate = timestamp,
+                RevisionDate = timestamp
+            };
+            newKeyPair.SetNewId();
+            await dbContext.UserSignatureKeyPairs.AddAsync(newKeyPair);
+        }
+
+        await dbContext.SaveChangesAsync();
+
+        // Update additional user data within the same transaction
+        if (updateUserDataActions != null)
+        {
+            foreach (var action in updateUserDataActions)
+            {
+                await action();
+            }
+        }
+        await transaction.CommitAsync();
+    }
+
     public async Task<IEnumerable<Core.Entities.User>> GetManyAsync(IEnumerable<Guid> ids)
     {
         using (var scope = ServiceScopeFactory.CreateScope())

src/Sql/dbo/Stored Procedures/User_UpdateAccountCryptographicState.sql

@@ -0,0 +1,65 @@
+CREATE PROCEDURE [dbo].[User_UpdateAccountCryptographicState]
+    @Id UNIQUEIDENTIFIER,
+    @PublicKey NVARCHAR(MAX),
+    @PrivateKey NVARCHAR(MAX),
+    @SignedPublicKey NVARCHAR(MAX) = NULL,
+    @SecurityState NVARCHAR(MAX) = NULL,
+    @SecurityVersion INT = NULL,
+    @SignatureKeyPairId UNIQUEIDENTIFIER = NULL,
+    @SignatureAlgorithm TINYINT = NULL,
+    @SigningKey VARCHAR(MAX) = NULL,
+    @VerifyingKey VARCHAR(MAX) = NULL,
+    @RevisionDate DATETIME2(7),
+    @AccountRevisionDate DATETIME2(7)
+AS
+BEGIN
+    SET NOCOUNT ON
+
+    UPDATE
+        [dbo].[User]
+    SET
+        [PublicKey] = @PublicKey,
+        [PrivateKey] = @PrivateKey,
+        [SignedPublicKey] = @SignedPublicKey,
+        [SecurityState] = @SecurityState,
+        [SecurityVersion] = @SecurityVersion,
+        [RevisionDate] = @RevisionDate,
+        [AccountRevisionDate] = @AccountRevisionDate
+    WHERE
+        [Id] = @Id
+
+    IF EXISTS (SELECT 1 FROM [dbo].[UserSignatureKeyPair] WHERE [UserId] = @Id)
+    BEGIN
+        UPDATE [dbo].[UserSignatureKeyPair]
+        SET
+            [SignatureAlgorithm] = @SignatureAlgorithm,
+            [SigningKey] = @SigningKey,
+            [VerifyingKey] = @VerifyingKey,
+            [RevisionDate] = @RevisionDate
+        WHERE
+            [UserId] = @Id
+    END
+    ELSE
+    BEGIN
+        INSERT INTO [dbo].[UserSignatureKeyPair]
+        (
+            [Id],
+            [UserId],
+            [SignatureAlgorithm],
+            [SigningKey],
+            [VerifyingKey],
+            [CreationDate],
+            [RevisionDate]
+        )
+        VALUES
+        (
+            @SignatureKeyPairId,
+            @Id,
+            @SignatureAlgorithm,
+            @SigningKey,
+            @VerifyingKey,
+            @RevisionDate,
+            @RevisionDate
+        )
+    END
+END

test/Infrastructure.EFIntegration.Test/Repositories/UserRepositoryTests.cs

@@ -2,6 +2,8 @@
 using Bit.Core.Auth.Entities;
 using Bit.Core.Entities;
 using Bit.Core.Enums;
+using Bit.Core.KeyManagement.Enums;
+using Bit.Core.KeyManagement.Models.Data;
 using Bit.Core.Models.Data;
 using Bit.Core.Test.AutoFixture.Attributes;
 using Bit.Infrastructure.EFIntegration.Test.AutoFixture;
@@ -313,4 +315,66 @@ public class UserRepositoryTests
         Assert.Equal(sqlUser.MasterPasswordHint, updatedUser.MasterPasswordHint);
         Assert.Equal(sqlUser.Email, updatedUser.Email);
     }
+
+    [CiSkippedTheory, EfUserAutoData]
+    public async Task UpdateAccountCryptographicStateAsync_Works_DataMatches(
+        User user,
+        List<EfRepo.UserRepository> suts,
+        SqlRepo.UserRepository sqlUserRepo)
+    {
+        // Test for V1 user (no signature key pair or security state)
+        var accountKeysDataV1 = new UserAccountKeysData
+        {
+            PublicKeyEncryptionKeyPairData = new PublicKeyEncryptionKeyPairData(
+                wrappedPrivateKey: "v1-wrapped-private-key",
+                publicKey: "v1-public-key"
+            )
+        };
+
+        foreach (var sut in suts)
+        {
+            var createdUser = await sut.CreateAsync(user);
+            sut.ClearChangeTracking();
+
+            await sut.SetV2AccountCryptographicStateAsync(createdUser.Id, accountKeysDataV1);
+            sut.ClearChangeTracking();
+
+            var updatedUser = await sut.GetByIdAsync(createdUser.Id);
+            Assert.Equal("v1-public-key", updatedUser.PublicKey);
+            Assert.Equal("v1-wrapped-private-key", updatedUser.PrivateKey);
+            Assert.Null(updatedUser.SignedPublicKey);
+            Assert.Null(updatedUser.SecurityState);
+            Assert.Null(updatedUser.SecurityVersion);
+        }
+
+        // Test for V2 user (with signature key pair and security state)
+        var accountKeysDataV2 = new UserAccountKeysData
+        {
+            PublicKeyEncryptionKeyPairData = new PublicKeyEncryptionKeyPairData(
+                wrappedPrivateKey: "v2-wrapped-private-key",
+                publicKey: "v2-public-key",
+                signedPublicKey: "v2-signed-public-key"
+            ),
+            SignatureKeyPairData = new SignatureKeyPairData(
+                signatureAlgorithm: SignatureAlgorithm.Ed25519,
+                wrappedSigningKey: "v2-wrapped-signing-key",
+                verifyingKey: "v2-verifying-key"
+            ),
+            SecurityStateData = new SecurityStateData
+            {
+                SecurityState = "v2-security-state",
+                SecurityVersion = 2
+            }
+        };
+
+        var sqlUser = await sqlUserRepo.CreateAsync(user);
+        await sqlUserRepo.SetV2AccountCryptographicStateAsync(sqlUser.Id, accountKeysDataV2);
+
+        var updatedSqlUser = await sqlUserRepo.GetByIdAsync(sqlUser.Id);
+        Assert.Equal("v2-public-key", updatedSqlUser.PublicKey);
+        Assert.Equal("v2-wrapped-private-key", updatedSqlUser.PrivateKey);
+        Assert.Equal("v2-signed-public-key", updatedSqlUser.SignedPublicKey);
+        Assert.Equal("v2-security-state", updatedSqlUser.SecurityState);
+        Assert.Equal(2, updatedSqlUser.SecurityVersion);
+    }
 }

util/Migrator/DbScripts/2025-12-08_00_User_UpdateAccountCryptographicState.sql

@@ -0,0 +1,72 @@
+IF OBJECT_ID('[dbo].[User_UpdateAccountCryptographicState]') IS NOT NULL
+BEGIN
+    DROP PROCEDURE [dbo].[User_UpdateAccountCryptographicState]
+END
+GO
+
+CREATE PROCEDURE [dbo].[User_UpdateAccountCryptographicState]
+    @Id UNIQUEIDENTIFIER,
+    @PublicKey NVARCHAR(MAX),
+    @PrivateKey NVARCHAR(MAX),
+    @SignedPublicKey NVARCHAR(MAX) = NULL,
+    @SecurityState NVARCHAR(MAX) = NULL,
+    @SecurityVersion INT = NULL,
+    @SignatureKeyPairId UNIQUEIDENTIFIER = NULL,
+    @SignatureAlgorithm TINYINT = NULL,
+    @SigningKey VARCHAR(MAX) = NULL,
+    @VerifyingKey VARCHAR(MAX) = NULL,
+    @RevisionDate DATETIME2(7),
+    @AccountRevisionDate DATETIME2(7)
+AS
+BEGIN
+    SET NOCOUNT ON
+
+    UPDATE
+        [dbo].[User]
+    SET
+        [PublicKey] = @PublicKey,
+        [PrivateKey] = @PrivateKey,
+        [SignedPublicKey] = @SignedPublicKey,
+        [SecurityState] = @SecurityState,
+        [SecurityVersion] = @SecurityVersion,
+        [RevisionDate] = @RevisionDate,
+        [AccountRevisionDate] = @AccountRevisionDate
+    WHERE
+        [Id] = @Id
+
+    IF EXISTS (SELECT 1 FROM [dbo].[UserSignatureKeyPair] WHERE [UserId] = @Id)
+    BEGIN
+        UPDATE [dbo].[UserSignatureKeyPair]
+        SET
+            [SignatureAlgorithm] = @SignatureAlgorithm,
+            [SigningKey] = @SigningKey,
+            [VerifyingKey] = @VerifyingKey,
+            [RevisionDate] = @RevisionDate
+        WHERE
+            [UserId] = @Id
+    END
+    ELSE
+    BEGIN
+        INSERT INTO [dbo].[UserSignatureKeyPair]
+        (
+            [Id],
+            [UserId],
+            [SignatureAlgorithm],
+            [SigningKey],
+            [VerifyingKey],
+            [CreationDate],
+            [RevisionDate]
+        )
+        VALUES
+        (
+            @SignatureKeyPairId,
+            @Id,
+            @SignatureAlgorithm,
+            @SigningKey,
+            @VerifyingKey,
+            @RevisionDate,
+            @RevisionDate
+        )
+    END
+END
+GO