diff --git a/keyserver/src/updaters/olm-session-updater.js b/keyserver/src/updaters/olm-session-updater.js --- a/keyserver/src/updaters/olm-session-updater.js +++ b/keyserver/src/updaters/olm-session-updater.js @@ -1,13 +1,13 @@ // @flow -import type { EncryptResult } from '@commapp/olm'; +import type { EncryptResult, Session as OlmSession } from '@commapp/olm'; import { ServerError } from 'lib/utils/errors.js'; import sleep from 'lib/utils/sleep.js'; import { SQL, dbQuery } from '../database/database.js'; import { fetchPickledOlmAccount } from '../fetchers/olm-account-fetchers.js'; -import { unpickleOlmSession } from '../utils/olm-utils.js'; +import { unpickleSessionAndUseCallback } from '../utils/olm-objects.js'; const maxOlmSessionUpdateAttemptTime = 30000; const olmSessionUpdateRetryDelay = 50; @@ -45,21 +45,28 @@ } const [{ version, pickled_olm_session: pickledSession }] = olmSessionResult; - const session = await unpickleOlmSession(pickledSession, picklingKey); const encryptedMessages: { [string]: EncryptResult } = {}; - for (const messageName in messagesToEncrypt) { - encryptedMessages[messageName] = session.encrypt( - messagesToEncrypt[messageName], - ); - } + const { + pickledOlmSession: { pickledSession: updatedSession }, + } = await unpickleSessionAndUseCallback( + { + picklingKey, + pickledSession, + }, + (olmSession: OlmSession) => { + for (const messageName in messagesToEncrypt) { + encryptedMessages[messageName] = olmSession.encrypt( + messagesToEncrypt[messageName], + ); + } + }, + ); if (dbPersistCondition && !dbPersistCondition(encryptedMessages)) { return { encryptedMessages, dbPersistConditionViolated: true }; } - const updatedSession = session.pickle(picklingKey); - const [transactionResult] = await dbQuery( SQL` START TRANSACTION; diff --git a/keyserver/src/utils/olm-objects.js b/keyserver/src/utils/olm-objects.js --- a/keyserver/src/utils/olm-objects.js +++ b/keyserver/src/utils/olm-objects.js @@ -1,6 +1,9 @@ // @flow -import olm, { type Account as OlmAccount } from '@commapp/olm'; +import olm, { + type Account as OlmAccount, + type Session as OlmSession, +} from '@commapp/olm'; import uuid from 'uuid'; import { ServerError } from 'lib/utils/errors.js'; @@ -56,4 +59,42 @@ }; } -export { unpickleAccountAndUseCallback, createPickledOlmAccount }; +export type PickledOlmSession = { + +picklingKey: string, + +pickledSession: string, +}; +async function unpickleSessionAndUseCallback( + pickledOlmSession: PickledOlmSession, + callback: (session: OlmSession, picklingKey: string) => Promise | T, +): Promise<{ +result: T, +pickledOlmSession: PickledOlmSession }> { + const { picklingKey, pickledSession } = pickledOlmSession; + + await olm.init(); + + const session = new olm.Session(); + session.unpickle(picklingKey, pickledSession); + + let result; + try { + result = await callback(session, picklingKey); + } catch (e) { + throw new ServerError(getMessageForException(e) ?? 'unknown_error'); + } + const updatedSession = session.pickle(picklingKey); + + session.free(); + + return { + result, + pickledOlmSession: { + ...pickledOlmSession, + pickledSession: updatedSession, + }, + }; +} + +export { + unpickleAccountAndUseCallback, + createPickledOlmAccount, + unpickleSessionAndUseCallback, +}; diff --git a/keyserver/src/utils/olm-utils.js b/keyserver/src/utils/olm-utils.js --- a/keyserver/src/utils/olm-utils.js +++ b/keyserver/src/utils/olm-utils.js @@ -1,9 +1,6 @@ // @flow -import type { - Account as OlmAccount, - Session as OlmSession, -} from '@commapp/olm'; +import type { Account as OlmAccount } from '@commapp/olm'; import olm from '@commapp/olm'; import invariant from 'invariant'; @@ -48,16 +45,6 @@ return session.pickle(accountPicklingKey); } -async function unpickleOlmSession( - pickledSession: string, - picklingKey: string, -): Promise { - await olm.init(); - const session = new olm.Session(); - session.unpickle(picklingKey, pickledSession); - return session; -} - async function markPrekeysAsPublished(): Promise { await Promise.all([ fetchCallUpdateOlmAccount('content', (contentAccount: OlmAccount) => { @@ -276,7 +263,6 @@ export { createPickledOlmSession, - unpickleOlmSession, uploadNewOneTimeKeys, getContentSigningKey, validateAndUploadAccountPrekeys,