core: add initial DEK support

packages/core/src/api/key-manager.ts

@@ -0,0 +1,145 @@
+/*
+This file is part of the Notesnook project (https://notesnook.com/)
+
+Copyright (C) 2023 Streetwriters (Private) Limited
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+import { Cipher, SerializedKey, SerializedKeyPair } from "@notesnook/crypto";
+import Database from ".";
+import { isCipher } from "../utils";
+
+const KEY_INFO = {
+  inboxKeys: {
+    type: "asymmetric"
+  },
+  attachmentsKey: {
+    type: "symmetric"
+  },
+  monographPasswordsKey: {
+    type: "symmetric"
+  },
+  dataEncryptionKey: {
+    type: "symmetric"
+  },
+  legacyDataEncryptionKey: {
+    type: "symmetric"
+  }
+} as const;
+
+export type KeyId = keyof typeof KEY_INFO;
+
+type WrapKeyReturnType<T extends SerializedKeyPair | SerializedKey> =
+  T extends SerializedKeyPair
+    ? { public: string; private: Cipher<"base64"> }
+    : Cipher<"base64">;
+
+type WrappedKey =
+  | Cipher<"base64">
+  | {
+      public: string;
+      private: Cipher<"base64">;
+    };
+
+type UnwrapKeyReturnType<T extends WrappedKey> = T extends {
+  public: string;
+  private: Cipher<"base64">;
+}
+  ? SerializedKeyPair
+  : SerializedKey;
+
+type KeyTypeFromId<TId extends KeyId> =
+  (typeof KEY_INFO)[TId]["type"] extends "symmetric"
+    ? Cipher<"base64">
+    : {
+        public: string;
+        private: Cipher<"base64">;
+      };
+
+export class KeyManager {
+  private cache = new Map<string, KeyTypeFromId<KeyId>>();
+  constructor(private readonly db: Database) {}
+
+  clearCache() {
+    this.cache.clear();
+  }
+
+  async get<TId extends KeyId>(
+    id: TId,
+    options: {
+      useCache?: boolean;
+      refetchUser?: boolean;
+    } = { refetchUser: true, useCache: true }
+  ): Promise<KeyTypeFromId<TId> | undefined> {
+    if (options.useCache && this.cache.has(id)) {
+      return this.cache.get(id) as KeyTypeFromId<TId>;
+    }
+    let user = await this.db.user.getUser();
+    if ((!user || !user[id]) && options.refetchUser) {
+      user = await this.db.user.fetchUser();
+    }
+    if (!user) return;
+
+    this.cache.set(id, user[id] as KeyTypeFromId<KeyId>);
+    return user[id] as KeyTypeFromId<TId>;
+  }
+
+  async unwrapKey<T extends WrappedKey>(
+    key: T,
+    wrappingKey: SerializedKey
+  ): Promise<UnwrapKeyReturnType<T>> {
+    if (isCipher(key))
+      return JSON.parse(
+        await this.db.storage().decrypt(wrappingKey, key)
+      ) as UnwrapKeyReturnType<T>;
+    else {
+      const privateKey = await this.db
+        .storage()
+        .decrypt(wrappingKey, key.private);
+      return {
+        publicKey: key.public,
+        privateKey
+      } as UnwrapKeyReturnType<T>;
+    }
+  }
+
+  async wrapKey<T extends SerializedKey | SerializedKeyPair>(
+    key: T,
+    wrappingKey: SerializedKey
+  ): Promise<WrapKeyReturnType<T>> {
+    if (!("publicKey" in key)) {
+      return (await this.db
+        .storage()
+        .encrypt(wrappingKey, JSON.stringify(key))) as WrapKeyReturnType<T>;
+    } else {
+      const encryptedPrivateKey = await this.db
+        .storage()
+        .encrypt(wrappingKey, (key as SerializedKeyPair).privateKey);
+      return {
+        public: (key as SerializedKeyPair).publicKey,
+        private: encryptedPrivateKey
+      } as WrapKeyReturnType<T>;
+    }
+  }
+
+  async rewrapKey<T extends WrappedKey>(
+    key: T,
+    oldWrappingKey: SerializedKey,
+    newWrappingKey: SerializedKey
+  ) {
+    const unwrappedKey = await this.unwrapKey(key, oldWrappingKey);
+    return await this.wrapKey(unwrappedKey, newWrappingKey);
+  }
+}

packages/core/src/api/sync/__tests__/collector.test.js

@@ -109,6 +109,223 @@ test("unlinked relation should get included in collector", () =>
     expect(items[0].items[0].id).toBe("cd93df7a4c64fbd5f100361d629ac5b5");
   }));
 
+test("collector should use latest key version for encryption", () =>
+  databaseTest().then(async (db) => {
+    await loginFakeUser(db);
+    const collector = new Collector(db);
+
+    const noteId = await db.notes.add(TEST_NOTE);
+
+    const items = await iteratorToArray(collector.collect(100, false));
+
+    // Find the note item
+    const noteItem = items.find((i) => i.type === "note");
+    expect(noteItem).toBeDefined();
+    expect(noteItem.items[0].keyVersion).toBeDefined();
+
+    // Should use the latest key version available
+    const keys = await db.user.getDataEncryptionKeys();
+    const latestKeyVersion = Math.max(...keys.map((k) => k.version));
+    expect(noteItem.items[0].keyVersion).toBe(latestKeyVersion);
+  }));
+
+test("collector should assign keyVersion to all encrypted items", () =>
+  databaseTest().then(async (db) => {
+    await loginFakeUser(db);
+    const collector = new Collector(db);
+
+    await db.notes.add(TEST_NOTE);
+    await db.notes.add({ ...TEST_NOTE, title: "Note 2" });
+    await db.notes.add({ ...TEST_NOTE, title: "Note 3" });
+
+    const items = await iteratorToArray(collector.collect(100, false));
+
+    // All items should have keyVersion set
+    for (const chunk of items) {
+      for (const item of chunk.items) {
+        expect(item.keyVersion).toBeDefined();
+        expect(typeof item.keyVersion).toBe("number");
+      }
+    }
+  }));
+
+test("sync roundtrip: items encrypted with keyVersion can be decrypted", () =>
+  databaseTest().then(async (db) => {
+    await loginFakeUser(db);
+    const { Sync } = await import("../index.ts");
+    const sync = new Sync(db);
+    const collector = new Collector(db);
+
+    const noteId = await db.notes.add({
+      ...TEST_NOTE,
+      title: "Sync Test Note"
+    });
+    const note = await db.notes.note(noteId);
+
+    const items = await iteratorToArray(collector.collect(100, false));
+    const noteChunk = items.find((i) => i.type === "note");
+
+    expect(noteChunk).toBeDefined();
+    expect(noteChunk.items[0].keyVersion).toBeDefined();
+
+    // Simulate receiving the same item back from server
+    const keys = await db.user.getDataEncryptionKeys();
+    await sync.processChunk(noteChunk, keys, { type: "fetch" });
+
+    // Verify the note is still intact
+    const syncedNote = await db.notes.note(noteId);
+    expect(syncedNote.title).toBe("Sync Test Note");
+    expect(syncedNote.id).toBe(note.id);
+  }));
+
+test("sync should handle mixed keyVersion items in same chunk", () =>
+  databaseTest().then(async (db) => {
+    await loginFakeUser(db);
+    const { Sync } = await import("../index.ts");
+    const sync = new Sync(db);
+
+    const keys = await db.user.getDataEncryptionKeys();
+
+    // Create mock items with different key versions
+    const note1 = JSON.stringify({
+      id: "note1",
+      type: "note",
+      title: "Note 1",
+      dateModified: Date.now()
+    });
+    const note2 = JSON.stringify({
+      id: "note2",
+      type: "note",
+      title: "Note 2",
+      dateModified: Date.now()
+    });
+
+    const cipher1 = await db.storage().encrypt(keys[0].key, note1);
+    const cipher2 =
+      keys.length > 1
+        ? await db.storage().encrypt(keys[1].key, note2)
+        : await db.storage().encrypt(keys[0].key, note2);
+
+    const chunk = {
+      type: "note",
+      count: 2,
+      items: [
+        { ...cipher1, id: "note1", v: 5, keyVersion: keys[0].version },
+        {
+          ...cipher2,
+          id: "note2",
+          v: 5,
+          keyVersion: keys.length > 1 ? keys[1].version : keys[0].version
+        }
+      ]
+    };
+
+    // Process the chunk with mixed key versions
+    await sync.processChunk(chunk, keys, { type: "fetch" });
+
+    // Verify both notes were decrypted correctly
+    const savedNote1 = await db.notes.note("note1");
+    const savedNote2 = await db.notes.note("note2");
+
+    expect(savedNote1).toBeDefined();
+    expect(savedNote2).toBeDefined();
+    expect(savedNote1.title).toBe("Note 1");
+    expect(savedNote2.title).toBe("Note 2");
+  }));
+
+test("sync should maintain stable ordering across decryptMulti", () =>
+  databaseTest().then(async (db) => {
+    await loginFakeUser(db);
+    const collector = new Collector(db);
+
+    // Create multiple notes with predictable order
+    const noteIds = [];
+    for (let i = 0; i < 5; i++) {
+      const id = await db.notes.add({
+        ...TEST_NOTE,
+        title: `Note ${i}`
+      });
+      noteIds.push(id);
+    }
+
+    const items = await iteratorToArray(collector.collect(100, false));
+    const noteChunk = items.find((i) => i.type === "note");
+
+    expect(noteChunk).toBeDefined();
+    expect(noteChunk.items).toHaveLength(5);
+
+    // Verify all items have IDs
+    const collectedIds = noteChunk.items.map((item) => item.id);
+    expect(collectedIds).toHaveLength(5);
+
+    // All IDs should be present
+    for (const id of noteIds) {
+      expect(collectedIds).toContain(id);
+    }
+
+    // Decrypt and verify ID mapping is preserved
+    const keys = await db.user.getDataEncryptionKeys();
+    const { Sync } = await import("../index.ts");
+    const sync = new Sync(db);
+
+    await sync.processChunk(noteChunk, keys, { type: "fetch" });
+
+    // Verify each note can be retrieved with correct content
+    for (let i = 0; i < 5; i++) {
+      const note = await db.notes.note(noteIds[i]);
+      expect(note).toBeDefined();
+      expect(note.title).toBe(`Note ${i}`);
+    }
+  }));
+
+test("sync should correctly select key based on keyVersion", () =>
+  databaseTest().then(async (db) => {
+    await loginFakeUser(db);
+    const { Sync } = await import("../index.ts");
+    const sync = new Sync(db);
+
+    const keys = await db.user.getDataEncryptionKeys();
+
+    // Create items encrypted with specific key versions
+    const testCases = keys.map((keyInfo, idx) => ({
+      id: `note${idx}`,
+      title: `Note with keyVersion ${keyInfo.version}`,
+      keyVersion: keyInfo.version,
+      key: keyInfo.key
+    }));
+
+    const chunks = [];
+    for (const testCase of testCases) {
+      const noteData = JSON.stringify({
+        id: testCase.id,
+        type: "note",
+        title: testCase.title,
+        dateModified: Date.now()
+      });
+      const cipher = await db.storage().encrypt(testCase.key, noteData);
+
+      chunks.push({
+        type: "note",
+        count: 1,
+        items: [
+          { ...cipher, id: testCase.id, v: 5, keyVersion: testCase.keyVersion }
+        ]
+      });
+    }
+
+    // Process each chunk
+    for (const chunk of chunks) {
+      await sync.processChunk(chunk, keys, { type: "fetch" });
+    }
+
+    // Verify each note was decrypted with the correct key
+    for (const testCase of testCases) {
+      const note = await db.notes.note(testCase.id);
+      expect(note).toBeDefined();
+      expect(note.title).toBe(testCase.title);
+    }
+  }));
+
 async function iteratorToArray(iterator) {
   let items = [];
   for await (const item of iterator) {

packages/core/src/api/sync/collector.ts

@@ -25,7 +25,8 @@ import {
   SyncItem,
   SyncTransferItem,
   SYNC_COLLECTIONS_MAP,
-  SYNC_ITEM_TYPES
+  SYNC_ITEM_TYPES,
+  KeyVersion
 } from "./types.js";
 import { Item, MaybeDeletedItem } from "../../types.js";
 
@@ -46,12 +47,17 @@ class Collector {
     chunkSize: number,
     isForceSync = false
   ): AsyncGenerator<SyncTransferItem, void, unknown> {
-    const key = await this.db.user.getEncryptionKey();
-    if (!key || !key.key || !key.salt) {
+    const keys = await this.db.user.getDataEncryptionKeys();
+    if (!keys || !keys.length) {
       EV.publish(EVENTS.userSessionExpired);
       throw new Error("User encryption key not generated. Please relogin.");
     }
 
+    // select the latest available key for encryption
+    const key = keys.reduce((max, current) =>
+      current.version > max.version ? current : max
+    );
+
     for (const itemType of SYNC_ITEM_TYPES) {
       const collectionKey = SYNC_COLLECTIONS_MAP[itemType];
       const collection = this.db[collectionKey].collection;
@@ -61,8 +67,8 @@ class Collector {
         if (!ids.length) continue;
         const ciphers = await this.db
           .storage()
-          .encryptMulti(key, syncableItems);
-        const items = toSyncItem(ids, ciphers);
+          .encryptMulti(key.key, syncableItems);
+        const items = toSyncItem(ids, ciphers, key.version);
         if (!items.length) continue;
         yield { items, type: itemType, count: items.length };
 
@@ -88,7 +94,11 @@ class Collector {
 }
 export default Collector;
 
-function toSyncItem(ids: string[], ciphers: Cipher<"base64">[]) {
+function toSyncItem(
+  ids: string[],
+  ciphers: Cipher<"base64">[],
+  keyVersion: KeyVersion
+) {
   if (ids.length !== ciphers.length)
     throw new Error("ids.length must be equal to ciphers.length");
 
@@ -98,6 +108,7 @@ function toSyncItem(ids: string[], ciphers: Cipher<"base64">[]) {
     const cipher = ciphers[i] as SyncItem;
     cipher.v = CURRENT_DATABASE_VERSION;
     cipher.id = id;
+    cipher.keyVersion = keyVersion;
     items.push(cipher);
   }
   return items;

packages/core/src/api/sync/index.ts

@@ -47,6 +47,8 @@ import {
   Notebook
 } from "../../types.js";
 import {
+  KEY_VERSION,
+  KeyVersion,
   SYNC_COLLECTIONS_MAP,
   SyncableItemType,
   SyncInboxItem,
@@ -55,7 +57,6 @@ import {
 import { DownloadableFile } from "../../database/fs.js";
 import { SyncDevices } from "./devices.js";
 import { DefaultColors } from "../../collections/colors.js";
-import { Monographs } from "../monographs.js";
 
 enum LogLevel {
   /** Log level for very low severity diagnostic messages. */
