From bb61b3df3ae9d40bdfb002c02a3c5f69695ef1e0 Mon Sep 17 00:00:00 2001
From: Thomas Avery <43214426+Thomas-Avery@users.noreply.github.com>
Date: Wed, 8 Jan 2025 16:41:02 -0600
Subject: [PATCH] [PM-15940] Add regen to SSO login (#12643)

* Add loginSuccessHandlerService to SSO login component

* Update regen service to handle SSO login
---
 libs/auth/src/angular/sso/sso.component.ts    |  7 ++-
 ...symmetric-key-regeneration.service.spec.ts | 50 +++++++++++++++++++
 ...ser-asymmetric-key-regeneration.service.ts | 21 +++++++-
 3 files changed, 72 insertions(+), 6 deletions(-)

diff --git a/libs/auth/src/angular/sso/sso.component.ts b/libs/auth/src/angular/sso/sso.component.ts
index aad0df4e39..3bcc56ae4c 100644
--- a/libs/auth/src/angular/sso/sso.component.ts
+++ b/libs/auth/src/angular/sso/sso.component.ts
@@ -12,6 +12,7 @@ import {
   TrustedDeviceUserDecryptionOption,
   UserDecryptionOptions,
   UserDecryptionOptionsServiceAbstraction,
+  LoginSuccessHandlerService,
 } from "@bitwarden/auth/common";
 import { ApiService } from "@bitwarden/common/abstractions/api.service";
 import { OrgDomainApiServiceAbstraction } from "@bitwarden/common/admin-console/abstractions/organization-domain/org-domain-api.service.abstraction";
@@ -35,7 +36,6 @@ import { LogService } from "@bitwarden/common/platform/abstractions/log.service"
 import { PlatformUtilsService } from "@bitwarden/common/platform/abstractions/platform-utils.service";
 import { ValidationService } from "@bitwarden/common/platform/abstractions/validation.service";
 import { Utils } from "@bitwarden/common/platform/misc/utils";
-import { SyncService } from "@bitwarden/common/vault/abstractions/sync/sync.service.abstraction";
 import {
   AsyncActionsModule,
   ButtonModule,
@@ -117,7 +117,7 @@ export class SsoComponent implements OnInit {
     private accountService: AccountService,
     private toastService: ToastService,
     private ssoComponentService: SsoComponentService,
-    private syncService: SyncService,
+    private loginSuccessHandlerService: LoginSuccessHandlerService,
   ) {
     environmentService.environment$.pipe(takeUntilDestroyed()).subscribe((env) => {
       this.redirectUri = env.getWebVaultUrl() + "/sso-connector.html";
@@ -378,8 +378,7 @@ export class SsoComponent implements OnInit {
 
       // Everything after the 2FA check is considered a successful login
       // Just have to figure out where to send the user
-
-      await this.syncService.fullSync(true);
+      await this.loginSuccessHandlerService.run(authResult.userId);
 
       // Save off the OrgSsoIdentifier for use in the TDE flows (or elsewhere)
       // - TDE login decryption options component
diff --git a/libs/key-management/src/user-asymmetric-key-regeneration/services/default-user-asymmetric-key-regeneration.service.spec.ts b/libs/key-management/src/user-asymmetric-key-regeneration/services/default-user-asymmetric-key-regeneration.service.spec.ts
index 77d7ebbb81..5312113d76 100644
--- a/libs/key-management/src/user-asymmetric-key-regeneration/services/default-user-asymmetric-key-regeneration.service.spec.ts
+++ b/libs/key-management/src/user-asymmetric-key-regeneration/services/default-user-asymmetric-key-regeneration.service.spec.ts
@@ -153,6 +153,56 @@ describe("regenerateIfNeeded", () => {
     expect(keyService.setPrivateKey).not.toHaveBeenCalled();
   });
 
+  it("should not regenerate when user symmetric key is unavailable", async () => {
+    const mockVerificationResponse: VerifyAsymmetricKeysResponse = {
+      privateKeyDecryptable: true,
+      validPrivateKey: false,
+    };
+    setupVerificationResponse(mockVerificationResponse, sdkService);
+    keyService.userKey$.mockReturnValue(of(undefined as unknown as UserKey));
+
+    await sut.regenerateIfNeeded(userId);
+
+    expect(
+      userAsymmetricKeysRegenerationApiService.regenerateUserAsymmetricKeys,
+    ).not.toHaveBeenCalled();
+    expect(keyService.setPrivateKey).not.toHaveBeenCalled();
+  });
+
+  it("should not regenerate when user's encrypted private key is unavailable", async () => {
+    const mockVerificationResponse: VerifyAsymmetricKeysResponse = {
+      privateKeyDecryptable: true,
+      validPrivateKey: false,
+    };
+    setupVerificationResponse(mockVerificationResponse, sdkService);
+    keyService.userEncryptedPrivateKey$.mockReturnValue(
+      of(undefined as unknown as EncryptedString),
+    );
+
+    await sut.regenerateIfNeeded(userId);
+
+    expect(
+      userAsymmetricKeysRegenerationApiService.regenerateUserAsymmetricKeys,
+    ).not.toHaveBeenCalled();
+    expect(keyService.setPrivateKey).not.toHaveBeenCalled();
+  });
+
+  it("should not regenerate when user's public key is unavailable", async () => {
+    const mockVerificationResponse: VerifyAsymmetricKeysResponse = {
+      privateKeyDecryptable: true,
+      validPrivateKey: false,
+    };
+    setupVerificationResponse(mockVerificationResponse, sdkService);
+    apiService.getUserPublicKey.mockResolvedValue(undefined as any);
+
+    await sut.regenerateIfNeeded(userId);
+
+    expect(
+      userAsymmetricKeysRegenerationApiService.regenerateUserAsymmetricKeys,
+    ).not.toHaveBeenCalled();
+    expect(keyService.setPrivateKey).not.toHaveBeenCalled();
+  });
+
   it("should regenerate when private key is decryptable and invalid", async () => {
     const mockVerificationResponse: VerifyAsymmetricKeysResponse = {
       privateKeyDecryptable: true,
diff --git a/libs/key-management/src/user-asymmetric-key-regeneration/services/default-user-asymmetric-key-regeneration.service.ts b/libs/key-management/src/user-asymmetric-key-regeneration/services/default-user-asymmetric-key-regeneration.service.ts
index ffaa3a8260..3110ebad63 100644
--- a/libs/key-management/src/user-asymmetric-key-regeneration/services/default-user-asymmetric-key-regeneration.service.ts
+++ b/libs/key-management/src/user-asymmetric-key-regeneration/services/default-user-asymmetric-key-regeneration.service.ts
@@ -49,14 +49,31 @@ export class DefaultUserAsymmetricKeysRegenerationService
   }
 
   private async shouldRegenerate(userId: UserId): Promise<boolean> {
-    const [userKey, userKeyEncryptedPrivateKey, publicKeyResponse] = await firstValueFrom(
+    const userKey = await firstValueFrom(this.keyService.userKey$(userId));
+
+    // For SSO logins from untrusted devices, the userKey will not be available, and the private key regeneration process should be skipped.
+    // In such cases, regeneration will occur on the following device login flow.
+    if (!userKey) {
+      this.logService.info(
+        "[UserAsymmetricKeyRegeneration] User symmetric key unavailable, skipping regeneration for the user.",
+      );
+      return false;
+    }
+
+    const [userKeyEncryptedPrivateKey, publicKeyResponse] = await firstValueFrom(
       combineLatest([
-        this.keyService.userKey$(userId),
         this.keyService.userEncryptedPrivateKey$(userId),
         this.apiService.getUserPublicKey(userId),
       ]),
     );
 
+    if (!userKeyEncryptedPrivateKey || !publicKeyResponse) {
+      this.logService.warning(
+        "[UserAsymmetricKeyRegeneration] User's asymmetric key initialization data is unavailable, skipping regeneration.",
+      );
+      return false;
+    }
+
     const verificationResponse = await firstValueFrom(
       this.sdkService.client$.pipe(
         map((sdk) => {