fix: fail E2EE sends without a session key and stop key distribution on logout

app/lib/encryption/encryption.test.ts

@@ -2,6 +2,7 @@
 // importing directly from the file. We exercise the real `encryptMessage` here.
 import encryption from './encryption';
 import database from '../database';
+import Deferred from './helpers/deferred';
 
 jest.unmock('./encryption');
 
@@ -232,7 +233,9 @@ describe('Encryption.canSendToRoom', () => {
 		jest.clearAllMocks();
 		mockHandshake.mockResolvedValue(undefined);
 		(encryption as any).roomInstances = {};
-		(encryption as any).readyPromise = Promise.resolve();
+		const readyPromise = new Deferred();
+		readyPromise.resolve();
+		(encryption as any).readyPromise = readyPromise;
 	});
 
 	it('allows rooms that are not encrypted', async () => {
@@ -422,3 +425,59 @@ describe('Encryption.decryptPendingSubscriptions', () => {
 		expect(failing.lastMessage).toEqual({ t: 'e2e', e2e: 'pending', msg: 'cipher' });
 	});
 });
+
+describe('Encryption key distribution', () => {
+	const db = (database as any).active;
+	const KEY_DISTRIBUTION_INTERVAL = 10000;
+	const distributionRuns = () => jest.mocked(db.get).mock.calls.filter(([name]: [string]) => name === 'subscriptions').length;
+
+	beforeEach(() => {
+		jest.useFakeTimers();
+		jest.spyOn(db, 'get').mockImplementation(() => ({ query: () => Promise.resolve([]) }));
+	});
+
+	afterEach(() => {
+		encryption.stop();
+		jest.useRealTimers();
+		jest.restoreAllMocks();
+	});
+
+	it('stops distributing room keys when the client stops', async () => {
+		await encryption.initiateKeyDistribution();
+		encryption.stop();
+
+		await jest.advanceTimersByTimeAsync(KEY_DISTRIBUTION_INTERVAL);
+
+		expect(distributionRuns()).toBe(1);
+	});
+
+	it('stops distributing room keys when the client stops during the first run', async () => {
+		const firstRun = deferredSubscriptions();
+		jest.mocked(db.get).mockImplementationOnce(() => ({ query: () => firstRun.promise }));
+
+		const initiating = encryption.initiateKeyDistribution();
+		encryption.stop();
+		firstRun.resolve([]);
+		await initiating;
+		await jest.advanceTimersByTimeAsync(KEY_DISTRIBUTION_INTERVAL);
+
+		expect(distributionRuns()).toBe(1);
+	});
+
+	it('distributes room keys right away when the client starts again after a stop', async () => {
+		await encryption.initiateKeyDistribution();
+		encryption.stop();
+
+		await encryption.initiateKeyDistribution();
+
+		expect(distributionRuns()).toBe(2);
+	});
+});
+
+const deferredSubscriptions = () => {
+	let resolve: (subscriptions: unknown[]) => void = () => undefined;
+	const promise = new Promise<unknown[]>(r => {
+		resolve = r;
+	});
+	return { promise, resolve };
+};

app/lib/encryption/encryption.ts

@@ -128,6 +128,10 @@ class Encryption {
 		this.privateKey = null;
 		this.publicKey = null;
 		this.roomInstances = {};
+		if (this.keyDistributionInterval) {
+			clearInterval(this.keyDistributionInterval);
+		}
+		this.keyDistributionInterval = null;
 		// Cancel ongoing encryption/decryption requests
 		this.readyPromise.reject();
 		// Reset Deferred
@@ -530,8 +534,8 @@ class Encryption {
 			}
 		};
 
-		await keyDistribution();
 		this.keyDistributionInterval = setInterval(keyDistribution, 10000);
+		await keyDistribution();
 	};
 
 	// Creating the instance is enough to generate room e2ee key

app/lib/methods/sendMessage.test.ts

@@ -423,6 +423,15 @@ describe('sendMessage', () => {
 			expect(mockDbBatch).not.toHaveBeenCalled();
 			expect(mockPost).not.toHaveBeenCalled();
 		});
+
+		it('stores the message as errored and posts nothing when the encrypted room has no session key', async () => {
+			mockEncryptMessage.mockResolvedValue(undefined);
+
+			await sendMessage(rid, 'hello', undefined, user);
+
+			expect(createdRecord('messages').status).toBe(messagesStatus.ERROR);
+			expect(mockPost).not.toHaveBeenCalled();
+		});
 	});
 
 	describe('sendMessage - draft', () => {
@@ -661,5 +670,17 @@ describe('sendMessage', () => {
 			expect(log).toHaveBeenCalledWith(expect.objectContaining({ message: 'encryption failed' }));
 			expect(mockPost).not.toHaveBeenCalled();
 		});
+
+		it('marks the message as errored again and posts nothing when the encrypted room has no session key', async () => {
+			mockEncryptMessage.mockResolvedValue(undefined);
+			const message = makeRecord('messages#messageId', { id: 'messageId', status: messagesStatus.ERROR });
+			mockGetCollection('messages').records.set('messageId', message);
+			message.update = jest.fn((updater: (m: FakeRecord) => void) => Promise.resolve(updater(message)));
+
+			await resendMessage(message as any);
+
+			expect(message.status).toBe(messagesStatus.ERROR);
+			expect(mockPost).not.toHaveBeenCalled();
+		});
 	});
 });

app/lib/methods/sendMessage.ts

@@ -84,6 +84,10 @@ export async function resendMessage(message: TMessageModel, tmid?: string) {
 			msg: message.msg,
 			...(tmid && { tmid })
 		} as IMessage);
+		if (!m) {
+			await changeMessageStatus(message.id, messagesStatus.ERROR, tmid);
+			return;
+		}
 
 		await sendMessageCall(m);
 	} catch (e) {
@@ -116,6 +120,7 @@ export async function sendMessage(
 		} as IMessage);
 
 		const messageDate = new Date();
+		const status = message ? messagesStatus.TEMP : messagesStatus.ERROR;
 
 		try {
 			await db.write(async () => {
@@ -173,7 +178,7 @@ export async function sendMessage(
 								tm.msg = msg;
 								tm.ts = messageDate;
 								tm._updatedAt = messageDate;
-								tm.status = messagesStatus.TEMP;
+								tm.status = status;
 								tm.u = {
 									_id: user.id || '1',
 									username: user.username,
@@ -200,7 +205,7 @@ export async function sendMessage(
 						m.msg = msg;
 						m.ts = messageDate;
 						m._updatedAt = messageDate;
-						m.status = messagesStatus.TEMP;
+						m.status = status;
 						m.u = {
 							_id: user.id || '1',
 							username: user.username,
@@ -239,7 +244,9 @@ export async function sendMessage(
 			return;
 		}
 
-		await sendMessageCall(message);
+		if (message) {
+			await sendMessageCall(message);
+		}
 	} catch (e) {
 		log(e);
 	}