From af7b2a5dd26202c16bc6308c605e78f7d432273a Mon Sep 17 00:00:00 2001 From: Adhiraj Singh Date: Wed, 20 Apr 2022 13:10:53 +0530 Subject: [PATCH 1/5] feat: limit number of message retries being sent --- src/Socket/messages-recv.ts | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/src/Socket/messages-recv.ts b/src/Socket/messages-recv.ts index 115a44e..7aa0d60 100644 --- a/src/Socket/messages-recv.ts +++ b/src/Socket/messages-recv.ts @@ -273,6 +273,11 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { } } + const willSendMessageAgain = (id: string) => { + const retryCount = msgRetryMap[id] || 0 + return retryCount < 5 + } + const sendMessagesAgain = async(key: proto.IMessageKey, ids: string[]) => { const msgs = await Promise.all( ids.map(id => ( @@ -291,6 +296,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { for(let i = 0; i < msgs.length;i++) { if(msgs[i]) { + msgRetryMap[ids[i]] = (msgRetryMap[ids[i]] || 0) + 1 await relayMessage(key.remoteJid, msgs[i], { messageId: ids[i], participant @@ -360,18 +366,22 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { } if(attrs.type === 'retry') { - // correctly set who is asking for the retry - key.participant = key.participant || attrs.from - if(key.fromMe) { - try { - logger.debug({ attrs, key }, 'recv retry request') - await sendMessagesAgain(key, ids) - } catch(error) { - logger.error({ key, ids, trace: error.stack }, 'error in sending message again') - shouldAck = false + if(willSendMessageAgain(key.id)) { + // correctly set who is asking for the retry + key.participant = key.participant || attrs.from + if(key.fromMe) { + try { + logger.debug({ attrs, key }, 'recv retry request') + await sendMessagesAgain(key, ids) + } catch(error) { + logger.error({ key, ids, trace: error.stack }, 'error in sending message again') + shouldAck = false + } + } else { + logger.info({ attrs, key }, 'recv retry for not fromMe message') } } else { - logger.info({ attrs, key }, 'recv retry for not fromMe message') + logger.info({ attrs, key }, 'will not send message again, as sent too many times') } } From ac17225cf34fbfd8803d5873acbafac29c4e5975 Mon Sep 17 00:00:00 2001 From: Adhiraj Singh Date: Wed, 20 Apr 2022 13:20:05 +0530 Subject: [PATCH 2/5] fix: use correct ID for retry recp --- src/Socket/messages-recv.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Socket/messages-recv.ts b/src/Socket/messages-recv.ts index 7aa0d60..808b497 100644 --- a/src/Socket/messages-recv.ts +++ b/src/Socket/messages-recv.ts @@ -366,7 +366,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { } if(attrs.type === 'retry') { - if(willSendMessageAgain(key.id)) { + if(willSendMessageAgain(ids[0])) { // correctly set who is asking for the retry key.participant = key.participant || attrs.from if(key.fromMe) { From 5655961d12552b676fe05fe1d7d08e2c27cb603d Mon Sep 17 00:00:00 2001 From: Adhiraj Singh Date: Wed, 20 Apr 2022 20:30:11 +0530 Subject: [PATCH 3/5] feat: allow passing of url info in text message --- src/Types/Message.ts | 1 + src/Utils/messages.ts | 23 ++++++++++++++--------- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/src/Types/Message.ts b/src/Types/Message.ts index acbc41a..a83afc7 100644 --- a/src/Types/Message.ts +++ b/src/Types/Message.ts @@ -103,6 +103,7 @@ export type AnyMediaMessageContent = ( export type AnyRegularMessageContent = ( ({ text: string + linkPreview?: WAUrlInfo } & Mentionable & Buttonable & Templatable & Listable) | AnyMediaMessageContent | diff --git a/src/Utils/messages.ts b/src/Utils/messages.ts index 772ee75..32c8e50 100644 --- a/src/Utils/messages.ts +++ b/src/Utils/messages.ts @@ -17,7 +17,7 @@ import { WAMessageContent, WAMessageStatus, WAProto, - WATextMessage + WATextMessage, } from '../Types' import { generateMessageID, unixTimestampSeconds } from './generics' import { downloadContentFromMessage, encryptedStream, generateThumbnail, getAudioDuration, MediaDownloadOptions } from './messages-media' @@ -245,21 +245,26 @@ export const generateWAMessageContent = async( ) => { let m: WAMessageContent = {} if('text' in message) { - const extContent = { ...message } as WATextMessage + const extContent = { text: message.text } as WATextMessage + + let urlInfo = message.linkPreview if(!!options.getUrlInfo && message.text.match(URL_REGEX)) { try { - const data = await options.getUrlInfo(message.text) - extContent.canonicalUrl = data['canonical-url'] - extContent.matchedText = data['matched-text'] - extContent.jpegThumbnail = data.jpegThumbnail - extContent.description = data.description - extContent.title = data.title - extContent.previewType = 0 + urlInfo = await options.getUrlInfo(message.text) } catch(error) { // ignore if fails options.logger?.warn({ trace: error.stack }, 'url generation failed') } } + if(urlInfo) { + extContent.canonicalUrl = urlInfo['canonical-url'] + extContent.matchedText = urlInfo['matched-text'] + extContent.jpegThumbnail = urlInfo.jpegThumbnail + extContent.description = urlInfo.description + extContent.title = urlInfo.title + extContent.previewType = 0 + } + m.extendedTextMessage = extContent } else if('contacts' in message) { const contactLen = message.contacts.contacts.length From d5e46b784bcd0496d7e17d12a998f55f3ffdd436 Mon Sep 17 00:00:00 2001 From: Adhiraj Singh Date: Wed, 20 Apr 2022 20:31:49 +0530 Subject: [PATCH 4/5] fix: only generate urlInfo if not provided --- src/Utils/messages.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Utils/messages.ts b/src/Utils/messages.ts index 32c8e50..6a238e6 100644 --- a/src/Utils/messages.ts +++ b/src/Utils/messages.ts @@ -248,7 +248,7 @@ export const generateWAMessageContent = async( const extContent = { text: message.text } as WATextMessage let urlInfo = message.linkPreview - if(!!options.getUrlInfo && message.text.match(URL_REGEX)) { + if(!urlInfo && !!options.getUrlInfo && message.text.match(URL_REGEX)) { try { urlInfo = await options.getUrlInfo(message.text) } catch(error) { // ignore if fails From a98484c38bc92055e3c660a86687ee0425bc5c5b Mon Sep 17 00:00:00 2001 From: Adhiraj Singh Date: Thu, 21 Apr 2022 00:52:17 +0530 Subject: [PATCH 5/5] refactor: split "downloadContentFromMessage" to "downloadEncryptedContent"" --- src/Types/Message.ts | 6 ++++++ src/Utils/messages-media.ts | 24 +++++++++++++++++++----- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/src/Types/Message.ts b/src/Types/Message.ts index a83afc7..bc55d89 100644 --- a/src/Types/Message.ts +++ b/src/Types/Message.ts @@ -178,3 +178,9 @@ export type WAMessageUpdate = { update: Partial, key: proto.IMessageK export type WAMessageCursor = { before: WAMessageKey | undefined } | { after: WAMessageKey | undefined } export type MessageUserReceiptUpdate = { key: proto.IMessageKey, receipt: MessageUserReceipt } + +export type MediaDecryptionKeyInfo = { + iv: Buffer + cipherKey: Buffer + macKey?: Buffer +} \ No newline at end of file diff --git a/src/Utils/messages-media.ts b/src/Utils/messages-media.ts index 48bf4cc..03b7ccb 100644 --- a/src/Utils/messages-media.ts +++ b/src/Utils/messages-media.ts @@ -11,7 +11,7 @@ import type { Logger } from 'pino' import { Readable, Transform } from 'stream' import { URL } from 'url' import { DEFAULT_ORIGIN, MEDIA_PATH_MAP } from '../Defaults' -import { CommonSocketConfig, DownloadableMessage, MediaConnInfo, MediaType, MessageType, WAGenericMediaMessage, WAMediaUpload, WAMediaUploadFunction, WAMessageContent, WAProto } from '../Types' +import { CommonSocketConfig, DownloadableMessage, MediaConnInfo, MediaDecryptionKeyInfo, MediaType, MessageType, WAGenericMediaMessage, WAMediaUpload, WAMediaUploadFunction, WAMessageContent, WAProto } from '../Types' import { hkdf } from './crypto' import { generateMessageID } from './generics' @@ -60,7 +60,7 @@ export const hkdfInfoKey = (type: MediaType) => { } /** generates all the keys required to encrypt/decrypt & sign a media message */ -export function getMediaKeys(buffer, mediaType: MediaType) { +export function getMediaKeys(buffer: Uint8Array | string, mediaType: MediaType): MediaDecryptionKeyInfo { if(typeof buffer === 'string') { buffer = Buffer.from(buffer.replace('data:;base64,', ''), 'base64') } @@ -186,6 +186,7 @@ export const toBuffer = async(stream: Readable) => { buff = Buffer.concat([ buff, chunk ]) } + stream.destroy() return buff } @@ -343,12 +344,26 @@ export type MediaDownloadOptions = { endByte?: number } -export const downloadContentFromMessage = async( +export const downloadContentFromMessage = ( { mediaKey, directPath, url }: DownloadableMessage, type: MediaType, - { startByte, endByte }: MediaDownloadOptions = { } + opts: MediaDownloadOptions = { } ) => { const downloadUrl = url || `https://${DEF_HOST}${directPath}` + const keys = getMediaKeys(mediaKey, type) + + return downloadEncryptedContent(downloadUrl, keys, opts) +} + +/** + * Decrypts and downloads an AES256-CBC encrypted file given the keys. + * Assumes the SHA256 of the plaintext is appended to the end of the ciphertext + * */ +export const downloadEncryptedContent = async( + downloadUrl: string, + { cipherKey, iv }: MediaDecryptionKeyInfo, + { startByte, endByte }: MediaDownloadOptions = { } +) => { let bytesFetched = 0 let startChunk = 0 let firstBlockIsIV = false @@ -386,7 +401,6 @@ export const downloadContentFromMessage = async( ) let remainingBytes = Buffer.from([]) - const { cipherKey, iv } = getMediaKeys(mediaKey, type) let aes: Crypto.Decipher