@@ -149,7 +150,7 @@ export default class SyncManager {
   }
 }
 
-class Sync {
+export class Sync {
   collector;
   merger;
   autoSync;
@@ -245,7 +246,7 @@ class Sync {
           "RequestFetchV3 failed, falling back to RequestFetchV2"
         );
         await this.connection?.invoke("RequestFetchV2", deviceId);
-      }
+      } else throw error;
     }
 
     if (this.conflictedNoteIds.length > 0) {
@@ -343,16 +344,51 @@ class Sync {
 
   async processChunk(
     chunk: SyncTransferItem,
-    key: SerializedKey,
+    keys: {
+      version: KeyVersion;
+      key: SerializedKey;
+    }[],
     options: SyncOptions
   ) {
     const itemType = chunk.type;
-    const decrypted = await this.db.storage().decryptMulti(key, chunk.items);
+    const decrypted: string[] = [];
+
+    // Pre-group items by keyVersion for O(1) lookups
+    const itemsByKeyVersion = new Map<KeyVersion, typeof chunk.items>();
+    const versionMap = new Map<string, number>();
+
+    for (const item of chunk.items) {
+      const keyVersion = item.keyVersion ?? KEY_VERSION.LEGACY;
+      const group = itemsByKeyVersion.get(keyVersion);
+      if (group) {
+        group.push(item);
+      } else {
+        itemsByKeyVersion.set(keyVersion, [item]);
+      }
+      versionMap.set(item.id, item.v);
+    }
+
+    for (const keyInfo of keys) {
+      const itemsToDecrypt = itemsByKeyVersion.get(keyInfo.version);
+      if (!itemsToDecrypt || itemsToDecrypt.length === 0) continue;
+
+      decrypted.push(
+        ...(await this.db.storage().decryptMulti(keyInfo.key, itemsToDecrypt))
+      );
+    }
 
     const deserialized: MaybeDeletedItem<Item>[] = [];
     for (let i = 0; i < decrypted.length; ++i) {
-      const decryptedItem = decrypted[i];
-      const version = chunk.items[i].v;
+      const decryptedItem = JSON.parse(decrypted[i]) as MaybeDeletedItem<Item>;
+      const version = versionMap.get(decryptedItem.id);
+      if (version === undefined) {
+        this.logger.error(
+          new Error(
+            `Version not found for item ${decryptedItem.id}. Skipping item.`
+          )
+        );
+        continue;
+      }
       const item = await deserializeItem(
         decryptedItem,
         itemType,
@@ -476,10 +512,15 @@ class Sync {
     this.connection.on("SendItems", async (chunk) => {
       if (this.connection?.state !== HubConnectionState.Connected) return false;
 
-      const key = await this.getKey();
-      if (!key) return false;
-
-      await this.processChunk(chunk, key, options);
+      const keys = await this.db.user.getDataEncryptionKeys();
+      if (!keys || !keys.length) {
+        this.logger.error(
+          new Error("User encryption keys not generated. Please relogin.")
+        );
+        EV.publish(EVENTS.userSessionExpired);
+        return false;
+      }
+      await this.processChunk(chunk, keys, options);
 
       sendSyncProgressEvent(this.db.eventManager, `download`, chunk.count);
 
@@ -513,18 +554,6 @@ class Sync {
     );
   }
 
-  private async getKey() {
-    const key = await this.db.user.getEncryptionKey();
-    if (!key?.key) {
-      this.logger.error(
-        new Error("User encryption key not generated. Please relogin.")
-      );
-      EV.publish(EVENTS.userSessionExpired);
-      return;
-    }
-    return key;
-  }
-
   private async checkConnection() {
     await this.syncConnectionMutex.runExclusive(async () => {
       try {
@@ -564,12 +593,11 @@ function promiseTimeout(ms: number, promise: Promise<unknown>) {
 }
 
 async function deserializeItem(
-  decryptedItem: string,
+  item: MaybeDeletedItem<Item>,
   type: SyncableItemType,
   version: number,
   database: Database
 ): Promise<MaybeDeletedItem<Item> | undefined> {
-  const item = JSON.parse(decryptedItem) as MaybeDeletedItem<Item>;
   item.remote = true;
   item.synced = true;
 

packages/core/src/api/sync/types.ts

@@ -19,9 +19,17 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import { Cipher } from "@notesnook/crypto";
 
+export const KEY_VERSION = {
+  LEGACY: 0,
+  DEK: 1
+} as const;
+
+export type KeyVersion = (typeof KEY_VERSION)[keyof typeof KEY_VERSION];
+
 export type SyncItem = {
   id: string;
   v: number;
+  keyVersion?: KeyVersion;
 } & Cipher<"base64">;
 
 export type SyncableItemType = keyof typeof SYNC_COLLECTIONS_MAP;

packages/core/src/api/user-manager.ts

@@ -24,8 +24,10 @@ import TokenManager from "./token-manager.js";
 import { EV, EVENTS } from "../common.js";
 import { HealthCheck } from "./healthcheck.js";
 import Database from "./index.js";
-import { SerializedKeyPair, SerializedKey, Cipher } from "@notesnook/crypto";
+import { SerializedKeyPair, SerializedKey } from "@notesnook/crypto";
 import { logger } from "../logger.js";
+import { KEY_VERSION, KeyVersion } from "./sync/types.js";
+import { KeyId, KeyManager } from "./key-manager.js";
 
 const ENDPOINTS = {
   signup: "/users",
@@ -42,11 +44,10 @@ const ENDPOINTS = {
 
 class UserManager {
   private tokenManager: TokenManager;
-  private cachedAttachmentKey?: SerializedKey;
-  private cachedMonographPasswordsKey?: SerializedKey;
-  private cachedInboxKeys?: SerializedKeyPair;
+  private keyManager: KeyManager;
   constructor(private readonly db: Database) {
-    this.tokenManager = new TokenManager(this.db.kv);
+    this.tokenManager = new TokenManager(db.kv);
+    this.keyManager = new KeyManager(db);
 
     EV.subscribe(EVENTS.userUnauthorized, async (url: string) => {
       if (url.includes("/connect/token") || !(await HealthCheck.auth())) return;
@@ -240,6 +241,17 @@ class UserManager {
     await this.db.setLastSynced(0);
     await this.db.syncer.devices.register();
 
+    // TODO: Uncomment this when we are done testing legacy password change
+    // support
+    // const masterKey = await this.getMasterKey();
+    // if (!masterKey) throw new Error("User encryption key not generated.");
+    // await this.updateUser({
+    //   dataEncryptionKey: await this.keyManager.wrapKey(
+    //     await this.db.crypto().generateRandomKey(),
+    //     masterKey
+    //   )
+    // });
+
     this.db.eventManager.publish(EVENTS.userLoggedIn, user);
   }
 
@@ -278,8 +290,7 @@ class UserManager {
     } catch (e) {
       logger.error(e, "Error logging out user.", { revoke, reason });
     } finally {
-      this.cachedAttachmentKey = undefined;
-      this.cachedInboxKeys = undefined;
+      this.keyManager.clearCache();
       await this.db.reset();
       this.db.eventManager.publish(EVENTS.userLoggedOut, reason);
       this.db.eventManager.publish(EVENTS.appRefreshRequested);
@@ -381,7 +392,7 @@ class UserManager {
   }
 
   changePassword(oldPassword: string, newPassword: string) {
-    return this._updatePassword("change_password", {
+    return this._updatePassword("change", {
       old_password: oldPassword,
       new_password: newPassword
     });
@@ -402,12 +413,46 @@ class UserManager {
   }
 
   resetPassword(newPassword: string) {
-    return this._updatePassword("reset_password", {
+    return this._updatePassword("reset", {
       new_password: newPassword
     });
   }
 
-  async getEncryptionKey(): Promise<SerializedKey | undefined> {
+  async getDataEncryptionKeys(): Promise<
+    { version: KeyVersion; key: SerializedKey }[] | undefined
+  > {
+    const masterKey = await this.getMasterKey();
+    if (!masterKey) return;
+
+    const dataEncryptionKey = await this.keyManager.get("dataEncryptionKey");
+    if (!dataEncryptionKey)
+      return [
+        {
+          key: masterKey,
+          version: KEY_VERSION.LEGACY
+        }
+      ];
+    const keys: { version: KeyVersion; key: SerializedKey }[] = [];
+
+    const legacyDataEncryptionKey = await this.keyManager.get(
+      "legacyDataEncryptionKey"
+    );
+    if (legacyDataEncryptionKey)
+      keys.push({
+        key: await this.keyManager.unwrapKey(
+          legacyDataEncryptionKey,
+          masterKey
+        ),
+        version: KEY_VERSION.LEGACY
+      });
+    keys.push({
+      key: await this.keyManager.unwrapKey(dataEncryptionKey, masterKey),
+      version: KEY_VERSION.DEK
+    });
+    return keys;
+  }
+
+  async getMasterKey(): Promise<SerializedKey | undefined> {
     const user = await this.getUser();
     if (!user) return;
     const key = await this.db.storage().getCryptoKey();
@@ -426,44 +471,28 @@ class UserManager {
     return { key, salt: user.salt };
   }
 
-  private async getUserKey<T>(config: {
-    getCache: () => T | undefined;
-    setCache: (key: T) => void;
-    userProperty: keyof User;
-    generateKey: () => Promise<T>;
-    errorContext: string;
-    decrypt: (user: User, userEncryptionKey: SerializedKey) => Promise<T>;
-    encrypt: (
-      key: T,
-      userEncryptionKey: SerializedKey
-    ) => Promise<Partial<User>>;
-  }): Promise<T | undefined> {
-    const cachedKey = config.getCache();
-    if (cachedKey) return cachedKey;
-
+  private async getUserKey(
+    id: KeyId,
+    config: {
+      generateKey: () => Promise<SerializedKey | SerializedKeyPair>;
+      errorContext: string;
+    }
+  ): Promise<SerializedKey | SerializedKeyPair | undefined> {
     try {
-      let user = await this.getUser();
-      if (!user) return;
+      const masterKey = await this.getMasterKey();
+      if (!masterKey) return;
 
-      if (!user[config.userProperty]) {
-        const token = await this.tokenManager.getAccessToken();
-        user = await http.get(`${constants.API_HOST}${ENDPOINTS.user}`, token);
-      }
-      if (!user) return;
+      const wrappedKey = await this.keyManager.get(id);
 
-      const userEncryptionKey = await this.getEncryptionKey();
-      if (!userEncryptionKey) return;
-
-      if (!user[config.userProperty]) {
+      if (!wrappedKey) {
         const key = await config.generateKey();
-        const updatePayload = await config.encrypt(key, userEncryptionKey);
-        await this.updateUser(updatePayload);
+        await this.updateUser({
+          [id]: await this.keyManager.wrapKey(key, masterKey)
+        });
         return key;
       }
 
-      const decryptedKey = await config.decrypt(user, userEncryptionKey);
-      config.setCache(decryptedKey);
-      return decryptedKey;
+      return await this.keyManager.unwrapKey(wrappedKey, masterKey);
     } catch (e) {
       logger.error(e, `Could not get ${config.errorContext}.`);
       if (e instanceof Error)
@@ -474,94 +503,27 @@ class UserManager {
   }
 
   async getAttachmentsKey() {
-    return this.getUserKey<SerializedKey>({
-      getCache: () => this.cachedAttachmentKey,
-      setCache: (key) => {
-        this.cachedAttachmentKey = key;
-      },
-      userProperty: "attachmentsKey",
+    return this.getUserKey("attachmentsKey", {
       generateKey: () => this.db.crypto().generateRandomKey(),
-      errorContext: "attachments encryption key",
-      encrypt: async (key, userEncryptionKey) => {
-        const encryptedKey = await this.db
-          .storage()
-          .encrypt(userEncryptionKey, JSON.stringify(key));
-        return { attachmentsKey: encryptedKey };
-      },
-      decrypt: async (user, userEncryptionKey) => {
-        const encryptedKey = user.attachmentsKey as Cipher<"base64">;
-        const plainData = await this.db
-          .storage()
-          .decrypt(userEncryptionKey, encryptedKey);
-        if (!plainData) throw new Error("Failed to decrypt attachments key");
-        return JSON.parse(plainData) as SerializedKey;
-      }
+      errorContext: "attachments encryption key"
     });
   }
 
   async getMonographPasswordsKey() {
-    return this.getUserKey<SerializedKey>({
-      getCache: () => this.cachedMonographPasswordsKey,
-      setCache: (key) => {
-        this.cachedMonographPasswordsKey = key;
-      },
-      userProperty: "monographPasswordsKey",
+    return this.getUserKey("monographPasswordsKey", {
       generateKey: () => this.db.crypto().generateRandomKey(),
-      errorContext: "monographs encryption key",
-      encrypt: async (key, userEncryptionKey) => {
-        const encryptedKey = await this.db
-          .storage()
-          .encrypt(userEncryptionKey, JSON.stringify(key));
-        return { monographPasswordsKey: encryptedKey };
-      },
-      decrypt: async (user, userEncryptionKey) => {
-        const encryptedKey = user.monographPasswordsKey as Cipher<"base64">;
-        const plainData = await this.db
-          .storage()
-          .decrypt(userEncryptionKey, encryptedKey);
-        if (!plainData)
-          throw new Error("Failed to decrypt monograph passwords key");
-        return JSON.parse(plainData) as SerializedKey;
-      }
+      errorContext: "monographs encryption key"
     });
   }
 
   async getInboxKeys() {
-    return this.getUserKey<SerializedKeyPair>({
-      getCache: () => this.cachedInboxKeys,
-      setCache: (key) => {
-        this.cachedInboxKeys = key;
-      },
-      userProperty: "inboxKeys",
+    return this.getUserKey("inboxKeys", {
       generateKey: () => this.db.crypto().generateCryptoKeyPair(),
-      errorContext: "inbox encryption keys",
-      encrypt: async (keys, userEncryptionKey) => {
-        const encryptedPrivateKey = await this.db
-          .storage()
-          .encrypt(userEncryptionKey, JSON.stringify(keys.privateKey));
-        return {
-          inboxKeys: {
-            public: keys.publicKey,
-            private: encryptedPrivateKey
-          }
-        };
-      },
-      decrypt: async (user, userEncryptionKey) => {
-        if (!user.inboxKeys) throw new Error("Inbox keys not found");
-        const decryptedPrivateKey = await this.db
-          .storage()
-          .decrypt(userEncryptionKey, user.inboxKeys.private);
-        return {
-          publicKey: user.inboxKeys.public,
-          privateKey: JSON.parse(decryptedPrivateKey)
-        };
-      }
+      errorContext: "inbox encryption keys"
     });
   }
 
   async hasInboxKeys() {
-    if (this.cachedInboxKeys) return true;
-
     const user = await this.getUser();
     if (!user) return false;
 
@@ -569,7 +531,7 @@ class UserManager {
   }
 
   async discardInboxKeys() {
-    this.cachedInboxKeys = undefined;
+    this.keyManager.clearCache();
 
     const user = await this.getUser();
     if (!user) return;
@@ -627,19 +589,20 @@ class UserManager {
   async verifyPassword(password: string) {
     try {
       const user = await this.getUser();
-      const key = await this.getEncryptionKey();
+      const key = await this.getMasterKey();
       if (!user || !key) return false;
 
       const cipher = await this.db.storage().encrypt(key, "notesnook");
       const plainText = await this.db.storage().decrypt({ password }, cipher);
       return plainText === "notesnook";
     } catch (e) {
+      logger.error(e);
       return false;
     }
   }
 
   async _updatePassword(
-    type: "change_password" | "reset_password",
+    type: "change" | "reset",
     data: {
       new_password: string;
       old_password?: string;
@@ -652,98 +615,97 @@ class UserManager {
 
     const { email, salt } = user;
 
-    let { new_password, old_password } = data;
+    const { new_password, old_password } = data;
     if (old_password && !(await this.verifyPassword(old_password)))
       throw new Error("Incorrect old password.");
 
+    const oldPassword = old_password
+      ? await this.db.storage().hash(old_password, email, {
+          usesFallback: await this.usesFallbackPWHash(old_password)
+        })
+      : null;
+
     if (!new_password) throw new Error("New password is required.");
 
-    data.encryptionKey = data.encryptionKey || (await this.getEncryptionKey());
+    data.encryptionKey = data.encryptionKey || (await this.getMasterKey());
+    if (!data.encryptionKey)
+      throw new Error("User encryption key not generated.");
 
-    await this.clearSessions();
+    const newMasterKey = await this.db
+      .storage()
+      .generateCryptoKey(new_password, salt);
 
-    if (data.encryptionKey) await this.db.sync({ type: "fetch", force: true });
+    const updateUserPayload: Partial<User> = {};
+    if (user.attachmentsKey) {
+      updateUserPayload.attachmentsKey = await this.keyManager.rewrapKey(
+        user.attachmentsKey,
+        data.encryptionKey,
+        newMasterKey
+      );
+    }
+    if (user.monographPasswordsKey) {
+      updateUserPayload.monographPasswordsKey = await this.keyManager.rewrapKey(
+        user.monographPasswordsKey,
+        data.encryptionKey,
+        newMasterKey
+      );
+    }
+    if (user.inboxKeys) {
+      updateUserPayload.inboxKeys = await this.keyManager.rewrapKey(
+        user.inboxKeys,
+        data.encryptionKey,
+        newMasterKey
+      );
+    }
 
-    if (old_password)
-      old_password = await this.db.storage().hash(old_password, email, {
-        usesFallback: await this.usesFallbackPWHash(old_password)
-      });
+    if (user.legacyDataEncryptionKey)
+      updateUserPayload.legacyDataEncryptionKey =
+        await this.keyManager.rewrapKey(
+          user.legacyDataEncryptionKey,
+          data.encryptionKey,
+          newMasterKey
+        );
+    if (user.dataEncryptionKey)
+      updateUserPayload.dataEncryptionKey = await this.keyManager.rewrapKey(
+        user.dataEncryptionKey,
+        data.encryptionKey,
+        newMasterKey
+      );
+    else {
+      updateUserPayload.dataEncryptionKey = await this.keyManager.wrapKey(
+        await this.db.crypto().generateRandomKey(),
+        newMasterKey
+      );
+      updateUserPayload.legacyDataEncryptionKey = await this.keyManager.wrapKey(
+        data.encryptionKey,
+        newMasterKey
+      );
+    }
 
-    // retrieve user keys before deriving a new encryption key
-    const oldUserKeys = {
-      attachmentsKey: await this.getAttachmentsKey(),
-      monographPasswordsKey: await this.getMonographPasswordsKey(),
-      inboxKeys: (await this.hasInboxKeys())
-        ? await this.getInboxKeys()
-        : undefined
-    } as const;
+    await http.patch.json(
+      `${constants.API_HOST}/users/password/${type}`,
+      {
+        oldPassword: oldPassword,
+        newPassword: await this.db.storage().hash(new_password, email),
+        userKeys: updateUserPayload
+      },
+      token
+    );
 
     await this.db.storage().deriveCryptoKey({
       password: new_password,
       salt
     });
 
-    if (!(await this.resetUser(false))) return;
-
-    await this.db.sync({ type: "send", force: true });
-
-    const userEncryptionKey = await this.getEncryptionKey();
-    if (userEncryptionKey) {
-      const updateUserPayload: Partial<User> = {};
-      if (oldUserKeys.attachmentsKey) {
-        user.attachmentsKey = await this.db
-          .storage()
-          .encrypt(
-            userEncryptionKey,
-            JSON.stringify(oldUserKeys.attachmentsKey)
-          );
-        updateUserPayload.attachmentsKey = user.attachmentsKey;
-      }
-      if (oldUserKeys.monographPasswordsKey) {
-        user.monographPasswordsKey = await this.db
-          .storage()
-          .encrypt(
-            userEncryptionKey,
-            JSON.stringify(oldUserKeys.monographPasswordsKey)
-          );
-        updateUserPayload.monographPasswordsKey = user.monographPasswordsKey;
-      }
-      if (oldUserKeys.inboxKeys) {
-        user.inboxKeys = {
-          public: oldUserKeys.inboxKeys.publicKey,
-          private: await this.db
-            .storage()
-            .encrypt(
-              userEncryptionKey,
-              JSON.stringify(oldUserKeys.inboxKeys.privateKey)
-            )
-        };
-        updateUserPayload.inboxKeys = user.inboxKeys;
-      }
-      if (Object.keys(updateUserPayload).length > 0) {
-        await this.updateUser(updateUserPayload);
-      }
-    }
-
-    if (new_password)
-      new_password = await this.db.storage().hash(new_password, email);
-
-    await http.patch(
-      `${constants.AUTH_HOST}${ENDPOINTS.patchUser}`,
-      {
-        type,
-        old_password,
-        new_password
-      },
-      token
-    );
+    this.keyManager.clearCache();
+    await this.setUser({ ...user, ...updateUserPayload });
 
     return true;
   }
 
   private async usesFallbackPWHash(password: string) {
     const user = await this.getUser();
-    const encryptionKey = await this.getEncryptionKey();
+    const encryptionKey = await this.getMasterKey();
     if (!user || !encryptionKey) return false;
     const fallbackCryptoKey = await this.db
       .storage()

packages/core/src/database/backup.ts

@@ -309,8 +309,8 @@ export default class Backup {
     if (encrypt && !user)
       throw new Error("Please login to create encrypted backups.");
 
-    const key = await this.db.user.getEncryptionKey();
-    if (encrypt && !key) throw new Error("No encryption key found.");
+    const key = await this.db.user.getMasterKey();
+    if (encrypt && !key) throw new Error("No master key found.");
 
     yield {
       type: "file",

packages/core/src/types.ts

@@ -599,6 +599,9 @@ export type User = {
   attachmentsKey?: Cipher<"base64">;
   monographPasswordsKey?: Cipher<"base64">;
   inboxKeys?: { public: string; private: Cipher<"base64"> };
+  dataEncryptionKey?: Cipher<"base64">;
+  legacyDataEncryptionKey?: Cipher<"base64">;
+
   marketingConsent?: boolean;
   storageUsed?: number;
   totalStorage?: number;