[PM-27315] Add account cryptographic state service (#17589)

* Update account init and save signed public key

* Add account cryptographic state service

* Fix build

* Cleanup

* Fix build

* Fix import

* Fix build on browser

* Fix

* Fix DI

* Fix

* Fix

* Fix

* Fix

* Fix

* Fix test

* Fix desktop build

* Fix

* Address nits

* Cleanup setting private key

* Add tests

* Add tests

* Add test coverage

* Relative imports

* Fix web build

* Cleanup setting of private key

libs/common/src/auth/models/response/identity-token.response.spec.ts

@@ -116,4 +116,36 @@ describe("IdentityTokenResponse", () => {
     const identityTokenResponse = new IdentityTokenResponse(response);
     expect(identityTokenResponse.userDecryptionOptions).toBeDefined();
   });
+
+  it("should create response with accountKeys not present", () => {
+    const response = {
+      access_token: accessToken,
+      token_type: tokenType,
+      AccountKeys: null as unknown,
+    };
+
+    const identityTokenResponse = new IdentityTokenResponse(response);
+    expect(identityTokenResponse.accountKeysResponseModel).toBeNull();
+  });
+
+  it("should create response with accountKeys present", () => {
+    const accountKeysData = {
+      publicKeyEncryptionKeyPair: {
+        publicKey: "testPublicKey",
+        wrappedPrivateKey: "testPrivateKey",
+      },
+    };
+
+    const response = {
+      access_token: accessToken,
+      token_type: tokenType,
+      AccountKeys: accountKeysData,
+    };
+
+    const identityTokenResponse = new IdentityTokenResponse(response);
+    expect(identityTokenResponse.accountKeysResponseModel).toBeDefined();
+    expect(
+      identityTokenResponse.accountKeysResponseModel?.publicKeyEncryptionKeyPair,
+    ).toBeDefined();
+  });
 });

libs/common/src/auth/models/response/identity-token.response.ts

@@ -5,6 +5,7 @@
 import { Argon2KdfConfig, KdfConfig, KdfType, PBKDF2KdfConfig } from "@bitwarden/key-management";
 
 import { EncString } from "../../../key-management/crypto/models/enc-string";
+import { PrivateKeysResponseModel } from "../../../key-management/keys/response/private-keys.response";
 import { BaseResponse } from "../../../models/response/base.response";
 
 import { MasterPasswordPolicyResponse } from "./master-password-policy.response";
