[PM-28901] Fix master key not being set to state after kdf update (#17990)

* Fix master key not being set to state after kdf update

* Fix cli build

* Fix test error

* Fix hash purpose

* Add test for master key being set

* Fix incorrect variable name

apps/cli/src/service-container/service-container.ts

@@ -982,7 +982,12 @@ export class ServiceContainer {
 
     this.masterPasswordApiService = new MasterPasswordApiService(this.apiService, this.logService);
     const changeKdfApiService = new DefaultChangeKdfApiService(this.apiService);
-    const changeKdfService = new DefaultChangeKdfService(changeKdfApiService, this.sdkService);
+    const changeKdfService = new DefaultChangeKdfService(
+      changeKdfApiService,
+      this.sdkService,
+      this.keyService,
+      this.masterPasswordService,
+    );
     this.encryptedMigrator = new DefaultEncryptedMigrator(
       this.kdfConfigService,
       changeKdfService,

libs/angular/src/services/jslib-services.module.ts

@@ -528,7 +528,7 @@ const safeProviders: SafeProvider[] = [
   safeProvider({
     provide: ChangeKdfService,
     useClass: DefaultChangeKdfService,
-    deps: [ChangeKdfApiService, SdkService],
+    deps: [ChangeKdfApiService, SdkService, KeyService, InternalMasterPasswordServiceAbstraction],
   }),
   safeProvider({
     provide: EncryptedMigrator,
@@ -1333,7 +1333,7 @@ const safeProviders: SafeProvider[] = [
   safeProvider({
     provide: ChangeKdfService,
     useClass: DefaultChangeKdfService,
-    deps: [ChangeKdfApiService, SdkService],
+    deps: [ChangeKdfApiService, SdkService, KeyService, InternalMasterPasswordServiceAbstraction],
   }),
   safeProvider({
     provide: AuthRequestServiceAbstraction,

libs/common/src/key-management/kdf/change-kdf.service.spec.ts

@@ -2,13 +2,14 @@ import { mock } from "jest-mock-extended";
 import { of } from "rxjs";
 
 // eslint-disable-next-line no-restricted-imports
-import { PBKDF2KdfConfig } from "@bitwarden/key-management";
+import { KeyService, PBKDF2KdfConfig } from "@bitwarden/key-management";
 
 import { makeEncString } from "../../../spec";
 import { KdfRequest } from "../../models/request/kdf.request";
 import { SdkService } from "../../platform/abstractions/sdk/sdk.service";
 import { UserId } from "../../types/guid";
 import { EncString } from "../crypto/models/enc-string";
+import { InternalMasterPasswordServiceAbstraction } from "../master-password/abstractions/master-password.service.abstraction";
 import {
   MasterKeyWrappedUserKey,
   MasterPasswordAuthenticationHash,
@@ -22,6 +23,8 @@ import { DefaultChangeKdfService } from "./change-kdf.service";
 describe("ChangeKdfService", () => {
   const changeKdfApiService = mock<ChangeKdfApiService>();
   const sdkService = mock<SdkService>();
+  const keyService = mock<KeyService>();
+  const masterPasswordService = mock<InternalMasterPasswordServiceAbstraction>();
 
   let sut: DefaultChangeKdfService;
 
@@ -48,7 +51,12 @@ describe("ChangeKdfService", () => {
 
   beforeEach(() => {
     sdkService.userClient$ = jest.fn((userId: UserId) => of(mockSdk)) as any;
-    sut = new DefaultChangeKdfService(changeKdfApiService, sdkService);
+    sut = new DefaultChangeKdfService(
+      changeKdfApiService,
+      sdkService,
+      keyService,
+      masterPasswordService,
+    );
   });
 
   afterEach(() => {
@@ -163,6 +171,20 @@ describe("ChangeKdfService", () => {
       expect(changeKdfApiService.updateUserKdfParams).toHaveBeenCalledWith(expectedRequest);
     });
 
+    it("should set master key and hash after KDF update", async () => {
+      const masterPassword = "masterPassword";
+      const mockMasterKey = {} as any;
+      const mockHash = "localHash";
+
+      keyService.makeMasterKey.mockResolvedValue(mockMasterKey);
+      keyService.hashMasterKey.mockResolvedValue(mockHash);
+
+      await sut.updateUserKdfParams(masterPassword, mockNewKdfConfig, mockUserId);
+
+      expect(masterPasswordService.setMasterKey).toHaveBeenCalledWith(mockMasterKey, mockUserId);
+      expect(masterPasswordService.setMasterKeyHash).toHaveBeenCalledWith(mockHash, mockUserId);
+    });
+
     it("should properly dispose of SDK resources", async () => {
       const masterPassword = "masterPassword";
       jest.spyOn(mockNewKdfConfig, "toSdkConfig").mockReturnValue({} as any);

libs/common/src/key-management/kdf/change-kdf.service.ts

@@ -1,12 +1,14 @@
 import { firstValueFrom, map } from "rxjs";
 
 import { assertNonNullish } from "@bitwarden/common/auth/utils";
+import { HashPurpose } from "@bitwarden/common/platform/enums";
 import { UserId } from "@bitwarden/common/types/guid";
 // eslint-disable-next-line no-restricted-imports
-import { KdfConfig } from "@bitwarden/key-management";
+import { KdfConfig, KeyService } from "@bitwarden/key-management";
 
 import { KdfRequest } from "../../models/request/kdf.request";
 import { SdkService } from "../../platform/abstractions/sdk/sdk.service";
+import { InternalMasterPasswordServiceAbstraction } from "../master-password/abstractions/master-password.service.abstraction";
 import {
   fromSdkAuthenticationData,
   MasterPasswordAuthenticationData,
@@ -20,6 +22,8 @@ export class DefaultChangeKdfService implements ChangeKdfService {
   constructor(
     private changeKdfApiService: ChangeKdfApiService,
     private sdkService: SdkService,
+    private keyService: KeyService,
+    private masterPasswordService: InternalMasterPasswordServiceAbstraction,
   ) {}
 
   async updateUserKdfParams(masterPassword: string, kdf: KdfConfig, userId: UserId): Promise<void> {
@@ -56,5 +60,19 @@ export class DefaultChangeKdfService implements ChangeKdfService {
     const request = new KdfRequest(authenticationData, unlockData);
     request.authenticateWith(oldAuthenticationData);
     await this.changeKdfApiService.updateUserKdfParams(request);
+
+    // Update the locally stored master key and hash, so that UV, etc. still works
+    const masterKey = await this.keyService.makeMasterKey(
+      masterPassword,
+      unlockData.salt,
+      unlockData.kdf,
+    );
+    const localMasterKeyHash = await this.keyService.hashMasterKey(
+      masterPassword,
+      masterKey,
+      HashPurpose.LocalAuthorization,
+    );
+    await this.masterPasswordService.setMasterKeyHash(localMasterKeyHash, userId);
+    await this.masterPasswordService.setMasterKey(masterKey, userId);
   }
 }