Add UpdateAccountCryptographicState repository function

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,10 @@ Task UpdateUserKeyAndEncryptedDataAsync(User user,
         IEnumerable<UpdateEncryptedDataForKeyRotation> updateDataActions);
     Task UpdateUserKeyAndEncryptedDataV2Async(User user,
         IEnumerable<UpdateEncryptedDataForKeyRotation> updateDataActions);
+    /// <summary>
+    /// Sets the account cryptographic state to a user in a single transaction. V1 and V2 account cryptographic states are supported,
+    /// and the data is created or replaced depending on whether it exists already.
+    /// </summary>
+    Task UpdateAccountCryptographicStateAsync(User user, UserAccountKeysData accountKeysData);
     Task DeleteManyAsync(IEnumerable<User> users);
 }

src/Infrastructure.Dapper/Repositories/UserRepository.cs

@@ -2,6 +2,7 @@
 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;
@@ -10,8 +11,6 @@
 using Microsoft.AspNetCore.DataProtection;
 using Microsoft.Data.SqlClient;
 
-#nullable enable
-
 namespace Bit.Infrastructure.Dapper.Repositories;
 
 public class UserRepository : Repository<User, Guid>, IUserRepository
@@ -288,6 +287,46 @@ await connection.ExecuteAsync(
         UnprotectData(user);
     }
 
+    public async Task UpdateAccountCryptographicStateAsync(User user, UserAccountKeysData accountKeysData)
+    {
+        await using var connection = new SqlConnection(ConnectionString);
+        await connection.OpenAsync();
+
+        await using var transaction = await connection.BeginTransactionAsync();
+        try
+        {
+            await using (var cmd = new SqlCommand("[dbo].[User_UpdateAccountCryptographicState]", connection, (SqlTransaction)transaction))
+            {
+                cmd.CommandType = CommandType.StoredProcedure;
+                cmd.Parameters.Add("@Id", SqlDbType.UniqueIdentifier).Value = user.Id;
+                cmd.Parameters.Add("@PublicKey", SqlDbType.NVarChar).Value = accountKeysData.PublicKeyEncryptionKeyPairData.PublicKey;
+                cmd.Parameters.Add("@PrivateKey", SqlDbType.NVarChar).Value = accountKeysData.PublicKeyEncryptionKeyPairData.WrappedPrivateKey;
+                cmd.Parameters.Add("@SignedPublicKey", SqlDbType.NVarChar).Value =
+                    accountKeysData.PublicKeyEncryptionKeyPairData.SignedPublicKey ?? (object)DBNull.Value;
+                cmd.Parameters.Add("@SecurityState", SqlDbType.NVarChar).Value =
+                    accountKeysData.SecurityStateData?.SecurityState ?? (object)DBNull.Value;
+                cmd.Parameters.Add("@SecurityVersion", SqlDbType.Int).Value =
+                    accountKeysData.SecurityStateData?.SecurityVersion ?? (object)DBNull.Value;
+                cmd.Parameters.Add("@SignatureAlgorithm", SqlDbType.TinyInt).Value =
+                    accountKeysData.SignatureKeyPairData != null ? (object)accountKeysData.SignatureKeyPairData.SignatureAlgorithm : DBNull.Value;
+                cmd.Parameters.Add("@SigningKey", SqlDbType.VarChar).Value =
+                    accountKeysData.SignatureKeyPairData?.WrappedSigningKey ?? (object)DBNull.Value;
+                cmd.Parameters.Add("@VerifyingKey", SqlDbType.VarChar).Value =
+                    accountKeysData.SignatureKeyPairData?.VerifyingKey ?? (object)DBNull.Value;
+                cmd.Parameters.Add("@RevisionDate", SqlDbType.DateTime2).Value = user.RevisionDate;
+                cmd.Parameters.Add("@AccountRevisionDate", SqlDbType.DateTime2).Value = user.AccountRevisionDate;
+                await cmd.ExecuteNonQueryAsync();
+            }
+
+            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,66 @@ public async Task UpdateUserKeyAndEncryptedDataV2Async(Core.Entities.User user,
         await transaction.CommitAsync();
     }
 
+    public async Task UpdateAccountCryptographicStateAsync(Core.Entities.User user, UserAccountKeysData 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(user.Id);
+        if (userEntity == null)
+        {
+            throw new ArgumentException("User not found", nameof(user));
+        }
+
+        // Update public key encryption key pair
+
+        userEntity.RevisionDate = user.RevisionDate;
+        userEntity.AccountRevisionDate = user.AccountRevisionDate;
+
+        // V1 + V2 user crypto changes
+        userEntity.PublicKey = accountKeysData.PublicKeyEncryptionKeyPairData.PublicKey;
+        userEntity.PrivateKey = accountKeysData.PublicKeyEncryptionKeyPairData.WrappedPrivateKey;
+
+        // V2 only changes
+        if (accountKeysData.IsV2Encryption())
+        {
+            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 == user.Id);
+            if (existingKeyPair != null)
+            {
+                existingKeyPair.SignatureAlgorithm = accountKeysData.SignatureKeyPairData!.SignatureAlgorithm;
+                existingKeyPair.SigningKey = accountKeysData.SignatureKeyPairData.WrappedSigningKey;
+                existingKeyPair.VerifyingKey = accountKeysData.SignatureKeyPairData.VerifyingKey;
+                existingKeyPair.RevisionDate = user.RevisionDate;
+            }
+            else
+            {
+                var newKeyPair = new UserSignatureKeyPair
+                {
+                    Id = Guid.NewGuid(),
+                    UserId = user.Id,
+                    SignatureAlgorithm = accountKeysData.SignatureKeyPairData!.SignatureAlgorithm,
+                    SigningKey = accountKeysData.SignatureKeyPairData.WrappedSigningKey,
+                    VerifyingKey = accountKeysData.SignatureKeyPairData.VerifyingKey,
+                    CreationDate = user.RevisionDate,
+                    RevisionDate = user.RevisionDate
+                };
+                await dbContext.UserSignatureKeyPairs.AddAsync(newKeyPair);
+            }
+        }
+
+        await dbContext.SaveChangesAsync();
+        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,68 @@
+CREATE PROCEDURE [dbo].[User_UpdateAccountCryptographicState]
+    @Id UNIQUEIDENTIFIER,
+    @PublicKey NVARCHAR(MAX),
+    @PrivateKey NVARCHAR(MAX),
+    @SignedPublicKey NVARCHAR(MAX) = NULL,
+    @SecurityState NVARCHAR(MAX) = NULL,
+    @SecurityVersion INT = 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
+
+    -- Update or insert signature key pair if provided
+    IF @SignatureAlgorithm IS NOT NULL AND @SigningKey IS NOT NULL AND @VerifyingKey IS NOT NULL
+    BEGIN
+        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
+            (
+                NEWID(),
+                @Id,
+                @SignatureAlgorithm,
+                @SigningKey,
+                @VerifyingKey,
+                @RevisionDate,
+                @RevisionDate
+            )
+        END
+    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,72 @@ public async Task UpdateUserKeyAndEncryptedDataAsync_Works_DataMatches(User user
         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();
+
+            createdUser.RevisionDate = DateTime.UtcNow;
+            createdUser.AccountRevisionDate = DateTime.UtcNow;
+
+            await sut.UpdateAccountCryptographicStateAsync(createdUser, 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);
+        sqlUser.RevisionDate = DateTime.UtcNow;
+        sqlUser.AccountRevisionDate = DateTime.UtcNow;
+
+        await sqlUserRepo.UpdateAccountCryptographicStateAsync(sqlUser, 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);
+    }
 }