Improve fallback key behaviour (#2037)

Author: uhoreg

package.json

@@ -75,7 +75,7 @@
     "@babel/preset-env": "^7.12.11",
     "@babel/preset-typescript": "^7.12.7",
     "@babel/register": "^7.12.10",
-    "@matrix-org/olm": "https://gitlab.matrix.org/api/v4/projects/27/packages/npm/@matrix-org/olm/-/@matrix-org/olm-3.2.3.tgz",
+    "@matrix-org/olm": "https://gitlab.matrix.org/api/v4/projects/27/packages/npm/@matrix-org/olm/-/@matrix-org/olm-3.2.7.tgz",
     "@types/bs58": "^4.0.1",
     "@types/jest": "^26.0.20",
     "@types/node": "12",

src/crypto/OlmDevice.ts

@@ -543,13 +543,25 @@ export class OlmDevice {
             'readonly', [IndexedDBCryptoStore.STORE_ACCOUNT],
             (txn) => {
                 this.getAccount(txn, (account: Account) => {
-                    result = JSON.parse(account.fallback_key());
+                    result = JSON.parse(account.unpublished_fallback_key());
                 });
             },
         );
         return result;
     }
 
+    public async forgetOldFallbackKey(): Promise<void> {
+        await this.cryptoStore.doTxn(
+            'readwrite', [IndexedDBCryptoStore.STORE_ACCOUNT],
+            (txn) => {
+                this.getAccount(txn, (account: Account) => {
+                    account.forget_old_fallback_key();
+                    this.storeAccount(txn, account);
+                });
+            },
+        );
+    }
+
     /**
      * Generate a new outbound session
      *

src/crypto/index.ts

@@ -248,6 +248,7 @@ export class Crypto extends EventEmitter {
 
     private oneTimeKeyCount: number;
     private needsNewFallback: boolean;
+    private fallbackCleanup?: number; // setTimeout ID
 
     /**
      * Cryptography bits
@@ -1850,8 +1851,23 @@ export class Crypto extends EventEmitter {
                 }
 
                 if (this.getNeedsNewFallback()) {
-                    logger.info("generating fallback key");
-                    await this.olmDevice.generateFallbackKey();
+                    const fallbackKeys = await this.olmDevice.getFallbackKey();
+                    // if fallbackKeys is non-empty, we've already generated a
+                    // fallback key, but it hasn't been published yet, so we
+                    // can use that instead of generating a new one
+                    if (!fallbackKeys.curve25519 ||
+                        Object.keys(fallbackKeys.curve25519).length == 0) {
+                        logger.info("generating fallback key");
+                        if (this.fallbackCleanup) {
+                            // cancel any pending fallback cleanup because generating
+                            // a new fallback key will already drop the old fallback
+                            // that would have been dropped, and we don't want to kill
+                            // the current key
+                            clearTimeout(this.fallbackCleanup);
+                            delete this.fallbackCleanup;
+                        }
+                        await this.olmDevice.generateFallbackKey();
+                    }
                 }
 
                 logger.info("calling uploadOneTimeKeys");
@@ -1898,8 +1914,9 @@ export class Crypto extends EventEmitter {
     private async uploadOneTimeKeys() {
         const promises = [];
 
-        const fallbackJson: Record<string, IOneTimeKey> = {};
+        let fallbackJson: Record<string, IOneTimeKey>;
         if (this.getNeedsNewFallback()) {
+            fallbackJson = {};
             const fallbackKeys = await this.olmDevice.getFallbackKey();
             for (const [keyId, key] of Object.entries(fallbackKeys.curve25519)) {
                 const k = { key, fallback: true };
@@ -1924,10 +1941,23 @@ export class Crypto extends EventEmitter {
 
         await Promise.all(promises);
 
-        const res = await this.baseApis.uploadKeysRequest({
+        const requestBody: Record<string, any> = {
             "one_time_keys": oneTimeJson,
-            "org.matrix.msc2732.fallback_keys": fallbackJson,
-        });
+        };
+
+        if (fallbackJson) {
+            requestBody["org.matrix.msc2732.fallback_keys"] = fallbackJson;
+            requestBody["fallback_keys"] = fallbackJson;
+        }
+
+        const res = await this.baseApis.uploadKeysRequest(requestBody);
+
+        if (fallbackJson) {
+            this.fallbackCleanup = setTimeout(() => {
+                delete this.fallbackCleanup;
+                this.olmDevice.forgetOldFallbackKey();
+            }, 60*60*1000);
+        }
 
         await this.olmDevice.markKeysAsPublished();
         return res;

src/sync.ts

@@ -1416,11 +1416,14 @@ export class SyncApi {
             const currentCount = data.device_one_time_keys_count.signed_curve25519 || 0;
             this.opts.crypto.updateOneTimeKeyCount(currentCount);
         }
-        if (this.opts.crypto && data["org.matrix.msc2732.device_unused_fallback_key_types"]) {
+        if (this.opts.crypto &&
+            (data["device_unused_fallback_key_types"] ||
+                data["org.matrix.msc2732.device_unused_fallback_key_types"])) {
             // The presence of device_unused_fallback_key_types indicates that the
             // server supports fallback keys. If there's no unused
             // signed_curve25519 fallback key we need a new one.
-            const unusedFallbackKeys = data["org.matrix.msc2732.device_unused_fallback_key_types"];
+            const unusedFallbackKeys = data["device_unused_fallback_key_types"] ||
+                data["org.matrix.msc2732.device_unused_fallback_key_types"];
             this.opts.crypto.setNeedsNewFallback(
                 unusedFallbackKeys instanceof Array &&
                 !unusedFallbackKeys.includes("signed_curve25519"),