Merge branch 'main' into auth/add-logout-reason

2026-02-10 05:30:01 +00:00 · 2026-01-11 11:24:18 -05:00
parent 449ed97b9e 404d925f84
commit a7902b2ba6
1569 changed files with 125258 additions and 31664 deletions
--- a/libs/auth/src/angular/input-password/input-password.mdx
+++ b/libs/auth/src/angular/input-password/input-password.mdx
@@ -1,4 +1,4 @@
-import { Meta, Story } from "@storybook/addon-docs";
+import { Meta, Story } from "@storybook/addon-docs/blocks";

 import * as stories from "./input-password.stories.ts";

--- a/libs/auth/src/angular/input-password/input-password.stories.ts
+++ b/libs/auth/src/angular/input-password/input-password.stories.ts
@@ -1,8 +1,8 @@
 import { importProvidersFrom } from "@angular/core";
 import { BrowserAnimationsModule } from "@angular/platform-browser/animations";
-import { action } from "@storybook/addon-actions";
 import { Meta, StoryObj, applicationConfig } from "@storybook/angular";
 import { of } from "rxjs";
+import { action } from "storybook/actions";
 import { ZXCVBNResult } from "zxcvbn";

 import { AuditService } from "@bitwarden/common/abstractions/audit.service";
--- a/libs/auth/src/angular/login-decryption-options/login-decryption-options.component.spec.ts
+++ b/libs/auth/src/angular/login-decryption-options/login-decryption-options.component.spec.ts
@@ -0,0 +1,377 @@
+// Mock asUuid to return the input value for test consistency
+jest.mock("@bitwarden/common/platform/abstractions/sdk/sdk.service", () => ({
+  asUuid: (x: any) => x,
+}));
+
+import { DestroyRef } from "@angular/core";
+import { FormBuilder } from "@angular/forms";
+import { Router } from "@angular/router";
+import { mock, MockProxy } from "jest-mock-extended";
+import { BehaviorSubject, of } from "rxjs";
+
+import {
+  LoginEmailServiceAbstraction,
+  LogoutService,
+  UserDecryptionOptionsServiceAbstraction,
+} from "@bitwarden/auth/common";
+import { ApiService } from "@bitwarden/common/abstractions/api.service";
+import { OrganizationApiServiceAbstraction } from "@bitwarden/common/admin-console/abstractions/organization/organization-api.service.abstraction";
+import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
+import { PasswordResetEnrollmentServiceAbstraction } from "@bitwarden/common/auth/abstractions/password-reset-enrollment.service.abstraction";
+import { SsoLoginServiceAbstraction } from "@bitwarden/common/auth/abstractions/sso-login.service.abstraction";
+import { ClientType } from "@bitwarden/common/enums";
+import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
+import { DeviceTrustServiceAbstraction } from "@bitwarden/common/key-management/device-trust/abstractions/device-trust.service.abstraction";
+import { SecurityStateService } from "@bitwarden/common/key-management/security-state/abstractions/security-state.service";
+import { SignedSecurityState } from "@bitwarden/common/key-management/types";
+import { AppIdService } from "@bitwarden/common/platform/abstractions/app-id.service";
+import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
+import { I18nService } from "@bitwarden/common/platform/abstractions/i18n.service";
+import { MessagingService } from "@bitwarden/common/platform/abstractions/messaging.service";
+import { PlatformUtilsService } from "@bitwarden/common/platform/abstractions/platform-utils.service";
+import { RegisterSdkService } from "@bitwarden/common/platform/abstractions/sdk/register-sdk.service";
+import { ValidationService } from "@bitwarden/common/platform/abstractions/validation.service";
+import { SymmetricCryptoKey } from "@bitwarden/common/platform/models/domain/symmetric-crypto-key";
+import { UserId } from "@bitwarden/common/types/guid";
+// eslint-disable-next-line no-restricted-imports
+import { AnonLayoutWrapperDataService, DialogService, ToastService } from "@bitwarden/components";
+import { KeyService } from "@bitwarden/key-management";
+
+import { LoginDecryptionOptionsComponent } from "./login-decryption-options.component";
+import { LoginDecryptionOptionsService } from "./login-decryption-options.service";
+
+describe("LoginDecryptionOptionsComponent", () => {
+  let component: LoginDecryptionOptionsComponent;
+  let accountService: MockProxy<AccountService>;
+  let anonLayoutWrapperDataService: MockProxy<AnonLayoutWrapperDataService>;
+  let apiService: MockProxy<ApiService>;
+  let destroyRef: MockProxy<DestroyRef>;
+  let deviceTrustService: MockProxy<DeviceTrustServiceAbstraction>;
+  let dialogService: MockProxy<DialogService>;
+  let formBuilder: FormBuilder;
+  let i18nService: MockProxy<I18nService>;
+  let keyService: MockProxy<KeyService>;
+  let loginDecryptionOptionsService: MockProxy<LoginDecryptionOptionsService>;
+  let loginEmailService: MockProxy<LoginEmailServiceAbstraction>;
+  let messagingService: MockProxy<MessagingService>;
+  let organizationApiService: MockProxy<OrganizationApiServiceAbstraction>;
+  let passwordResetEnrollmentService: MockProxy<PasswordResetEnrollmentServiceAbstraction>;
+  let platformUtilsService: MockProxy<PlatformUtilsService>;
+  let router: MockProxy<Router>;
+  let ssoLoginService: MockProxy<SsoLoginServiceAbstraction>;
+  let toastService: MockProxy<ToastService>;
+  let userDecryptionOptionsService: MockProxy<UserDecryptionOptionsServiceAbstraction>;
+  let validationService: MockProxy<ValidationService>;
+  let logoutService: MockProxy<LogoutService>;
+  let registerSdkService: MockProxy<RegisterSdkService>;
+  let securityStateService: MockProxy<SecurityStateService>;
+  let appIdService: MockProxy<AppIdService>;
+  let configService: MockProxy<ConfigService>;
+  let accountCryptographicStateService: MockProxy<any>;
+
+  const mockUserId = "user-id-123" as UserId;
+  const mockEmail = "test@example.com";
+  const mockOrgId = "org-id-456";
+
+  beforeEach(() => {
+    accountService = mock<AccountService>();
+    anonLayoutWrapperDataService = mock<AnonLayoutWrapperDataService>();
+    apiService = mock<ApiService>();
+    destroyRef = mock<DestroyRef>();
+    deviceTrustService = mock<DeviceTrustServiceAbstraction>();
+    dialogService = mock<DialogService>();
+    formBuilder = new FormBuilder();
+    i18nService = mock<I18nService>();
+    keyService = mock<KeyService>();
+    loginDecryptionOptionsService = mock<LoginDecryptionOptionsService>();
+    loginEmailService = mock<LoginEmailServiceAbstraction>();
+    messagingService = mock<MessagingService>();
+    organizationApiService = mock<OrganizationApiServiceAbstraction>();
+    passwordResetEnrollmentService = mock<PasswordResetEnrollmentServiceAbstraction>();
+    platformUtilsService = mock<PlatformUtilsService>();
+    router = mock<Router>();
+    ssoLoginService = mock<SsoLoginServiceAbstraction>();
+    toastService = mock<ToastService>();
+    userDecryptionOptionsService = mock<UserDecryptionOptionsServiceAbstraction>();
+    validationService = mock<ValidationService>();
+    logoutService = mock<LogoutService>();
+    registerSdkService = mock<RegisterSdkService>();
+    securityStateService = mock<SecurityStateService>();
+    appIdService = mock<AppIdService>();
+    configService = mock<ConfigService>();
+    accountCryptographicStateService = mock();
+
+    // Setup default mocks
+    accountService.activeAccount$ = new BehaviorSubject({
+      id: mockUserId,
+      email: mockEmail,
+      name: "Test User",
+      emailVerified: true,
+      creationDate: new Date(),
+    });
+    platformUtilsService.getClientType.mockReturnValue(ClientType.Browser);
+    deviceTrustService.getShouldTrustDevice.mockResolvedValue(true);
+    i18nService.t.mockImplementation((key: string) => key);
+
+    component = new LoginDecryptionOptionsComponent(
+      accountService,
+      anonLayoutWrapperDataService,
+      apiService,
+      destroyRef,
+      deviceTrustService,
+      dialogService,
+      formBuilder,
+      i18nService,
+      keyService,
+      loginDecryptionOptionsService,
+      loginEmailService,
+      messagingService,
+      organizationApiService,
+      passwordResetEnrollmentService,
+      platformUtilsService,
+      router,
+      ssoLoginService,
+      toastService,
+      userDecryptionOptionsService,
+      validationService,
+      logoutService,
+      registerSdkService,
+      securityStateService,
+      appIdService,
+      configService,
+      accountCryptographicStateService,
+    );
+  });
+
+  describe("createUser with feature flag enabled", () => {
+    let mockPostKeysForTdeRegistration: jest.Mock;
+    let mockRegistration: any;
+    let mockAuth: any;
+    let mockSdkValue: any;
+    let mockSdkRef: any;
+    let mockSdk: any;
+    let mockDeviceKey: string;
+    let mockDeviceKeyObj: SymmetricCryptoKey;
+    let mockUserKeyBytes: Uint8Array;
+    let mockPrivateKey: string;
+    let mockSignedPublicKey: string;
+    let mockSigningKey: string;
+    let mockSecurityState: SignedSecurityState;
+
+    beforeEach(async () => {
+      // Mock asUuid to return the input value for test consistency
+      jest.mock("@bitwarden/common/platform/abstractions/sdk/sdk.service", () => ({
+        asUuid: (x: any) => x,
+      }));
+      (Symbol as any).dispose = Symbol("dispose");
+
+      mockPrivateKey = "mock-private-key";
+      mockSignedPublicKey = "mock-signed-public-key";
+      mockSigningKey = "mock-signing-key";
+      mockSecurityState = {
+        signature: "mock-signature",
+        payload: {
+          version: 2,
+          timestamp: Date.now(),
+          privateKeyHash: "mock-hash",
+        },
+      } as any;
+      const deviceKeyBytes = new Uint8Array(32).fill(5);
+      mockDeviceKey = Buffer.from(deviceKeyBytes).toString("base64");
+      mockDeviceKeyObj = SymmetricCryptoKey.fromString(mockDeviceKey);
+      mockUserKeyBytes = new Uint8Array(64);
+
+      mockPostKeysForTdeRegistration = jest.fn().mockResolvedValue({
+        account_cryptographic_state: {
+          V2: {
+            private_key: mockPrivateKey,
+            signed_public_key: mockSignedPublicKey,
+            signing_key: mockSigningKey,
+            security_state: mockSecurityState,
+          },
+        },
+        device_key: mockDeviceKey,
+        user_key: mockUserKeyBytes,
+      });
+
+      mockRegistration = {
+        post_keys_for_tde_registration: mockPostKeysForTdeRegistration,
+      };
+
+      mockAuth = {
+        registration: jest.fn().mockReturnValue(mockRegistration),
+      };
+
+      mockSdkValue = {
+        auth: jest.fn().mockReturnValue(mockAuth),
+      };
+
+      mockSdkRef = {
+        value: mockSdkValue,
+        [Symbol.dispose]: jest.fn(),
+      };
+
+      mockSdk = {
+        take: jest.fn().mockReturnValue(mockSdkRef),
+      };
+
+      registerSdkService.registerClient$ = jest.fn((userId: UserId) => of(mockSdk)) as any;
+
+      // Setup for new user state
+      userDecryptionOptionsService.userDecryptionOptionsById$.mockReturnValue(
+        of({
+          trustedDeviceOption: {
+            hasAdminApproval: false,
+            hasLoginApprovingDevice: false,
+            hasManageResetPasswordPermission: false,
+            isTdeOffboarding: false,
+          },
+          hasMasterPassword: false,
+          keyConnectorOption: undefined,
+        }),
+      );
+
+      ssoLoginService.getActiveUserOrganizationSsoIdentifier.mockResolvedValue("org-identifier");
+      organizationApiService.getAutoEnrollStatus.mockResolvedValue({
+        id: mockOrgId,
+        resetPasswordEnabled: true,
+      } as any);
+
+      // Initialize component to set up new user state
+      await component.ngOnInit();
+    });
+
+    it("should use SDK v2 registration when feature flag is enabled", async () => {
+      // Arrange
+      configService.getFeatureFlag.mockResolvedValue(true);
+      loginDecryptionOptionsService.handleCreateUserSuccess.mockResolvedValue(undefined);
+      router.navigate.mockResolvedValue(true);
+      appIdService.getAppId.mockResolvedValue("mock-app-id");
+      organizationApiService.getKeys.mockResolvedValue({
+        publicKey: "mock-org-public-key",
+        privateKey: "mock-org-private-key",
+      } as any);
+
+      // Act
+      await component["createUser"]();
+
+      // Assert
+      expect(configService.getFeatureFlag).toHaveBeenCalledWith(
+        FeatureFlag.PM27279_V2RegistrationTdeJit,
+      );
+      expect(appIdService.getAppId).toHaveBeenCalled();
+      expect(organizationApiService.getKeys).toHaveBeenCalledWith(mockOrgId);
+      expect(registerSdkService.registerClient$).toHaveBeenCalledWith(mockUserId);
+
+      // Verify SDK registration was called with correct parameters
+      expect(mockSdkValue.auth).toHaveBeenCalled();
+      expect(mockAuth.registration).toHaveBeenCalled();
+      expect(mockPostKeysForTdeRegistration).toHaveBeenCalledWith({
+        org_id: mockOrgId,
+        org_public_key: "mock-org-public-key",
+        user_id: mockUserId,
+        device_identifier: "mock-app-id",
+        trust_device: true,
+      });
+
+      const expectedDeviceKey = mockDeviceKeyObj;
+      const expectedUserKey = new SymmetricCryptoKey(new Uint8Array(mockUserKeyBytes));
+
+      // Verify keys were set
+      expect(keyService.setPrivateKey).toHaveBeenCalledWith(mockPrivateKey, mockUserId);
+      expect(keyService.setSignedPublicKey).toHaveBeenCalledWith(mockSignedPublicKey, mockUserId);
+      expect(keyService.setUserSigningKey).toHaveBeenCalledWith(mockSigningKey, mockUserId);
+      expect(securityStateService.setAccountSecurityState).toHaveBeenCalledWith(
+        mockSecurityState,
+        mockUserId,
+      );
+      expect(accountCryptographicStateService.setAccountCryptographicState).toHaveBeenCalledWith(
+        expect.objectContaining({
+          V2: {
+            private_key: mockPrivateKey,
+            signed_public_key: mockSignedPublicKey,
+            signing_key: mockSigningKey,
+            security_state: mockSecurityState,
+          },
+        }),
+        mockUserId,
+      );
+
+      expect(validationService.showError).not.toHaveBeenCalled();
+
+      // Verify device and user keys were persisted
+      expect(deviceTrustService.setDeviceKey).toHaveBeenCalledWith(
+        mockUserId,
+        expect.any(SymmetricCryptoKey),
+      );
+      expect(keyService.setUserKey).toHaveBeenCalledWith(
+        expect.any(SymmetricCryptoKey),
+        mockUserId,
+      );
+
+      const [, deviceKeyArg] = deviceTrustService.setDeviceKey.mock.calls[0];
+      const [userKeyArg] = keyService.setUserKey.mock.calls[0];
+
+      expect((deviceKeyArg as SymmetricCryptoKey).keyB64).toBe(expectedDeviceKey.keyB64);
+      expect((userKeyArg as SymmetricCryptoKey).keyB64).toBe(expectedUserKey.keyB64);
+
+      // Verify success toast and navigation
+      expect(toastService.showToast).toHaveBeenCalledWith({
+        variant: "success",
+        title: null,
+        message: "accountSuccessfullyCreated",
+      });
+      expect(loginDecryptionOptionsService.handleCreateUserSuccess).toHaveBeenCalled();
+      expect(router.navigate).toHaveBeenCalledWith(["/tabs/vault"]);
+    });
+
+    it("should use legacy registration when feature flag is disabled", async () => {
+      // Arrange
+      configService.getFeatureFlag.mockResolvedValue(false);
+
+      const mockPublicKey = "mock-public-key";
+      const mockPrivateKey = {
+        encryptedString: "mock-encrypted-private-key",
+      } as any;
+
+      keyService.initAccount.mockResolvedValue({
+        publicKey: mockPublicKey,
+        privateKey: mockPrivateKey,
+      } as any);
+
+      apiService.postAccountKeys.mockResolvedValue(undefined);
+      passwordResetEnrollmentService.enroll.mockResolvedValue(undefined);
+      deviceTrustService.trustDevice.mockResolvedValue(undefined);
+      loginDecryptionOptionsService.handleCreateUserSuccess.mockResolvedValue(undefined);
+      router.navigate.mockResolvedValue(true);
+
+      // Act
+      await component["createUser"]();
+
+      // Assert
+      expect(configService.getFeatureFlag).toHaveBeenCalledWith(
+        FeatureFlag.PM27279_V2RegistrationTdeJit,
+      );
+      expect(keyService.initAccount).toHaveBeenCalledWith(mockUserId);
+      expect(apiService.postAccountKeys).toHaveBeenCalledWith(
+        expect.objectContaining({
+          publicKey: mockPublicKey,
+          encryptedPrivateKey: mockPrivateKey.encryptedString,
+        }),
+      );
+      expect(passwordResetEnrollmentService.enroll).toHaveBeenCalledWith(mockOrgId);
+      expect(deviceTrustService.trustDevice).toHaveBeenCalledWith(mockUserId);
+
+      // Verify success toast
+      expect(toastService.showToast).toHaveBeenCalledWith({
+        variant: "success",
+        title: null,
+        message: "accountSuccessfullyCreated",
+      });
+
+      // Verify navigation
+      expect(loginDecryptionOptionsService.handleCreateUserSuccess).toHaveBeenCalled();
+      expect(router.navigate).toHaveBeenCalledWith(["/tabs/vault"]);
+    });
+  });
+});
--- a/libs/auth/src/angular/login-decryption-options/login-decryption-options.component.ts
+++ b/libs/auth/src/angular/login-decryption-options/login-decryption-options.component.ts
@@ -5,7 +5,17 @@ import { Component, DestroyRef, OnInit } from "@angular/core";
 import { takeUntilDestroyed } from "@angular/core/rxjs-interop";
 import { FormBuilder, FormControl, ReactiveFormsModule } from "@angular/forms";
 import { Router } from "@angular/router";
