Move Olm account to IndexedDB

src/crypto/OlmDevice.js

@@ -79,8 +79,9 @@ function checkPayloadLength(payloadString) {
  * @property {string} deviceCurve25519Key   Curve25519 key for the account
  * @property {string} deviceEd25519Key      Ed25519 key for the account
  */
-function OlmDevice(sessionStore) {
+function OlmDevice(sessionStore, cryptoStore) {
     this._sessionStore = sessionStore;
+    this._cryptoStore = cryptoStore;
     this._pickleKey = "DEFAULT_KEY";
 
     // don't know these until we load the account from storage in init()
@@ -124,7 +125,7 @@ OlmDevice.prototype.init = async function() {
     let e2eKeys;
     const account = new Olm.Account();
     try {
-        _initialise_account(this._sessionStore, this._pickleKey, account);
+        await _initialise_account(this._cryptoStore, this._pickleKey, account);
         e2eKeys = JSON.parse(account.identity_keys());
 
         this._maxOneTimeKeys = account.max_number_of_one_time_keys();
@@ -137,16 +138,16 @@ OlmDevice.prototype.init = async function() {
 };
 
 
-function _initialise_account(sessionStore, pickleKey, account) {
-    const e2eAccount = sessionStore.getEndToEndAccount();
-    if (e2eAccount !== null) {
-        account.unpickle(pickleKey, e2eAccount);
+async function _initialise_account(cryptoStore, pickleKey, account) {
+    const accountTxn = await cryptoStore.endToEndAccountTransaction();
+    if (accountTxn.account !== null) {
+        account.unpickle(pickleKey, accountTxn.account);
         return;
     }
 
     account.create();
     const pickled = account.pickle(pickleKey);
-    sessionStore.storeEndToEndAccount(pickled);
+    await accountTxn.save(pickled);
 }
 
 /**
@@ -158,21 +159,36 @@ OlmDevice.getOlmVersion = function() {
 
 
 /**
- * extract our OlmAccount from the session store and call the given function
+ * extract our OlmAccount from the crypto store and call the given function
+ * with the account object and a 'save' function which returns a promise.
+ * The function will not be awaited upon and the save function must be
+ * called before the function returns, or not at all.
  *
  * @param {function} func
  * @return {object} result of func
  * @private
  */
-OlmDevice.prototype._getAccount = function(func) {
-    const account = new Olm.Account();
+OlmDevice.prototype._getAccount = async function(func) {
+    let result;
+
+    let account = null;
     try {
-        const pickledAccount = this._sessionStore.getEndToEndAccount();
-        account.unpickle(this._pickleKey, pickledAccount);
-        return func(account);
+        const accountTxn = await this._cryptoStore.endToEndAccountTransaction();
+        // Olm has a limited stack size so we must tightly control the number of
+        // Olm account objects in existence at any given time: once created, it
+        // must be destroyed again before we await.
+        account = new Olm.Account();
+        account.unpickle(this._pickleKey, accountTxn.account);
+
+        result = func(account, () => {
+                const pickledAccount = account.pickle(this._pickleKey);
+                return accountTxn.save(pickledAccount);
+            }
+        );
     } finally {
-        account.free();
+        if (account !== null) account.free();
     }
+    return result;
 };
 
 
@@ -182,9 +198,9 @@ OlmDevice.prototype._getAccount = function(func) {
  * @param {OlmAccount} account
  * @private
  */
-OlmDevice.prototype._saveAccount = function(account) {
+OlmDevice.prototype._saveAccount = async function(account) {
     const pickledAccount = account.pickle(this._pickleKey);
-    this._sessionStore.storeEndToEndAccount(pickledAccount);
+    await this._cryptoStore.storeEndToEndAccount(pickledAccount);
 };
 
 
@@ -250,7 +266,7 @@ OlmDevice.prototype._getUtility = function(func) {
  * @return {Promise<string>} base64-encoded signature
  */
 OlmDevice.prototype.sign = async function(message) {
-    return this._getAccount(function(account) {
+    return await this._getAccount(function(account) {
         return account.sign(message);
     });
 };
@@ -263,7 +279,7 @@ OlmDevice.prototype.sign = async function(message) {
  * key.
  */
 OlmDevice.prototype.getOneTimeKeys = async function() {
-    return this._getAccount(function(account) {
+    return await this._getAccount(function(account) {
         return JSON.parse(account.one_time_keys());
     });
 };
@@ -283,9 +299,9 @@ OlmDevice.prototype.maxNumberOfOneTimeKeys = function() {
  */
 OlmDevice.prototype.markKeysAsPublished = async function() {
     const self = this;
-    this._getAccount(function(account) {
+    await this._getAccount(function(account, save) {
         account.mark_keys_as_published();
-        self._saveAccount(account);
+        return save();
     });
 };
 
@@ -296,9 +312,9 @@ OlmDevice.prototype.markKeysAsPublished = async function() {
  */
 OlmDevice.prototype.generateOneTimeKeys = async function(numKeys) {
     const self = this;
-    this._getAccount(function(account) {
+    return this._getAccount(function(account, save) {
         account.generate_one_time_keys(numKeys);
-        self._saveAccount(account);
+        return save();
     });
 };
 
@@ -315,11 +331,12 @@ OlmDevice.prototype.createOutboundSession = async function(
     theirIdentityKey, theirOneTimeKey,
 ) {
     const self = this;
-    return this._getAccount(function(account) {
+    return await this._getAccount(async function(account, save) {
         const session = new Olm.Session();
         try {
             session.create_outbound(account, theirIdentityKey, theirOneTimeKey);
-            self._saveSession(theirIdentityKey, session);
+            await save();
+            await self._saveSession(theirIdentityKey, session);
             return session.session_id();
         } finally {
             session.free();
@@ -349,12 +366,12 @@ OlmDevice.prototype.createInboundSession = async function(
     }
 
     const self = this;
-    return this._getAccount(function(account) {
+    return await this._getAccount(async function(account, save) {
         const session = new Olm.Session();
         try {
             session.create_inbound_from(account, theirDeviceIdentityKey, ciphertext);
             account.remove_one_time_keys(session);
-            self._saveAccount(account);
+            await save();
 
             const payloadString = session.decrypt(message_type, ciphertext);
 

src/crypto/index.js

@@ -67,7 +67,7 @@ function Crypto(baseApis, sessionStore, userId, deviceId,
     this._clientStore = clientStore;
     this._cryptoStore = cryptoStore;
 
-    this._olmDevice = new OlmDevice(sessionStore);
+    this._olmDevice = new OlmDevice(sessionStore, cryptoStore);
     this._deviceList = new DeviceList(baseApis, sessionStore, this._olmDevice);
 
     // the last time we did a check for the number of one-time-keys on the

src/crypto/store/indexeddb-crypto-store-backend.js

@@ -1,7 +1,7 @@
 import Promise from 'bluebird';
 import utils from '../../utils';
 
-export const VERSION = 1;
+export const VERSION = 2;
 
 /**
  * Implementation of a CryptoStore which is backed by an existing
@@ -257,6 +257,42 @@ export class Backend {
         };
         return promiseifyTxn(txn);
     }
+
+    /**
+     * Load the end to end account for the logged-in user, giving an object
+     * that has the base64 encoded account string and a method for saving
+     * the account string back to the database. This allows the account
+     * to be read and writen atomically.
+     * @return {Promise<Object>} Object
+     * @return {Promise<Object>.account} Base64 encoded account.
+     * @return {Promise<Object>.save} Function to save account data back.
+     *     Takes base64 encoded account data, returns a promise.
+     */
+    endToEndAccountTransaction() {
+        const txn = this._db.transaction("account", "readwrite");
+        const objectStore = txn.objectStore("account");
+
+
+        return new Promise((resolve, reject) => {
+            const getReq = objectStore.get("-");
+            // We resolve on success here rather than on complete:
+            // the caller may wish to save the account back, which needs
+            // to be done while the transaction is still open (ie. before
+            // oncomplete)
+            getReq.onsuccess = function() {
+                resolve({
+                    account: getReq.result || null,
+                    save: (newData) => {
+                        const saveReq = objectStore.put(newData, "-");
+                        return promiseifyTxn(txn);
+                    },
+                });
+            };
+            getReq.onerror = reject;
+        });
+
+        return promiseifyTxn(txn).then(() => returnObj);
+    }
 }
 
 export function upgradeDatabase(db, oldVersion) {
@@ -267,6 +303,9 @@ export function upgradeDatabase(db, oldVersion) {
     if (oldVersion < 1) { // The database did not previously exist.
         createDatabase(db);
     }
+    if (oldVersion < 2) {
+        createV2Tables(db);
+    }
     // Expand as needed.
 }
 
@@ -283,6 +322,10 @@ function createDatabase(db) {
     outgoingRoomKeyRequestsStore.createIndex("state", "state");
 }
 
+function createV2Tables(db) {
+    db.createObjectStore("account");
+}
+
 function promiseifyTxn(txn) {
     return new Promise((resolve, reject) => {
         txn.oncomplete = resolve;

src/crypto/store/indexeddb-crypto-store.js

@@ -220,4 +220,20 @@ export default class IndexedDBCryptoStore {
             return backend.deleteOutgoingRoomKeyRequest(requestId, expectedState);
         });
     }
+
+    /**
+     * Load the end to end account for the logged-in user, giving an object
+     * that has the base64 encoded account string and a method for saving
+     * the account string back to the database. This allows the account
+     * to be read and writen atomically.
+     * @return {Promise<Object>} Object
+     * @return {Promise<Object>.account} Base64 encoded account.
+     * @return {Promise<Object>.save} Function to save account data back.
+     *     Takes base64 encoded account data, returns a promise.
+     */
+    endToEndAccountTransaction() {
+        return this._connect().then((backend) => {
+            return backend.endToEndAccountTransaction();
+        });
+    }
 }

src/crypto/store/memory-crypto-store.js

@@ -30,6 +30,7 @@ import utils from '../../utils';
 export default class MemoryCryptoStore {
     constructor() {
         this._outgoingRoomKeyRequests = [];
+        this._account = null;
     }
 
     /**
@@ -196,4 +197,23 @@ export default class MemoryCryptoStore {
 
         return Promise.resolve(null);
     }
+
+    /**
+     * Load the end to end account for the logged-in user, giving an object
+     * that has the base64 encoded account string and a method for saving
+     * the account string back to the database. This allows the account
+     * to be read and writen atomically.
+     * @return {Promise<Object>} Object
+     * @return {Promise<Object>.account} Base64 encoded account.
+     * @return {Promise<Object>.save} Function to save account data back.
+     *     Takes base64 encoded account data, returns a promise.
+     */
+    endToEndAccountTransaction() {
+        return Promise.resolve({
+            account: this._account,
+            save: (newData) => {
+                this._account = newData;
+            },
+        });
+    }
 }