Add support for dehydrated devices

src/Lifecycle.js

@@ -186,6 +186,8 @@ export function attemptTokenLogin(queryParams, defaultDeviceDisplayName) {
         console.log("Logged in with token");
         return _clearStorage().then(() => {
             _persistCredentialsToLocalStorage(creds);
+            // remember that we just logged in
+            sessionStorage.setItem("mx_fresh_login", true);
             return true;
         });
     }).catch((err) => {
@@ -312,6 +314,9 @@ async function _restoreFromLocalStorage(opts) {
             console.log("No pickle key available");
         }
 
+        const freshLogin = sessionStorage.getItem("mx_fresh_login");
+        sessionStorage.removeItem("mx_fresh_login");
+
         console.log(`Restoring session for ${userId}`);
         await _doSetLoggedIn({
             userId: userId,
@@ -321,6 +326,7 @@ async function _restoreFromLocalStorage(opts) {
             identityServerUrl: isUrl,
             guest: isGuest,
             pickleKey: pickleKey,
+            freshLogin: freshLogin,
         }, false);
         return true;
     } else {
@@ -364,6 +370,7 @@ async function _handleLoadSessionFailure(e) {
  * @returns {Promise} promise which resolves to the new MatrixClient once it has been started
  */
 export async function setLoggedIn(credentials) {
+    credentials.freshLogin = true;
     stopMatrixClient();
     const pickleKey = credentials.userId && credentials.deviceId
           ? await PlatformPeg.get().createPickleKey(credentials.userId, credentials.deviceId)
@@ -429,6 +436,7 @@ async function _doSetLoggedIn(credentials, clearStorage) {
         " guest: " + credentials.guest +
         " hs: " + credentials.homeserverUrl +
         " softLogout: " + softLogout,
+        " freshLogin: " + credentials.freshLogin,
     );
 
     // This is dispatched to indicate that the user is still in the process of logging in
@@ -462,10 +470,28 @@ async function _doSetLoggedIn(credentials, clearStorage) {
 
     Analytics.setLoggedIn(credentials.guest, credentials.homeserverUrl);
 
+    MatrixClientPeg.replaceUsingCreds(credentials);
+    const client = MatrixClientPeg.get();
+
+    if (credentials.freshLogin && SettingsStore.getValue("feature_dehydration")) {
+        // If we just logged in, try to rehydrate a device instead of using a
+        // new device.  If it succeeds, we'll get a new device ID, so make sure
+        // we persist that ID to localStorage
+        const newDeviceId = await client.rehydrateDevice();
+        if (newDeviceId) {
+            credentials.deviceId = newDeviceId;
+        }
+
+        delete credentials.freshLogin;
+    }
+
     if (localStorage) {
         try {
             _persistCredentialsToLocalStorage(credentials);
 
+            // make sure we don't think that it's a fresh login any more
+            sessionStorage.removeItem("mx_fresh_login");
+
             // The user registered as a PWLU (PassWord-Less User), the generated password
             // is cached here such that the user can change it at a later time.
             if (credentials.password) {
@@ -482,12 +508,10 @@ async function _doSetLoggedIn(credentials, clearStorage) {
         console.warn("No local storage available: can't persist session!");
     }
 
-    MatrixClientPeg.replaceUsingCreds(credentials);
-
     dis.dispatch({ action: 'on_logged_in' });
 
     await startMatrixClient(/*startSyncing=*/!softLogout);
-    return MatrixClientPeg.get();
+    return client;
 }
 
 function _showStorageEvictedDialog() {

src/MatrixClientPeg.ts

@@ -31,7 +31,7 @@ import {verificationMethods} from 'matrix-js-sdk/src/crypto';
 import MatrixClientBackedSettingsHandler from "./settings/handlers/MatrixClientBackedSettingsHandler";
 import * as StorageManager from './utils/StorageManager';
 import IdentityAuthClient from './IdentityAuthClient';
-import { crossSigningCallbacks } from './SecurityManager';
+import { crossSigningCallbacks, tryToUnlockSecretStorageWithDehydrationKey } from './SecurityManager';
 import {SHOW_QR_CODE_METHOD} from "matrix-js-sdk/src/crypto/verification/QRCode";
 
 export interface IMatrixClientCreds {
@@ -42,6 +42,7 @@ export interface IMatrixClientCreds {
     accessToken: string;
     guest: boolean;
     pickleKey?: string;
+    freshLogin?: boolean;
 }
 
 // TODO: Move this to the js-sdk
@@ -192,6 +193,7 @@ class _MatrixClientPeg implements IMatrixClientPeg {
                 this.matrixClient.setCryptoTrustCrossSignedDevices(
                     !SettingsStore.getValue('e2ee.manuallyVerifyAllSessions'),
                 );
+                await tryToUnlockSecretStorageWithDehydrationKey(this.matrixClient);
                 StorageManager.setCryptoInitialised(true);
             }
         } catch (e) {

src/SecurityManager.js

@@ -24,15 +24,21 @@ import {encodeBase64} from "matrix-js-sdk/src/crypto/olmlib";
 import { isSecureBackupRequired } from './utils/WellKnownUtils';
 import AccessSecretStorageDialog from './components/views/dialogs/security/AccessSecretStorageDialog';
 import RestoreKeyBackupDialog from './components/views/dialogs/security/RestoreKeyBackupDialog';
+import SettingsStore from "./settings/SettingsStore";
 
 // This stores the secret storage private keys in memory for the JS SDK. This is
 // only meant to act as a cache to avoid prompting the user multiple times
 // during the same single operation. Use `accessSecretStorage` below to scope a
 // single secret storage operation, as it will clear the cached keys once the
 // operation ends.
 let secretStorageKeys = {};
+let secretStorageKeyInfo = {};
 let secretStorageBeingAccessed = false;
 
+let nonInteractive = false;
+
+let dehydrationCache = {};
+
 function isCachingAllowed() {
     return secretStorageBeingAccessed;
 }
@@ -66,6 +72,20 @@ async function confirmToDismiss() {
     return !sure;
 }
 
+function makeInputToKey(keyInfo) {
+    return async ({ passphrase, recoveryKey }) => {
+        if (passphrase) {
+            return deriveKey(
+                passphrase,
+                keyInfo.passphrase.salt,
+                keyInfo.passphrase.iterations,
+            );
+        } else {
+            return decodeRecoveryKey(recoveryKey);
+        }
+    };
+}
+
 async function getSecretStorageKey({ keys: keyInfos }, ssssItemName) {
     const keyInfoEntries = Object.entries(keyInfos);
     if (keyInfoEntries.length > 1) {
@@ -78,17 +98,18 @@ async function getSecretStorageKey({ keys: keyInfos }, ssssItemName) {
         return [keyId, secretStorageKeys[keyId]];
     }
 
-    const inputToKey = async ({ passphrase, recoveryKey }) => {
-        if (passphrase) {
-            return deriveKey(
-                passphrase,
-                keyInfo.passphrase.salt,
-                keyInfo.passphrase.iterations,
-            );
-        } else {
-            return decodeRecoveryKey(recoveryKey);
+    if (dehydrationCache.key) {
+        if (await MatrixClientPeg.get().checkSecretStorageKey(dehydrationCache.key, keyInfo)) {
+            cacheSecretStorageKey(keyId, dehydrationCache.key, keyInfo);
+            return [keyId, dehydrationCache.key];
         }
-    };
+    }
+
+    if (nonInteractive) {
+        throw new Error("Could not unlock non-interactively");
+    }
+
+    const inputToKey = makeInputToKey(keyInfo);
     const { finished } = Modal.createTrackedDialog("Access Secret Storage dialog", "",
         AccessSecretStorageDialog,
         /* props= */
@@ -118,14 +139,56 @@ async function getSecretStorageKey({ keys: keyInfos }, ssssItemName) {
     const key = await inputToKey(input);
 
     // Save to cache to avoid future prompts in the current session
-    cacheSecretStorageKey(keyId, key);
+    cacheSecretStorageKey(keyId, key, keyInfo);
 
     return [keyId, key];
 }
 
-function cacheSecretStorageKey(keyId, key) {
+export async function getDehydrationKey(keyInfo, checkFunc) {
+    const inputToKey = makeInputToKey(keyInfo);
+    const { finished } = Modal.createTrackedDialog("Access Secret Storage dialog", "",
+        AccessSecretStorageDialog,
+        /* props= */
+        {
+            keyInfo,
+            checkPrivateKey: async (input) => {
+                const key = await inputToKey(input);
+                try {
+                    checkFunc(key);
+                    return true;
+                } catch (e) {
+                    return false;
+                }
+            },
+        },
+        /* className= */ null,
+        /* isPriorityModal= */ false,
+        /* isStaticModal= */ false,
+        /* options= */ {
+            onBeforeClose: async (reason) => {
+                if (reason === "backgroundClick") {
+                    return confirmToDismiss();
+                }
+                return true;
+            },
+        },
+    );
+    const [input] = await finished;
+    if (!input) {
+        throw new AccessCancelledError();
+    }
+    const key = await inputToKey(input);
+
+    // need to copy the key because rehydration (unpickling) will clobber it
+    dehydrationCache = {key: new Uint8Array(key), keyInfo};
+
+    return key;
+}
+
+function cacheSecretStorageKey(keyId, key, keyInfo) {
     if (isCachingAllowed()) {
         secretStorageKeys[keyId] = key;
+        secretStorageKeyInfo[keyId] = keyInfo;
     }
 }
 
@@ -176,6 +239,7 @@ export const crossSigningCallbacks = {
     getSecretStorageKey,
     cacheSecretStorageKey,
     onSecretRequested,
+    getDehydrationKey,
 };
 
 export async function promptForBackupPassphrase() {
@@ -262,6 +326,18 @@ export async function accessSecretStorage(func = async () => { }, forceReset = f
             await cli.bootstrapSecretStorage({
                 getKeyBackupPassphrase: promptForBackupPassphrase,
             });
+
+            const keyId = Object.keys(secretStorageKeys)[0];
+            if (keyId && SettingsStore.getValue("feature_dehydration")) {
+                const dehydrationKeyInfo =
+                      secretStorageKeyInfo[keyId] && secretStorageKeyInfo[keyId].passphrase
+                      ? {passphrase: secretStorageKeyInfo[keyId].passphrase}
+                      : {};
+                console.log("Setting dehydration key");
+                await cli.setDehydrationKey(secretStorageKeys[keyId], dehydrationKeyInfo, "Backup device");
+            } else {
+                console.log("Not setting dehydration key: no SSSS key found");
+            }
         }
 
         // `return await` needed here to ensure `finally` block runs after the
@@ -272,6 +348,57 @@ export async function accessSecretStorage(func = async () => { }, forceReset = f
         secretStorageBeingAccessed = false;
         if (!isCachingAllowed()) {
             secretStorageKeys = {};
+            secretStorageKeyInfo = {};
+        }
+    }
+}
+
+// FIXME: this function name is a bit of a mouthful
+export async function tryToUnlockSecretStorageWithDehydrationKey(client) {
+    const key = dehydrationCache.key;
+    let restoringBackup = false;
+    if (key && await client.isSecretStorageReady()) {
+        console.log("Trying to set up cross-signing using dehydration key");
+        secretStorageBeingAccessed = true;
+        nonInteractive = true;
+        try {
+            await client.checkOwnCrossSigningTrust();
+
+            // we also need to set a new dehydrated device to replace the
+            // device we rehydrated
+            const dehydrationKeyInfo =
+                  dehydrationCache.keyInfo && dehydrationCache.keyInfo.passphrase
+                  ? {passphrase: dehydrationCache.keyInfo.passphrase}
+                  : {};
+            await client.setDehydrationKey(key, dehydrationKeyInfo, "Backup device");
+
+            // and restore from backup
+            const backupInfo = await client.getKeyBackupVersion();
+            if (backupInfo) {
+                restoringBackup = true;
+                // don't await, because this can take a long time
+                client.restoreKeyBackupWithSecretStorage(backupInfo)
+                    .finally(() => {
+                        secretStorageBeingAccessed = false;
+                        nonInteractive = false;
+                        if (!isCachingAllowed()) {
+                            secretStorageKeys = {};
+                            secretStorageKeyInfo = {};
+                        }
+                    });
+            }
+        } finally {
+            dehydrationCache = {};
+            // the secret storage cache is needed for restoring from backup, so
+            // don't clear it yet if we're restoring from backup
+            if (!restoringBackup) {
+                secretStorageBeingAccessed = false;
+                nonInteractive = false;
+                if (!isCachingAllowed()) {
+                    secretStorageKeys = {};
+                    secretStorageKeyInfo = {};
+                }
+            }
         }
     }
 }

src/i18n/strings/en_EN.json

@@ -452,6 +452,7 @@
     "Support adding custom themes": "Support adding custom themes",
     "Show message previews for reactions in DMs": "Show message previews for reactions in DMs",
     "Show message previews for reactions in all rooms": "Show message previews for reactions in all rooms",
+    "Offline encrypted messaging using dehydrated devices": "Offline encrypted messaging using dehydrated devices",
     "Enable advanced debugging for the room list": "Enable advanced debugging for the room list",
     "Show info about bridges in room settings": "Show info about bridges in room settings",
     "Font size": "Font size",

src/settings/Settings.ts

@@ -186,6 +186,12 @@ export const SETTINGS: {[setting: string]: ISetting} = {
         supportedLevels: LEVELS_FEATURE,
         default: false,
     },
+    "feature_dehydration": {
+        isFeature: true,
+        displayName: _td("Offline encrypted messaging using dehydrated devices"),
+        supportedLevels: LEVELS_FEATURE,
+        default: false,
+    },
     "advancedRoomListLogging": {
         // TODO: Remove flag before launch: https://github.com/vector-im/element-web/issues/14231
         displayName: _td("Enable advanced debugging for the room list"),