@@ -19,6 +20,7 @@ export class IdentityTokenResponse extends BaseResponse {
 
   // Decryption Information
   privateKey: string; // userKeyEncryptedPrivateKey
+  accountKeysResponseModel: PrivateKeysResponseModel | null = null;
   key?: EncString; // masterKeyEncryptedUserKey
   twoFactorToken: string;
   kdfConfig: KdfConfig;
@@ -52,6 +54,11 @@ export class IdentityTokenResponse extends BaseResponse {
     }
 
     this.privateKey = this.getResponseProperty("PrivateKey");
+    if (this.getResponseProperty("AccountKeys") != null) {
+      this.accountKeysResponseModel = new PrivateKeysResponseModel(
+        this.getResponseProperty("AccountKeys"),
+      );
+    }
     const key = this.getResponseProperty("Key");
     if (key) {
       this.key = new EncString(key);

libs/common/src/key-management/account-cryptography/account-cryptographic-state.service.ts

@@ -0,0 +1,22 @@
+import { Observable } from "rxjs";
+
+import { WrappedAccountCryptographicState } from "@bitwarden/sdk-internal";
+import { UserId } from "@bitwarden/user-core";
+
+export abstract class AccountCryptographicStateService {
+  /**
+   * Emits the provided user's account cryptographic state or null if there is no account cryptographic state present for the user.
+   */
+  abstract accountCryptographicState$(
+    userId: UserId,
+  ): Observable<WrappedAccountCryptographicState | null>;
+
+  /**
+   * Sets the account cryptographic state.
+   * This is not yet validated, and is only validated upon SDK initialization.
+   */
+  abstract setAccountCryptographicState(
+    accountCryptographicState: WrappedAccountCryptographicState,
+    userId: UserId,
+  ): Promise<void>;
+}

libs/common/src/key-management/account-cryptography/default-account-cryptographic-state.service.spec.ts

@@ -0,0 +1,133 @@
+import { firstValueFrom } from "rxjs";
+
+import { WrappedAccountCryptographicState } from "@bitwarden/sdk-internal";
+import { FakeStateProvider } from "@bitwarden/state-test-utils";
+import { UserId } from "@bitwarden/user-core";
+
+import { FakeAccountService, mockAccountServiceWith } from "../../../spec";
+
+import {
+  ACCOUNT_CRYPTOGRAPHIC_STATE,
+  DefaultAccountCryptographicStateService,
+} from "./default-account-cryptographic-state.service";
+
+describe("DefaultAccountCryptographicStateService", () => {
+  let service: DefaultAccountCryptographicStateService;
+  let stateProvider: FakeStateProvider;
+  let accountService: FakeAccountService;
+
+  const mockUserId = "user-id" as UserId;
+
+  beforeEach(() => {
+    accountService = mockAccountServiceWith(mockUserId);
+    stateProvider = new FakeStateProvider(accountService);
+    service = new DefaultAccountCryptographicStateService(stateProvider);
+  });
+
+  describe("accountCryptographicState$", () => {
+    it("returns null when no state is set", async () => {
+      const result = await firstValueFrom(service.accountCryptographicState$(mockUserId));
+      expect(result).toBeNull();
+    });
+
+    it("returns the account cryptographic state when set (V1)", async () => {
+      const mockState: WrappedAccountCryptographicState = {
+        V1: {
+          private_key: "test-wrapped-state" as any,
+        },
+      };
+      await stateProvider.setUserState(ACCOUNT_CRYPTOGRAPHIC_STATE, mockState, mockUserId);
+      const result = await firstValueFrom(service.accountCryptographicState$(mockUserId));
+      expect(result).toEqual(mockState);
+    });
+
+    it("returns the account cryptographic state when set (V2)", async () => {
+      const mockState: WrappedAccountCryptographicState = {
+        V2: {
+          private_key: "test-wrapped-private-key" as any,
+          signing_key: "test-wrapped-signing-key" as any,
+          signed_public_key: "test-signed-public-key" as any,
+          security_state: "test-security-state",
+        },
+      };
+      await stateProvider.setUserState(ACCOUNT_CRYPTOGRAPHIC_STATE, mockState, mockUserId);
+      const result = await firstValueFrom(service.accountCryptographicState$(mockUserId));
+      expect(result).toEqual(mockState);
+    });
+
+    it("emits updated state when state changes", async () => {
+      const mockState1: any = {
+        V1: {
+          private_key: "test-state-1" as any,
+        },
+      };
+      const mockState2: any = {
+        V1: {
+          private_key: "test-state-2" as any,
+        },
+      };
+
+      await stateProvider.setUserState(ACCOUNT_CRYPTOGRAPHIC_STATE, mockState1, mockUserId);
+
+      const observable = service.accountCryptographicState$(mockUserId);
+      const results: (WrappedAccountCryptographicState | null)[] = [];
+      const subscription = observable.subscribe((state) => results.push(state));
+
+      await stateProvider.setUserState(ACCOUNT_CRYPTOGRAPHIC_STATE, mockState2, mockUserId);
+
+      subscription.unsubscribe();
+
+      expect(results).toHaveLength(2);
+      expect(results[0]).toEqual(mockState1);
+      expect(results[1]).toEqual(mockState2);
+    });
+  });
+
+  describe("setAccountCryptographicState", () => {
+    it("sets the account cryptographic state", async () => {
+      const mockState: WrappedAccountCryptographicState = {
+        V1: {
+          private_key: "test-wrapped-state" as any,
+        },
+      };
+
+      await service.setAccountCryptographicState(mockState, mockUserId);
+
+      const result = await firstValueFrom(service.accountCryptographicState$(mockUserId));
+
+      expect(result).toEqual(mockState);
+    });
+
+    it("overwrites existing state", async () => {
+      const mockState1: WrappedAccountCryptographicState = {
+        V1: {
+          private_key: "test-state-1" as any,
+        },
+      };
+      const mockState2: WrappedAccountCryptographicState = {
+        V1: {
+          private_key: "test-state-2" as any,
+        },
+      };
+
+      await service.setAccountCryptographicState(mockState1, mockUserId);
+      await service.setAccountCryptographicState(mockState2, mockUserId);
+
+      const result = await firstValueFrom(service.accountCryptographicState$(mockUserId));
+
+      expect(result).toEqual(mockState2);
+    });
+  });
+
+  describe("ACCOUNT_CRYPTOGRAPHIC_STATE key definition", () => {
+    it("deserializer returns object as-is", () => {
+      const mockState: any = {
+        V1: {
+          private_key: "test" as any,
+        },
+      };
+      const result = ACCOUNT_CRYPTOGRAPHIC_STATE.deserializer(mockState);
+      expect(result).toBe(mockState);
+    });
+  });
+});

libs/common/src/key-management/account-cryptography/default-account-cryptographic-state.service.ts

@@ -0,0 +1,35 @@
+import { Observable } from "rxjs";
+
+import { WrappedAccountCryptographicState } from "@bitwarden/sdk-internal";
+import { CRYPTO_DISK, StateProvider, UserKeyDefinition } from "@bitwarden/state";
+import { UserId } from "@bitwarden/user-core";
+
+import { AccountCryptographicStateService } from "./account-cryptographic-state.service";
+
+export const ACCOUNT_CRYPTOGRAPHIC_STATE = new UserKeyDefinition<WrappedAccountCryptographicState>(
+  CRYPTO_DISK,
+  "accountCryptographicState",
+  {
+    deserializer: (obj) => obj as WrappedAccountCryptographicState,
+    clearOn: ["logout"],
+  },
+);
+
+export class DefaultAccountCryptographicStateService implements AccountCryptographicStateService {
+  constructor(protected stateProvider: StateProvider) {}
+
+  accountCryptographicState$(userId: UserId): Observable<WrappedAccountCryptographicState | null> {
+    return this.stateProvider.getUserState$(ACCOUNT_CRYPTOGRAPHIC_STATE, userId);
+  }
+
+  async setAccountCryptographicState(
+    accountCryptographicState: WrappedAccountCryptographicState,
+    userId: UserId,
+  ): Promise<void> {
+    await this.stateProvider.setUserState(
+      ACCOUNT_CRYPTOGRAPHIC_STATE,
+      accountCryptographicState,
+      userId,
+    );
+  }
+}

libs/common/src/key-management/keys/response/private-keys.response.ts

@@ -1,3 +1,5 @@
+import { SignedPublicKey, WrappedAccountCryptographicState } from "@bitwarden/sdk-internal";
+
 import { SecurityStateResponse } from "../../security-state/response/security-state.response";
 
 import { PublicKeyEncryptionKeyPairResponse } from "./public-key-encryption-key-pair.response";
@@ -52,4 +54,31 @@ export class PrivateKeysResponseModel {
       );
     }
   }
