diff --git a/ts/session/sending/MessageSender.ts b/ts/session/sending/MessageSender.ts index 78d5763c3..58c070d72 100644 --- a/ts/session/sending/MessageSender.ts +++ b/ts/session/sending/MessageSender.ts @@ -6,6 +6,7 @@ import { SignalService } from '../../protobuf'; import { UserUtil } from '../../util'; import { MessageEncrypter } from '../crypto'; import { lokiMessageAPI, lokiPublicChatAPI } from '../../window'; +import pRetry from 'p-retry'; // ================ Regular ================ @@ -14,13 +15,10 @@ export function canSendToSnode(): boolean { return Boolean(lokiMessageAPI); } -export async function send({ - device, - plainTextBuffer, - encryption, - timestamp, - ttl, -}: RawMessage): Promise { +export async function send( + { device, plainTextBuffer, encryption, timestamp, ttl }: RawMessage, + retries: number = 3 +): Promise { if (!canSendToSnode()) { throw new Error('lokiMessageAPI is not initialized.'); } @@ -33,8 +31,14 @@ export async function send({ const envelope = await buildEnvelope(envelopeType, timestamp, cipherText); const data = wrapEnvelope(envelope); - // TODO: Somehow differentiate between Retryable and Regular erros - return lokiMessageAPI.sendMessage(device, data, timestamp, ttl); + // pRetry doesn't count the first call as a retry + return pRetry( + async () => lokiMessageAPI.sendMessage(device, data, timestamp, ttl), + { + retries: retries - 1, + factor: 1, + } + ); } async function buildEnvelope( diff --git a/ts/test/session/sending/MessageSender_test.ts b/ts/test/session/sending/MessageSender_test.ts index 67da38618..8a6525ec8 100644 --- a/ts/test/session/sending/MessageSender_test.ts +++ b/ts/test/session/sending/MessageSender_test.ts @@ -21,11 +21,25 @@ describe('MessageSender', () => { TestUtils.restoreStubs(); }); + describe('canSendToSnode', () => { + it('should return the correct value', () => { + const stub = TestUtils.stubWindow('lokiMessageAPI', undefined); + expect(MessageSender.canSendToSnode()).to.equal( + false, + 'We cannot send if lokiMessageAPI is not set' + ); + stub.set(sandbox.createStubInstance(LokiMessageAPI)); + expect(MessageSender.canSendToSnode()).to.equal( + true, + 'We can send if lokiMessageAPI is set' + ); + }); + }); + describe('send', () => { const ourNumber = 'ourNumber'; let lokiMessageAPIStub: sinon.SinonStubbedInstance; - let messageEncyrptReturnEnvelopeType = - SignalService.Envelope.Type.CIPHERTEXT; + let encryptStub: sinon.SinonStub<[string, Uint8Array, EncryptionType]>; beforeEach(() => { // We can do this because LokiMessageAPI has a module export in it @@ -34,76 +48,87 @@ describe('MessageSender', () => { }); TestUtils.stubWindow('lokiMessageAPI', lokiMessageAPIStub); + encryptStub = sandbox.stub(MessageEncrypter, 'encrypt').resolves({ + envelopeType: SignalService.Envelope.Type.CIPHERTEXT, + cipherText: crypto.randomBytes(10), + }); + sandbox.stub(UserUtil, 'getCurrentDevicePubKey').resolves(ourNumber); - sandbox - .stub(MessageEncrypter, 'encrypt') - .callsFake(async (_device, plainTextBuffer, _type) => ({ - envelopeType: messageEncyrptReturnEnvelopeType, - cipherText: plainTextBuffer, - })); }); - it('should pass the correct values to lokiMessageAPI', async () => { - const device = '0'; - const timestamp = Date.now(); - const ttl = 100; - - await MessageSender.send({ + describe('retry', () => { + const rawMessage = { identifier: '1', - device, + device: '0', plainTextBuffer: crypto.randomBytes(10), encryption: EncryptionType.Signal, - timestamp, - ttl, + timestamp: Date.now(), + ttl: 100, + }; + + it('should not retry if an error occurred during encryption', async () => { + encryptStub.throws(new Error('Failed to encrypt.')); + const promise = MessageSender.send(rawMessage); + await expect(promise).is.rejectedWith('Failed to encrypt.'); + expect(lokiMessageAPIStub.sendMessage.callCount).to.equal(0); }); - const args = lokiMessageAPIStub.sendMessage.getCall(0).args; - expect(args[0]).to.equal(device); - expect(args[2]).to.equal(timestamp); - expect(args[3]).to.equal(ttl); - }); + it('should only call lokiMessageAPI once if no errors occured', async () => { + await MessageSender.send(rawMessage); + expect(lokiMessageAPIStub.sendMessage.callCount).to.equal(1); + }); - it('should correctly build the envelope', async () => { - messageEncyrptReturnEnvelopeType = SignalService.Envelope.Type.CIPHERTEXT; + it('should only retry the specified amount of times before throwing', async () => { + lokiMessageAPIStub.sendMessage.throws(new Error('API error')); + const retries = 2; + const promise = MessageSender.send(rawMessage, retries); + await expect(promise).is.rejectedWith('API error'); + expect(lokiMessageAPIStub.sendMessage.callCount).to.equal(retries); + }); - // This test assumes the encryption stub returns the plainText passed into it. - const plainTextBuffer = crypto.randomBytes(10); - const timestamp = Date.now(); + it('should not throw error if successful send occurs within the retry limit', async () => { + lokiMessageAPIStub.sendMessage + .onFirstCall() + .throws(new Error('API error')); + await MessageSender.send(rawMessage, 3); + expect(lokiMessageAPIStub.sendMessage.callCount).to.equal(2); + }); + }); - await MessageSender.send({ - identifier: '1', - device: '0', - plainTextBuffer, - encryption: EncryptionType.Signal, - timestamp, - ttl: 1, + describe('logic', () => { + let messageEncyrptReturnEnvelopeType = + SignalService.Envelope.Type.CIPHERTEXT; + + beforeEach(() => { + encryptStub.callsFake(async (_device, plainTextBuffer, _type) => ({ + envelopeType: messageEncyrptReturnEnvelopeType, + cipherText: plainTextBuffer, + })); }); - const data = lokiMessageAPIStub.sendMessage.getCall(0).args[1]; - const webSocketMessage = SignalService.WebSocketMessage.decode(data); - expect(webSocketMessage.request?.body).to.not.equal( - undefined, - 'Request body should not be undefined' - ); - expect(webSocketMessage.request?.body).to.not.equal( - null, - 'Request body should not be null' - ); + it('should pass the correct values to lokiMessageAPI', async () => { + const device = '0'; + const timestamp = Date.now(); + const ttl = 100; - const envelope = SignalService.Envelope.decode( - webSocketMessage.request?.body as Uint8Array - ); - expect(envelope.type).to.equal(SignalService.Envelope.Type.CIPHERTEXT); - expect(envelope.source).to.equal(ourNumber); - expect(envelope.sourceDevice).to.equal(1); - expect(toNumber(envelope.timestamp)).to.equal(timestamp); - expect(envelope.content).to.deep.equal(plainTextBuffer); - }); + await MessageSender.send({ + identifier: '1', + device, + plainTextBuffer: crypto.randomBytes(10), + encryption: EncryptionType.Signal, + timestamp, + ttl, + }); + + const args = lokiMessageAPIStub.sendMessage.getCall(0).args; + expect(args[0]).to.equal(device); + expect(args[2]).to.equal(timestamp); + expect(args[3]).to.equal(ttl); + }); - describe('UNIDENTIFIED_SENDER', () => { - it('should set the envelope source to be empty', async () => { + it('should correctly build the envelope', async () => { messageEncyrptReturnEnvelopeType = - SignalService.Envelope.Type.UNIDENTIFIED_SENDER; + SignalService.Envelope.Type.CIPHERTEXT; // This test assumes the encryption stub returns the plainText passed into it. const plainTextBuffer = crypto.randomBytes(10); @@ -132,13 +157,53 @@ describe('MessageSender', () => { const envelope = SignalService.Envelope.decode( webSocketMessage.request?.body as Uint8Array ); - expect(envelope.type).to.equal( - SignalService.Envelope.Type.UNIDENTIFIED_SENDER - ); - expect(envelope.source).to.equal( - '', - 'envelope source should be empty in UNIDENTIFIED_SENDER' - ); + expect(envelope.type).to.equal(SignalService.Envelope.Type.CIPHERTEXT); + expect(envelope.source).to.equal(ourNumber); + expect(envelope.sourceDevice).to.equal(1); + expect(toNumber(envelope.timestamp)).to.equal(timestamp); + expect(envelope.content).to.deep.equal(plainTextBuffer); + }); + + describe('UNIDENTIFIED_SENDER', () => { + it('should set the envelope source to be empty', async () => { + messageEncyrptReturnEnvelopeType = + SignalService.Envelope.Type.UNIDENTIFIED_SENDER; + + // This test assumes the encryption stub returns the plainText passed into it. + const plainTextBuffer = crypto.randomBytes(10); + const timestamp = Date.now(); + + await MessageSender.send({ + identifier: '1', + device: '0', + plainTextBuffer, + encryption: EncryptionType.Signal, + timestamp, + ttl: 1, + }); + + const data = lokiMessageAPIStub.sendMessage.getCall(0).args[1]; + const webSocketMessage = SignalService.WebSocketMessage.decode(data); + expect(webSocketMessage.request?.body).to.not.equal( + undefined, + 'Request body should not be undefined' + ); + expect(webSocketMessage.request?.body).to.not.equal( + null, + 'Request body should not be null' + ); + + const envelope = SignalService.Envelope.decode( + webSocketMessage.request?.body as Uint8Array + ); + expect(envelope.type).to.equal( + SignalService.Envelope.Type.UNIDENTIFIED_SENDER + ); + expect(envelope.source).to.equal( + '', + 'envelope source should be empty in UNIDENTIFIED_SENDER' + ); + }); }); }); });