refactor: split downloadMediaMessage into functional components

src/LegacySocket/messages.ts

@@ -2,7 +2,7 @@ import { Boom } from '@hapi/boom'
 import { proto } from '../../WAProto'
 import { WA_DEFAULT_EPHEMERAL } from '../Defaults'
 import { AnyMessageContent, Chat, GroupMetadata, LegacySocketConfig, MediaConnInfo, MessageUpdateType, MessageUserReceipt, MessageUserReceiptUpdate, MiscMessageGenerationOptions, ParticipantAction, WAFlag, WAMessage, WAMessageCursor, WAMessageKey, WAMessageStatus, WAMessageStubType, WAMessageUpdate, WAMetric, WAUrlInfo } from '../Types'
-import { decryptMediaMessageBuffer, extractMessageContent, generateWAMessage, getWAUploadToServer, toNumber } from '../Utils'
+import { downloadMediaMessage, generateWAMessage, getWAUploadToServer, MediaDownloadOptions, toNumber } from '../Utils'
 import { areJidsSameUser, BinaryNode, getBinaryNodeMessages, isJidGroup, jidNormalizedUser } from '../WABinary'
 import makeChatsSocket from './chats'
 
@@ -446,28 +446,9 @@ const makeMessagesSocket = (config: LegacySocketConfig) => {
 
 			return Object.values(info)
 		},
-		downloadMediaMessage: async(message: WAMessage, type: 'buffer' | 'stream' = 'buffer') => {
+		downloadMediaMessage: async(message: WAMessage, type: 'buffer' | 'stream' = 'buffer', options: MediaDownloadOptions = { }) => {
-			const downloadMediaMessage = async() => {
-				const mContent = extractMessageContent(message.message)
-				if(!mContent) {
-					throw new Boom('No message present', { statusCode: 400, data: message })
-				}
-
-				const stream = await decryptMediaMessageBuffer(mContent)
-				if(type === 'buffer') {
-					let buffer = Buffer.from([])
-					for await (const chunk of stream) {
-						buffer = Buffer.concat([buffer, chunk])
-					}
-
-					return buffer
-				}
-
-				return stream
-			}
-
 			try {
-				const result = await downloadMediaMessage()
+				const result = await downloadMediaMessage(message, type, options)
 				return result
 			} catch(error) {
 				if(error.message.includes('404')) { // media needs to be updated
@@ -475,7 +456,7 @@ const makeMessagesSocket = (config: LegacySocketConfig) => {
 
 					await updateMediaMessage(message)
 
-					const result = await downloadMediaMessage()
+					const result = await downloadMediaMessage(message, type, options)
 					return result
 				}
 

src/Utils/messages-media.ts

@@ -338,7 +338,7 @@ const toSmallestChunkSize = (num: number) => {
 	return Math.floor(num / AES_CHUNK_SIZE) * AES_CHUNK_SIZE
 }
 
-type MediaDownloadOptions = {
+export type MediaDownloadOptions = {
     startByte?: number
     endByte?: number
 }
@@ -446,47 +446,6 @@ export const downloadContentFromMessage = async(
 	return fetched.pipe(output, { end: true })
 }
 
-/**
- * Decode a media message (video, image, document, audio) & return decrypted buffer
- * @param message the media message you want to decode
- */
-export async function decryptMediaMessageBuffer(message: WAMessageContent): Promise<Readable> {
-	/*
-        One can infer media type from the key in the message
-        it is usually written as [mediaType]Message. Eg. imageMessage, audioMessage etc.
-    */
-	const type = Object.keys(message)[0] as MessageType
-	if(
-		!type ||
-		type === 'conversation' ||
-		type === 'extendedTextMessage'
-	) {
-		throw new Boom(`no media message for "${type}"`, { statusCode: 400 })
-	}
-
-	if(type === 'locationMessage' || type === 'liveLocationMessage') {
-		const buffer = Buffer.from(message[type].jpegThumbnail)
-		const readable = new Readable({ read: () => {} })
-		readable.push(buffer)
-		readable.push(null)
-		return readable
-	}
-
-	let messageContent: WAGenericMediaMessage
-	if(message.productMessage) {
-		const product = message.productMessage.product?.productImage
-		if(!product) {
-			throw new Boom('product has no image', { statusCode: 400 })
-		}
-
-		messageContent = product
-	} else {
-		messageContent = message[type]
-	}
-
-	return downloadContentFromMessage(messageContent, type.replace('Message', '') as MediaType)
-}
-
 export function extensionForMediaMessage(message: WAMessageContent) {
 	const getExtension = (mimetype: string) => mimetype.split(';')[0].split('/')[1]
 	const type = Object.keys(message)[0] as MessageType

src/Utils/messages.ts

@@ -21,6 +21,7 @@ import {
 } from '../Types'
 import { generateMessageID, unixTimestampSeconds } from './generics'
 import { encryptedStream, generateThumbnail, getAudioDuration } from './messages-media'
+import { downloadContentFromMessage, MediaDownloadOptions } from '.'
 
 type MediaUploadData = {
 	media: WAMediaUpload
@@ -535,3 +536,32 @@ export const updateMessageWithReceipt = (msg: WAMessage, receipt: MessageUserRec
 		msg.userReceipt.push(receipt)
 	}
 }
+
+/**
+ * Downloads the given message. Throws an error if it's not a media message
+ */
+export const downloadMediaMessage = async(message: WAMessage, type: 'buffer' | 'stream', options: MediaDownloadOptions) => {
+	const mContent = extractMessageContent(message.message)
+	if(!mContent) {
+		throw new Boom('No message present', { statusCode: 400, data: message })
+	}
+
+	const contentType = getContentType(mContent)
+	const mediaType = contentType.replace('Message', '') as MediaType
+	const media = mContent[contentType]
+	if(typeof media !== 'object' || !('url' in media)) {
+		throw new Boom(`"${contentType}" message is not a media message`)
+	}
+
+	const stream = await downloadContentFromMessage(media, mediaType, options)
+	if(type === 'buffer') {
+		let buffer = Buffer.from([])
+		for await (const chunk of stream) {
+			buffer = Buffer.concat([buffer, chunk])
+		}
+
+		return buffer
+	}
+
+	return stream
+}