+
+  toWrappedAccountCryptographicState(): WrappedAccountCryptographicState {
+    if (this.signatureKeyPair === null && this.securityState === null) {
+      // V1 user
+      return {
+        V1: {
+          private_key: this.publicKeyEncryptionKeyPair.wrappedPrivateKey,
+        },
+      };
+    } else if (this.signatureKeyPair !== null && this.securityState !== null) {
+      // V2 user
+      return {
+        V2: {
+          private_key: this.publicKeyEncryptionKeyPair.wrappedPrivateKey,
+          signing_key: this.signatureKeyPair.wrappedSigningKey,
+          signed_public_key: this.publicKeyEncryptionKeyPair.signedPublicKey as SignedPublicKey,
+          security_state: this.securityState.securityState as string,
+        },
+      };
+    } else {
+      throw new Error("Both signatureKeyPair and securityState must be present or absent together");
+    }
+  }
+
+  isV2Encryption(): boolean {
+    return this.signatureKeyPair !== null && this.securityState !== null;
+  }
 }

libs/common/src/platform/services/key-state/user-key.state.spec.ts

@@ -1,4 +1,4 @@
-import { EncString } from "../../../key-management/crypto/models/enc-string";
+import { EncryptedString, EncString } from "../../../key-management/crypto/models/enc-string";
 import { EncryptionType } from "../../enums";
 import { Utils } from "../../misc/utils";
 