-import { catchError, defer, firstValueFrom, from, map, of, switchMap, throwError } from "rxjs";
+import {
+  catchError,
+  concatMap,
+  defer,
+  firstValueFrom,
+  from,
+  map,
+  of,
+  switchMap,
+  throwError,
+} from "rxjs";

 import { JslibModule } from "@bitwarden/angular/jslib.module";
 import {
@@ -20,13 +30,27 @@ import { AccountService } from "@bitwarden/common/auth/abstractions/account.serv
 import { PasswordResetEnrollmentServiceAbstraction } from "@bitwarden/common/auth/abstractions/password-reset-enrollment.service.abstraction";
 import { SsoLoginServiceAbstraction } from "@bitwarden/common/auth/abstractions/sso-login.service.abstraction";
 import { ClientType } from "@bitwarden/common/enums";
+import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
+import { AccountCryptographicStateService } from "@bitwarden/common/key-management/account-cryptography/account-cryptographic-state.service";
 import { DeviceTrustServiceAbstraction } from "@bitwarden/common/key-management/device-trust/abstractions/device-trust.service.abstraction";
+import { SecurityStateService } from "@bitwarden/common/key-management/security-state/abstractions/security-state.service";
+import {
+  SignedPublicKey,
+  SignedSecurityState,
+  WrappedSigningKey,
+} from "@bitwarden/common/key-management/types";
 import { KeysRequest } from "@bitwarden/common/models/request/keys.request";
+import { AppIdService } from "@bitwarden/common/platform/abstractions/app-id.service";
+import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
 import { I18nService } from "@bitwarden/common/platform/abstractions/i18n.service";
 import { MessagingService } from "@bitwarden/common/platform/abstractions/messaging.service";
 import { PlatformUtilsService } from "@bitwarden/common/platform/abstractions/platform-utils.service";
+import { RegisterSdkService } from "@bitwarden/common/platform/abstractions/sdk/register-sdk.service";
+import { asUuid } from "@bitwarden/common/platform/abstractions/sdk/sdk.service";
 import { ValidationService } from "@bitwarden/common/platform/abstractions/validation.service";
+import { SymmetricCryptoKey } from "@bitwarden/common/platform/models/domain/symmetric-crypto-key";
 import { UserId } from "@bitwarden/common/types/guid";
+import { DeviceKey, UserKey } from "@bitwarden/common/types/key";
 // This import has been flagged as unallowed for this class. It may be involved in a circular dependency loop.
 // eslint-disable-next-line no-restricted-imports
 import {
@@ -40,6 +64,7 @@ import {
  TypographyModule,
 } from "@bitwarden/components";
 import { KeyService } from "@bitwarden/key-management";
+import { OrganizationId as SdkOrganizationId, UserId as SdkUserId } from "@bitwarden/sdk-internal";

 import { LoginDecryptionOptionsService } from "./login-decryption-options.service";

@@ -112,6 +137,11 @@ export class LoginDecryptionOptionsComponent implements OnInit {
    private userDecryptionOptionsService: UserDecryptionOptionsServiceAbstraction,
    private validationService: ValidationService,
    private logoutService: LogoutService,
+    private registerSdkService: RegisterSdkService,
+    private securityStateService: SecurityStateService,
+    private appIdService: AppIdService,
+    private configService: ConfigService,
+    private accountCryptographicStateService: AccountCryptographicStateService,
  ) {
    this.clientType = this.platformUtilsService.getClientType();
  }
@@ -135,7 +165,7 @@ export class LoginDecryptionOptionsComponent implements OnInit {

    try {
      const userDecryptionOptions = await firstValueFrom(
-        this.userDecryptionOptionsService.userDecryptionOptions$,
+        this.userDecryptionOptionsService.userDecryptionOptionsById$(this.activeAccountId),
      );

      if (
@@ -251,9 +281,85 @@ export class LoginDecryptionOptionsComponent implements OnInit {
    }

    try {
-      const { publicKey, privateKey } = await this.keyService.initAccount(this.activeAccountId);
-      const keysRequest = new KeysRequest(publicKey, privateKey.encryptedString);
-      await this.apiService.postAccountKeys(keysRequest);
+      const useSdkV2Creation = await this.configService.getFeatureFlag(
+        FeatureFlag.PM27279_V2RegistrationTdeJit,
+      );
+      if (useSdkV2Creation) {
+        const deviceIdentifier = await this.appIdService.getAppId();
+        const userId = this.activeAccountId;
+        const organizationId = this.newUserOrgId;
+
+        const orgKeyResponse = await this.organizationApiService.getKeys(organizationId);
+        const register_result = await firstValueFrom(
+          this.registerSdkService.registerClient$(userId).pipe(
+            concatMap(async (sdk) => {
+              if (!sdk) {
+                throw new Error("SDK not available");
+              }
+
+              using ref = sdk.take();
+              return await ref.value
+                .auth()
+                .registration()
+                .post_keys_for_tde_registration({
+                  org_id: asUuid<SdkOrganizationId>(organizationId),
+                  org_public_key: orgKeyResponse.publicKey,
+                  user_id: asUuid<SdkUserId>(userId),
+                  device_identifier: deviceIdentifier,
+                  trust_device: this.formGroup.value.rememberDevice,
+                });
+            }),
+          ),
+        );
+        // The keys returned here can only be v2 keys, since the SDK only implements returning V2 keys.
+        if ("V1" in register_result.account_cryptographic_state) {
+          throw new Error("Unexpected V1 account cryptographic state");
+        }
+
+        // Note: When SDK state management matures, these should be moved into post_keys_for_tde_registration
+        // Set account cryptography state
+        await this.accountCryptographicStateService.setAccountCryptographicState(
+          register_result.account_cryptographic_state,
+          userId,
+        );
+        // Legacy individual states
+        await this.keyService.setPrivateKey(
+          register_result.account_cryptographic_state.V2.private_key,
+          userId,
+        );
+        await this.keyService.setSignedPublicKey(
+          register_result.account_cryptographic_state.V2.signed_public_key as SignedPublicKey,
+          userId,
+        );
+        await this.keyService.setUserSigningKey(
+          register_result.account_cryptographic_state.V2.signing_key as WrappedSigningKey,
+          userId,
+        );
+        await this.securityStateService.setAccountSecurityState(
+          register_result.account_cryptographic_state.V2.security_state as SignedSecurityState,
+          userId,
+        );
+
+        // TDE unlock
+        await this.deviceTrustService.setDeviceKey(
+          userId,
+          SymmetricCryptoKey.fromString(register_result.device_key) as DeviceKey,
+        );
+
+        // Set user key - user is now unlocked
+        await this.keyService.setUserKey(
+          SymmetricCryptoKey.fromString(register_result.user_key) as UserKey,
+          userId,
+        );
+      } else {
+        const { publicKey, privateKey } = await this.keyService.initAccount(this.activeAccountId);
+        const keysRequest = new KeysRequest(publicKey, privateKey.encryptedString);
+        await this.apiService.postAccountKeys(keysRequest);
+        await this.passwordResetEnrollmentService.enroll(this.newUserOrgId);
+        if (this.formGroup.value.rememberDevice) {
+          await this.deviceTrustService.trustDevice(this.activeAccountId);
+        }
+      }

      this.toastService.showToast({
        variant: "success",
@@ -261,12 +367,6 @@ export class LoginDecryptionOptionsComponent implements OnInit {
        message: this.i18nService.t("accountSuccessfullyCreated"),
      });

-      await this.passwordResetEnrollmentService.enroll(this.newUserOrgId);
-
-      if (this.formGroup.value.rememberDevice) {
-        await this.deviceTrustService.trustDevice(this.activeAccountId);
-      }
-
      await this.loginDecryptionOptionsService.handleCreateUserSuccess();

      if (this.clientType === ClientType.Desktop) {
--- a/libs/auth/src/angular/login-via-auth-request/login-via-auth-request.component.ts
+++ b/libs/auth/src/angular/login-via-auth-request/login-via-auth-request.component.ts
@@ -822,7 +822,7 @@ export class LoginViaAuthRequestComponent implements OnInit, OnDestroy {
  }

  private async handleSuccessfulLoginNavigation(userId: UserId) {
-    await this.loginSuccessHandlerService.run(userId);
+    await this.loginSuccessHandlerService.run(userId, null);
    await this.router.navigate(["vault"]);
  }
 }
--- a/libs/auth/src/angular/login/default-login-component.service.ts
+++ b/libs/auth/src/angular/login/default-login-component.service.ts
@@ -33,19 +33,27 @@ export class DefaultLoginComponentService implements LoginComponentService {
   */
  async redirectToSsoLogin(email: string): Promise<void | null> {
    // Set the state that we'll need to verify the SSO login when we get the code back
-    const [state, codeChallenge] = await this.setSsoPreLoginState();
-
-    // Set the email address in state. This is used in 2 places:
-    // 1. On the web client, on the SSO component we need the email address to look up
-    //    the org SSO identifier. The email address is passed via query param for the other clients.
-    // 2. On all clients, after authentication on the originating client the SSO component
-    //    will need to look up 2FA Remember token by email.
-    await this.ssoLoginService.setSsoEmail(email);
+    const [state, codeChallenge] = await this.setSsoPreLoginState(email);

    // Finally, we redirect to the SSO login page. This will be handled by each client implementation of this service.
    await this.redirectToSso(email, state, codeChallenge);
  }

+  /**
+   * Redirects the user to the SSO login page, either via route or in a new browser window.
+   * @param email The email address of the user attempting to log in
+   */
+  async redirectToSsoLoginWithOrganizationSsoIdentifier(
+    email: string,
+    orgSsoIdentifier: string,
+  ): Promise<void | null> {
+    // Set the state that we'll need to verify the SSO login when we get the code back
+    const [state, codeChallenge] = await this.setSsoPreLoginState(email);
+
+    // Finally, we redirect to the SSO login page. This will be handled by each client implementation of this service.
+    await this.redirectToSso(email, state, codeChallenge, orgSsoIdentifier);
+  }
+
  /**
   * No-op implementation of redirectToSso
   */
@@ -53,6 +61,7 @@ export class DefaultLoginComponentService implements LoginComponentService {
    email: string,
    state: string,
    codeChallenge: string,
+    orgSsoIdentifier?: string,
  ): Promise<void> {
    return;
  }
@@ -65,9 +74,9 @@ export class DefaultLoginComponentService implements LoginComponentService {
  }

  /**
-   * Sets the state required for verifying SSO login after completion
+   * Set the state that we'll need to verify the SSO login when we get the authorization code back
   */
-  private async setSsoPreLoginState(): Promise<[string, string]> {
+  private async setSsoPreLoginState(email: string): Promise<[string, string]> {
    // Generate SSO params
    const passwordOptions: any = {
      type: "password",
@@ -93,6 +102,13 @@ export class DefaultLoginComponentService implements LoginComponentService {
    await this.ssoLoginService.setSsoState(state);
    await this.ssoLoginService.setCodeVerifier(codeVerifier);

+    // Set the email address in state. This is used in 2 places:
+    // 1. On the web client, on the SSO component we need the email address to look up
+    //    the org SSO identifier. The email address is passed via query param for the other clients.
+    // 2. On all clients, after authentication on the originating client the SSO component
+    //    will need to look up 2FA Remember token by email.
+    await this.ssoLoginService.setSsoEmail(email);
+
    return [state, codeChallenge];
  }
 }
--- a/libs/auth/src/angular/login/login-component.service.ts
+++ b/libs/auth/src/angular/login/login-component.service.ts
@@ -35,6 +35,14 @@ export abstract class LoginComponentService {
   */
  redirectToSsoLogin: (email: string) => Promise<void | null>;

+  /**
+   * Redirects the user to the SSO login page with organization SSO identifier, either via route or in a new browser window.
+   */
+  redirectToSsoLoginWithOrganizationSsoIdentifier: (
+    email: string,
+    orgSsoIdentifier: string | null | undefined,
+  ) => Promise<void | null>;
+
  /**
   * Shows the back button.
   */
--- a/libs/auth/src/angular/login/login.component.ts
+++ b/libs/auth/src/angular/login/login.component.ts
@@ -381,8 +381,26 @@ export class LoginComponent implements OnInit, OnDestroy {
      return;
    }

+    // redirect to SSO if ssoOrganizationIdentifier is present in token response
+    if (authResult.requiresSso) {
+      const email = this.formGroup?.value?.email;
+      if (!email) {
+        this.toastService.showToast({
+          variant: "error",
+          title: this.i18nService.t("errorOccurred"),
+          message: this.i18nService.t("emailRequiredForSsoLogin"),
+        });
+        return;
+      }
+      await this.loginComponentService.redirectToSsoLoginWithOrganizationSsoIdentifier(
+        email,
+        authResult.ssoOrganizationIdentifier,
+      );
+      return;
+    }
+
    // User logged in successfully so execute side effects
-    await this.loginSuccessHandlerService.run(authResult.userId);
+    await this.loginSuccessHandlerService.run(authResult.userId, authResult.masterPassword);

    // Determine where to send the user next
    // The AuthGuard will handle routing to change-password based on state
--- a/libs/auth/src/angular/new-device-verification/default-new-device-verification-component.service.ts
+++ b/libs/auth/src/angular/new-device-verification/default-new-device-verification-component.service.ts
@@ -1,8 +1,6 @@
 import { NewDeviceVerificationComponentService } from "./new-device-verification-component.service";

-export class DefaultNewDeviceVerificationComponentService
-  implements NewDeviceVerificationComponentService
-{
+export class DefaultNewDeviceVerificationComponentService implements NewDeviceVerificationComponentService {
  showBackButton() {
    return true;
  }
--- a/libs/auth/src/angular/new-device-verification/new-device-verification.component.ts
+++ b/libs/auth/src/angular/new-device-verification/new-device-verification.component.ts
@@ -152,9 +152,7 @@ export class NewDeviceVerificationComponent implements OnInit, OnDestroy {
        return;
      }

-      // FIXME: Verify that this floating promise is intentional. If it is, add an explanatory comment and ensure there is proper error handling.
-      // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      this.loginSuccessHandlerService.run(authResult.userId);
+      await this.loginSuccessHandlerService.run(authResult.userId, authResult.masterPassword);

      // TODO: PM-22663 use the new service to handle routing.
      const activeUserId = await firstValueFrom(this.accountService.activeAccount$.pipe(getUserId));
--- a/libs/auth/src/angular/registration/registration-finish/registration-finish.component.ts
+++ b/libs/auth/src/angular/registration/registration-finish/registration-finish.component.ts
@@ -206,7 +206,10 @@ export class RegistrationFinishComponent implements OnInit, OnDestroy {
        return;
      }

-      await this.loginSuccessHandlerService.run(authenticationResult.userId);
+      await this.loginSuccessHandlerService.run(
+        authenticationResult.userId,
+        authenticationResult.masterPassword ?? null,
+      );

      if (this.premiumInterest) {
        await this.premiumInterestStateService.setPremiumInterest(
--- a/libs/auth/src/angular/registration/registration-start/registration-start.mdx
+++ b/libs/auth/src/angular/registration/registration-start/registration-start.mdx
@@ -1,4 +1,4 @@
-import { Meta, Story, Controls } from "@storybook/addon-docs";
+import { Meta, Story, Controls } from "@storybook/addon-docs/blocks";

 import * as stories from "./registration-start.stories";

--- a/libs/auth/src/angular/sso/sso.component.ts
+++ b/libs/auth/src/angular/sso/sso.component.ts
@@ -437,7 +437,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.loginSuccessHandlerService.run(authResult.userId);
+      await this.loginSuccessHandlerService.run(authResult.userId, null);

      // Save off the OrgSsoIdentifier for use in the TDE flows (or elsewhere)
      // - TDE login decryption options component
@@ -460,7 +460,7 @@ export class SsoComponent implements OnInit {

      // must come after 2fa check since user decryption options aren't available if 2fa is required
      const userDecryptionOpts = await firstValueFrom(
-        this.userDecryptionOptionsService.userDecryptionOptions$,
+        this.userDecryptionOptionsService.userDecryptionOptionsById$(authResult.userId),
      );

      const tdeEnabled = userDecryptionOpts.trustedDeviceOption
@@ -478,7 +478,7 @@ export class SsoComponent implements OnInit {
        !userDecryptionOpts.hasMasterPassword &&
        userDecryptionOpts.keyConnectorOption === undefined;

-      if (requireSetPassword || authResult.resetMasterPassword) {
+      if (requireSetPassword) {
        // Change implies going no password -> password in this case
        return await this.handleChangePasswordRequired(orgSsoIdentifier);
      }
--- a/libs/auth/src/angular/two-factor-auth/child-components/two-factor-auth-webauthn/default-two-factor-auth-webauthn-component.service.ts
+++ b/libs/auth/src/angular/two-factor-auth/child-components/two-factor-auth-webauthn/default-two-factor-auth-webauthn-component.service.ts
@@ -1,8 +1,6 @@
 import { TwoFactorAuthWebAuthnComponentService } from "./two-factor-auth-webauthn-component.service";

-export class DefaultTwoFactorAuthWebAuthnComponentService
-  implements TwoFactorAuthWebAuthnComponentService
-{
+export class DefaultTwoFactorAuthWebAuthnComponentService implements TwoFactorAuthWebAuthnComponentService {
  /**
   * Default implementation is to not open in a new tab.
   */
--- a/libs/auth/src/angular/two-factor-auth/two-factor-auth.component.spec.ts
+++ b/libs/auth/src/angular/two-factor-auth/two-factor-auth.component.spec.ts
@@ -176,7 +176,9 @@ describe("TwoFactorAuthComponent", () => {
    selectedUserDecryptionOptions = new BehaviorSubject<UserDecryptionOptions>(
      mockUserDecryptionOpts.withMasterPassword,
    );
-    mockUserDecryptionOptionsService.userDecryptionOptions$ = selectedUserDecryptionOptions;
+    mockUserDecryptionOptionsService.userDecryptionOptionsById$.mockReturnValue(
+      selectedUserDecryptionOptions,
+    );

    TestBed.configureTestingModule({
      declarations: [TestTwoFactorComponent],
@@ -419,6 +421,7 @@ describe("TwoFactorAuthComponent", () => {
            keyConnectorUrl:
              mockUserDecryptionOpts.noMasterPasswordWithKeyConnector.keyConnectorOption!
                .keyConnectorUrl,
+            organizationSsoIdentifier: "test-sso-id",
          }),
        );
        const authResult = new AuthResult();
--- a/libs/auth/src/angular/two-factor-auth/two-factor-auth.component.ts
+++ b/libs/auth/src/angular/two-factor-auth/two-factor-auth.component.ts
@@ -450,7 +450,7 @@ export class TwoFactorAuthComponent implements OnInit, OnDestroy {
    }

    // User is fully logged in so handle any post login logic before executing navigation
-    await this.loginSuccessHandlerService.run(authResult.userId);
+    await this.loginSuccessHandlerService.run(authResult.userId, authResult.masterPassword);

    // Save off the OrgSsoIdentifier for use in the TDE flows
    // - TDE login decryption options component
@@ -473,7 +473,7 @@ export class TwoFactorAuthComponent implements OnInit, OnDestroy {
    }

    const userDecryptionOpts = await firstValueFrom(
-      this.userDecryptionOptionsService.userDecryptionOptions$,
+      this.userDecryptionOptionsService.userDecryptionOptionsById$(authResult.userId),
    );

    const tdeEnabled = await this.isTrustedDeviceEncEnabled(userDecryptionOpts.trustedDeviceOption);
@@ -487,7 +487,7 @@ export class TwoFactorAuthComponent implements OnInit, OnDestroy {
      !userDecryptionOpts.hasMasterPassword && userDecryptionOpts.keyConnectorOption === undefined;

    // New users without a master password must set a master password before advancing.
-    if (requireSetPassword || authResult.resetMasterPassword) {
+    if (requireSetPassword) {
      // Change implies going no password -> password in this case
      return await this.handleChangePasswordRequired(this.orgSsoIdentifier);
    }