From a1ba5d7dba9ef7e3db4bc24cafe450c6d513fe34 Mon Sep 17 00:00:00 2001
From: Bruno Windels <brunow@matrix.org>
Date: Fri, 28 Aug 2020 13:58:17 +0200
Subject: [PATCH] between syncs, generate more otks if needed, and upload them

---
 src/matrix/Session.js      | 10 ++++++++++
 src/matrix/Sync.js         |  6 ++++++
 src/matrix/e2ee/Account.js | 25 +++++++++++++++++++++++++
 3 files changed, 41 insertions(+)

diff --git a/src/matrix/Session.js b/src/matrix/Session.js
index 90d29b61..bc52227a 100644
--- a/src/matrix/Session.js
+++ b/src/matrix/Session.js
@@ -60,6 +60,7 @@ export class Session {
                 }
                 await txn.complete();
             }
+            await this._e2eeAccount.generateOTKsIfNeeded(this._storage);
             await this._e2eeAccount.uploadKeys(this._storage);
         }
     }
@@ -178,6 +179,15 @@ export class Session {
         }
     }
 
+    async afterSyncCompleted() {
+        const needsToUploadOTKs = await this._e2eeAccount.generateOTKsIfNeeded(this._storage);
+        if (needsToUploadOTKs) {
+            // TODO: we could do this in parallel with sync if it proves to be too slow
+            // but I'm not sure how to not swallow errors in that case
+            await this._e2eeAccount.uploadKeys(this._storage);
+        }
+    }
+
     get syncToken() {
         return this._syncInfo?.token;
     }
diff --git a/src/matrix/Sync.js b/src/matrix/Sync.js
index e14451a9..c7aaaa99 100644
--- a/src/matrix/Sync.js
+++ b/src/matrix/Sync.js
@@ -100,6 +100,12 @@ export class Sync {
                     this._status.set(SyncStatus.Stopped);
                 }
             }
+            try {
+                await this._session.afterSyncCompleted();
+            } catch (err) {
+                console.err("error during after sync completed, continuing to sync.",  err.stack);
+                // swallowing error here apart from logging
+            }
         }
     }
 
diff --git a/src/matrix/e2ee/Account.js b/src/matrix/e2ee/Account.js
index 90fabc72..b8c39826 100644
--- a/src/matrix/e2ee/Account.js
+++ b/src/matrix/e2ee/Account.js
@@ -95,6 +95,31 @@ export class Account {
         }
     }
 
+    async generateOTKsIfNeeded(storage) {
+        const maxOTKs = this._account.max_number_of_one_time_keys();
+        const limit = maxOTKs / 2;
+        if (this._serverOTKCount < limit) {
+            // TODO: cache unpublishedOTKCount, so we don't have to parse this JSON on every sync iteration
+            // for now, we only determine it when serverOTKCount is sufficiently low, which is should rarely be,
+            // and recheck
+            const oneTimeKeys = JSON.parse(this._account.one_time_keys());
+            const oneTimeKeysEntries = Object.entries(oneTimeKeys.curve25519);
+            const unpublishedOTKCount = oneTimeKeysEntries.length;
+            const totalOTKCount = this._serverOTKCount + unpublishedOTKCount;
+            if (totalOTKCount < limit) {
+                // we could in theory also generated the keys and store them in
+                // writeSync, but then we would have to clone the account to avoid side-effects.
+                await this._updateSessionStorage(storage, sessionStore => {
+                    const newKeyCount = maxOTKs - totalOTKCount;
+                    this._account.generate_one_time_keys(newKeyCount);
+                    sessionStore.set(ACCOUNT_SESSION_KEY, this._account.pickle(this._pickleKey));
+                });
+                return true;
+            }
+        }
+        return false;
+    }
+
     writeSync(deviceOneTimeKeysCount, txn) {
         // we only upload signed_curve25519 otks
         const otkCount = deviceOneTimeKeysCount.signed_curve25519;