@@ -27,7 +27,9 @@ describe("Encrypted private key", () => {
   it("should deserialize encrypted private key", () => {
     const encryptedPrivateKey = makeEncString().encryptedString;
 
-    const result = sut.deserializer(JSON.parse(JSON.stringify(encryptedPrivateKey)));
+    const result = sut.deserializer(
+      JSON.parse(JSON.stringify(encryptedPrivateKey as unknown)) as unknown as EncryptedString,
+    );
 
     expect(result).toEqual(encryptedPrivateKey);
   });

libs/common/src/platform/sync/default-sync.service.spec.ts

@@ -11,8 +11,6 @@ import {
   UserDecryptionOptions,
   UserDecryptionOptionsServiceAbstraction,
 } from "@bitwarden/auth/common";
-import { EncString } from "@bitwarden/common/key-management/crypto/models/enc-string";
-import { SecurityStateService } from "@bitwarden/common/key-management/security-state/abstractions/security-state.service";
 // 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 { KdfConfigService, KeyService, PBKDF2KdfConfig } from "@bitwarden/key-management";
@@ -29,6 +27,8 @@ import { TokenService } from "../../auth/abstractions/token.service";
 import { AuthenticationStatus } from "../../auth/enums/authentication-status";
 import { DomainSettingsService } from "../../autofill/services/domain-settings.service";
 import { BillingAccountProfileStateService } from "../../billing/abstractions";
+import { AccountCryptographicStateService } from "../../key-management/account-cryptography/account-cryptographic-state.service";
+import { EncString } from "../../key-management/crypto/models/enc-string";
 import { KeyConnectorService } from "../../key-management/key-connector/abstractions/key-connector.service";
 import { InternalMasterPasswordServiceAbstraction } from "../../key-management/master-password/abstractions/master-password.service.abstraction";
 import {
@@ -36,6 +36,7 @@ import {
   MasterPasswordSalt,
   MasterPasswordUnlockData,
 } from "../../key-management/master-password/types/master-password.types";
+import { SecurityStateService } from "../../key-management/security-state/abstractions/security-state.service";
 import { SendApiService } from "../../tools/send/services/send-api.service.abstraction";
 import { InternalSendService } from "../../tools/send/services/send.service.abstraction";
 import { UserId } from "../../types/guid";
@@ -76,6 +77,7 @@ describe("DefaultSyncService", () => {
   let stateProvider: MockProxy<StateProvider>;
   let securityStateService: MockProxy<SecurityStateService>;
   let kdfConfigService: MockProxy<KdfConfigService>;
+  let accountCryptographicStateService: MockProxy<AccountCryptographicStateService>;
 
   let sut: DefaultSyncService;
 
@@ -107,6 +109,7 @@ describe("DefaultSyncService", () => {
     stateProvider = mock();
     securityStateService = mock();
     kdfConfigService = mock();
+    accountCryptographicStateService = mock();
 
     sut = new DefaultSyncService(
       masterPasswordAbstraction,
@@ -135,6 +138,7 @@ describe("DefaultSyncService", () => {
       stateProvider,
       securityStateService,
       kdfConfigService,
+      accountCryptographicStateService,
     );
   });
 

libs/common/src/platform/sync/default-sync.service.ts

@@ -9,8 +9,9 @@ import {
   CollectionDetailsResponse,
   CollectionService,
 } from "@bitwarden/admin-console/common";
-// This import has been flagged as unallowed for this class. It may be involved in a circular dependency loop.
+import { AccountCryptographicStateService } from "@bitwarden/common/key-management/account-cryptography/account-cryptographic-state.service";
 import { SecurityStateService } from "@bitwarden/common/key-management/security-state/abstractions/security-state.service";
+// 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 { KdfConfigService, KeyService } from "@bitwarden/key-management";
 
@@ -101,6 +102,7 @@ export class DefaultSyncService extends CoreSyncService {
     stateProvider: StateProvider,
     private securityStateService: SecurityStateService,
     private kdfConfigService: KdfConfigService,
+    private accountCryptographicStateService: AccountCryptographicStateService,
   ) {
     super(
       tokenService,
@@ -239,12 +241,18 @@ export class DefaultSyncService extends CoreSyncService {
 
     // Cleanup: Only the first branch should be kept after the server always returns accountKeys https://bitwarden.atlassian.net/browse/PM-21768
     if (response.accountKeys != null) {
+      await this.accountCryptographicStateService.setAccountCryptographicState(
+        response.accountKeys.toWrappedAccountCryptographicState(),
+        response.id,
+      );
+
+      // V1 and V2 users
       await this.keyService.setPrivateKey(
         response.accountKeys.publicKeyEncryptionKeyPair.wrappedPrivateKey,
         response.id,
       );
-      if (response.accountKeys.signatureKeyPair !== null) {
-        // User is V2 user
+      // V2 users only
+      if (response.accountKeys.isV2Encryption()) {
         await this.keyService.setUserSigningKey(
           response.accountKeys.signatureKeyPair.wrappedSigningKey,
           response.id,