From 100602344371201af7e0cf6fb8f56dbb4f8f615a Mon Sep 17 00:00:00 2001
From: Matt Czech <matt@livefront.com>
Date: Tue, 30 Apr 2024 12:13:31 -0500
Subject: [PATCH] BIT-2304: Fix pin unlock for new TDE accounts (#626)

---
 .../Core/Auth/Repositories/AuthRepository.swift           | 7 +++++++
 .../Core/Auth/Repositories/AuthRepositoryTests.swift      | 8 ++++++++
 2 files changed, 15 insertions(+)

diff --git a/BitwardenShared/Core/Auth/Repositories/AuthRepository.swift b/BitwardenShared/Core/Auth/Repositories/AuthRepository.swift
index 2bdcd2aef..803dec7fa 100644
--- a/BitwardenShared/Core/Auth/Repositories/AuthRepository.swift
+++ b/BitwardenShared/Core/Auth/Repositories/AuthRepository.swift
@@ -394,6 +394,13 @@ extension DefaultAuthRepository: AuthRepository {
             publicKey: registrationKeys.publicKey
         ))
 
+        try await stateService.setAccountEncryptionKeys(
+            AccountEncryptionKeys(
+                encryptedPrivateKey: registrationKeys.privateKey,
+                encryptedUserKey: nil
+            )
+        )
+
         try await organizationUserAPIService.organizationUserResetPasswordEnrollment(
             organizationId: enrollStatus.id,
             requestModel: OrganizationUserResetPasswordEnrollmentRequestModel(
diff --git a/BitwardenShared/Core/Auth/Repositories/AuthRepositoryTests.swift b/BitwardenShared/Core/Auth/Repositories/AuthRepositoryTests.swift
index eb4a6edc6..9523a6698 100644
--- a/BitwardenShared/Core/Auth/Repositories/AuthRepositoryTests.swift
+++ b/BitwardenShared/Core/Auth/Repositories/AuthRepositoryTests.swift
@@ -171,6 +171,10 @@ class AuthRepositoryTests: BitwardenTestCase { // swiftlint:disable:this type_bo
         XCTAssertEqual(clientService.mockAuth.makeRegisterTdeKeysEmail, "user@bitwarden.com")
         XCTAssertEqual(clientService.mockAuth.makeRegisterTdeKeysOrgPublicKey, "MIIBIjAN...2QIDAQAB")
         XCTAssertEqual(clientService.mockAuth.makeRegisterTdeKeysRememberDevice, true)
+        XCTAssertEqual(
+            stateService.accountEncryptionKeys["1"],
+            AccountEncryptionKeys(encryptedPrivateKey: "privateKey", encryptedUserKey: nil)
+        )
     }
 
     /// `createNewSsoUser()` creates a new account for sso JIT user and don't trust device.
@@ -200,6 +204,10 @@ class AuthRepositoryTests: BitwardenTestCase { // swiftlint:disable:this type_bo
         XCTAssertNil(trustDeviceService.trustDeviceWithExistingKeysValue)
         XCTAssertEqual(clientService.mockAuth.makeRegisterTdeKeysOrgPublicKey, "MIIBIjAN...2QIDAQAB")
         XCTAssertEqual(clientService.mockAuth.makeRegisterTdeKeysRememberDevice, false)
+        XCTAssertEqual(
+            stateService.accountEncryptionKeys["1"],
+            AccountEncryptionKeys(encryptedPrivateKey: "privateKey", encryptedUserKey: nil)
+        )
     }
 
     /// `deleteAccount()` deletes the active account and removes it from the state.