Skip to content
119 changes: 112 additions & 7 deletions app/lib/encryption/encryption.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
// Bypass the global mock of `app/lib/encryption` declared in jest.setup.js by
// importing directly from the file. We exercise the real `encryptMessage` here.
import encryption from './encryption';
import database from '../database';

jest.unmock('./encryption');

Expand Down Expand Up @@ -48,14 +49,38 @@ jest.mock('../store/auxStore', () => ({
}));

const mockSubFind = jest.fn();
jest.mock('../database', () => ({
__esModule: true,
default: {
active: {
get: () => ({ find: (rid: string) => mockSubFind(rid) })
// Rows returned by `collection.query(...).fetch()`, keyed by collection name.
const mockQueryRows: Record<string, unknown[]> = {};
const mockDbBatch = jest.fn((...args: any[]) => {
// db.batch commits prepared records, clearing their pending state (like the real writer).
args.flat().forEach((item: any) => {
if (item && typeof item === 'object' && '_preparedState' in item) {
item._preparedState = null;
}
}
}));
});
return Promise.resolve(undefined);
});
jest.mock('../database', () => {
let writerQueue: Promise<unknown> = Promise.resolve();
return {
__esModule: true,
default: {
active: {
get: (name: string) => ({
find: (rid: string) => mockSubFind(rid),
query: () => ({ fetch: () => Promise.resolve(mockQueryRows[name] ?? []) })
}),
// Serialized writer lock, like WatermelonDB's.
write: (callback: () => Promise<void>) => {
const run = writerQueue.then(() => callback());
writerQueue = run.catch(() => undefined);
return run;
},
batch: (...args: unknown[]) => mockDbBatch(...args)
}
}
};
});

const mockRoomEncrypt = jest.fn();
const mockHasSessionKey = jest.fn();
Expand Down Expand Up @@ -141,3 +166,83 @@ describe('Encryption.encryptMessage', () => {
expect(mockRoomEncrypt).not.toHaveBeenCalled();
});
});

describe('Encryption.decryptPendingMessages', () => {
const rid = 'r1';

// Mimics a WatermelonDB Model: prepareUpdate throws while a previous prepared
// update has not been committed yet.
const makeMessageRecord = (id: string) => {
const record: any = {
id,
t: 'e2e',
msg: 'cipher',
subscription: { id: rid },
_preparedState: null as string | null,
prepareUpdate(recordUpdater: (m: any) => void) {
if (record._preparedState) {
throw new Error(`Cannot update a record with pending changes (messages#${id})`);
}
recordUpdater(record);
record._preparedState = 'update';
return record;
}
};
return record;
};

const deferred = () => {
let resolve: () => void = () => undefined;
const promise = new Promise<void>(r => {
resolve = r;
});
return { promise, resolve };
};

beforeEach(() => {
jest.clearAllMocks();
mockQueryRows.messages = [];
mockQueryRows.threads = [];
mockQueryRows.thread_messages = [];
});

afterEach(() => {
jest.restoreAllMocks();
});

it('does not throw "pending changes" when a concurrent writer touches the same message mid-decrypt', async () => {
const record = makeMessageRecord('m1');
mockQueryRows.messages = [record];
jest.spyOn(encryption, 'decryptMessage').mockResolvedValue({ msg: 'plain', e2e: 'done' } as any);

const db = (database as any).active;

// Hold the writer lock — as a new incoming message being saved would — and then update
// the very record decryptPendingMessages is about to prepare.
const concurrentGate = deferred();
const concurrentWrite = db.write(async () => {
await concurrentGate.promise;
await db.batch([
record.prepareUpdate((m: any) => {
m.msg = 'written by another writer';
})
]);
});

const decrypting = encryption.decryptPendingMessages(rid);

// Give an unlocked implementation the chance to prepare now — before the concurrent
// writer runs — and hold the record pending until its own batch.
await new Promise(resolve => setImmediate(resolve));
concurrentGate.resolve();

await expect(Promise.all([concurrentWrite, decrypting])).resolves.toBeDefined();

// The decrypted update reached db.batch and nothing was left prepared-but-uncommitted.
const committed = mockDbBatch.mock.calls.map(call => call.flat()).some(items => items.includes(record));
expect(committed).toBe(true);
expect(record.msg).toBe('plain');
expect(record.e2e).toBe('done');
expect(record._preparedState).toBeNull();
Comment thread
OtavioStasiak marked this conversation as resolved.
});
});
32 changes: 19 additions & 13 deletions app/lib/encryption/encryption.ts
Original file line number Diff line number Diff line change
Expand Up @@ -335,12 +335,13 @@ class Encryption {
const threadMessagesToDecrypt = await threadMessagesCollection.query(...whereClause).fetch();

// Concat messages/threads/threadMessages
let toDecrypt: (TThreadModel | TThreadMessageModel | TMessageModel)[] = [
const toDecrypt: (TThreadModel | TThreadMessageModel | TMessageModel)[] = [
...messagesToDecrypt,
...threadsToDecrypt,
...threadMessagesToDecrypt
];
toDecrypt = (await Promise.all(

const decrypted = await Promise.all(
toDecrypt.map(async message => {
const { t, msg, tmsg, attachments, content } = message;
let newMessage: Partial<TMessageModel> = {};
Comment thread
OtavioStasiak marked this conversation as resolved.
Expand All @@ -357,20 +358,25 @@ class Encryption {
} as IMessage);
}

try {
return message.prepareUpdate(
protectedFunction((m: TMessageModel) => {
Object.assign(m, newMessage);
})
);
} catch {
return null;
}
return { message, newMessage };
})
)) as (TThreadModel | TThreadMessageModel)[];
);

await db.write(async () => {
Comment thread
OtavioStasiak marked this conversation as resolved.
await db.batch(toDecrypt);
const prepared = decrypted
.map(({ message, newMessage }) => {
try {
return message.prepareUpdate(
protectedFunction((m: TMessageModel) => {
Object.assign(m, newMessage);
})
);
} catch {
return null;
}
Comment thread
OtavioStasiak marked this conversation as resolved.
Outdated
})
.filter((record): record is TThreadModel | TThreadMessageModel => record !== null);
Comment thread
OtavioStasiak marked this conversation as resolved.
Outdated
await db.batch(prepared);
});
} catch (e) {
log(e);
Expand Down
Loading