feat: add media retry for MD

src/Socket/messages-recv.ts

@@ -2,7 +2,7 @@
 import { proto } from '../../WAProto'
 import { KEY_BUNDLE_TYPE, MIN_PREKEY_COUNT } from '../Defaults'
 import { BaileysEventMap, MessageReceiptType, MessageUserReceipt, SocketConfig, WACallEvent, WAMessageStubType } from '../Types'
-import { debouncedTimeout, decodeMessageStanza, delay, encodeBigEndian, generateSignalPubKey, getCallStatusFromNode, getNextPreKeys, getStatusFromReceiptType, normalizeMessageContent, unixTimestampSeconds, xmppPreKey, xmppSignedPreKey } from '../Utils'
+import { debouncedTimeout, decodeMediaRetryNode, decodeMessageStanza, delay, encodeBigEndian, generateSignalPubKey, getCallStatusFromNode, getNextPreKeys, getStatusFromReceiptType, normalizeMessageContent, unixTimestampSeconds, xmppPreKey, xmppSignedPreKey } from '../Utils'
 import { makeKeyedMutex, makeMutex } from '../Utils/make-mutex'
 import processMessage, { cleanMessage } from '../Utils/process-message'
 import { areJidsSameUser, BinaryNode, BinaryNodeAttributes, getAllBinaryNodeChildren, getBinaryNodeChild, getBinaryNodeChildren, isJidGroup, isJidUser, jidDecode, jidEncode, jidNormalizedUser, S_WHATSAPP_NET } from '../WABinary'
@@ -171,12 +171,13 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 		const normalizedContent = !!msg.message ? normalizeMessageContent(msg.message) : undefined
 		const isAnyHistoryMsg = !!normalizedContent?.protocolMessage?.historySyncNotification
 		if(isAnyHistoryMsg) {
-			const jid = jidEncode(jidDecode(msg.key.remoteJid!).user, 'c.us')
-			await sendReceipt(jid, undefined, [msg.key.id], 'hist_sync')
 			// we only want to sync app state once we've all the history
 			// restart the app state sync timeout
 			logger.debug('restarting app sync timeout')
 			appStateSyncTimeout.start()
+
+			const jid = jidEncode(jidDecode(msg.key.remoteJid!).user, 'c.us')
+			await sendReceipt(jid, undefined, [msg.key.id], 'hist_sync')
 		}
 
 		return newEvents
@@ -208,8 +209,9 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 	const processNotification = async(node: BinaryNode): Promise<Partial<proto.IWebMessageInfo>> => {
 		const result: Partial<proto.IWebMessageInfo> = { }
 		const [child] = getAllBinaryNodeChildren(node)
+		const nodeType = node.attrs.type
 
-		if(node.attrs.type === 'w:gp2') {
+		if(nodeType === 'w:gp2') {
 			switch (child?.tag) {
 			case 'create':
 				const metadata = extractGroupMetadata(child)
@@ -271,6 +273,9 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 				break
 
 			}
+		} else if(nodeType === 'mediaretry') {
+			const event = decodeMediaRetryNode(node)
+			ev.emit('messages.media-update', [event])
 		} else {
 			switch (child.tag) {
 			case 'devices':

src/Socket/messages-send.ts

@@ -3,7 +3,7 @@ import NodeCache from 'node-cache'
 import { proto } from '../../WAProto'
 import { WA_DEFAULT_EPHEMERAL } from '../Defaults'
 import { AnyMessageContent, MediaConnInfo, MessageReceiptType, MessageRelayOptions, MiscMessageGenerationOptions, SocketConfig, WAMessageKey } from '../Types'
-import { aggregateMessageKeysNotFromMe, encodeWAMessage, encryptSenderKeyMsgSignalProto, encryptSignalProto, extractDeviceJids, generateMessageID, generateWAMessage, getWAUploadToServer, jidToSignalProtocolAddress, parseAndInjectE2ESessions, unixTimestampSeconds } from '../Utils'
+import { aggregateMessageKeysNotFromMe, assertMediaContent, bindWaitForEvent, decryptMediaRetryData, encodeWAMessage, encryptMediaRetryRequest, encryptSenderKeyMsgSignalProto, encryptSignalProto, extractDeviceJids, generateMessageID, generateWAMessage, getUrlFromDirectPath, getWAUploadToServer, jidToSignalProtocolAddress, parseAndInjectE2ESessions, unixTimestampSeconds } from '../Utils'
 import { getUrlInfo } from '../Utils/link-preview'
 import { BinaryNode, BinaryNodeAttributes, getBinaryNodeChild, getBinaryNodeChildren, isJidGroup, isJidUser, jidDecode, jidEncode, jidNormalizedUser, JidWithDevice, reduceBinaryNodeToDictionary, S_WHATSAPP_NET } from '../WABinary'
 import { makeGroupsSocket } from './groups'
@@ -454,6 +454,8 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
 	const waUploadToServer = getWAUploadToServer(config, refreshMediaConn)
 
+	const waitForMsgMediaUpdate = bindWaitForEvent(ev, 'messages.media-update')
+
 	return {
 		...sock,
 		assertSessions,
@@ -464,6 +466,43 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 		refreshMediaConn,
 	    waUploadToServer,
 		fetchPrivacySettings,
+		updateMediaMessage: async(message: proto.IWebMessageInfo) => {
+			const content = assertMediaContent(message.message)
+			const mediaKey = content.mediaKey!
+			const meId = authState.creds.me!.id
+			const node = encryptMediaRetryRequest(message.key, mediaKey, meId)
+
+			let error: Error | undefined = undefined
+			await Promise.all(
+				[
+					sendNode(node),
+					waitForMsgMediaUpdate(update => {
+						const result = update.find(c => c.key.id === message.key.id)
+						if(result) {
+							if(result.error) {
+								error = result.error
+							} else {
+								try {
+									const media = decryptMediaRetryData(result.media!, mediaKey, result.key.id)
+									content.directPath = media.directPath
+									content.url = getUrlFromDirectPath(content.directPath!)
+								} catch(err) {
+									error = err
+								}
+							}
+
+							return true
+						}
+					})
+				]
+			)
+
+			if(error) {
+				throw error
+			}
+
+			return message
+		},
 		sendMessage: async(
 			jid: string,
 			content: AnyMessageContent,

src/Types/Events.ts

@@ -1,3 +1,4 @@
+import type { Boom } from '@hapi/boom'
 import type EventEmitter from 'events'
 import { proto } from '../../WAProto'
 import { AuthenticationCreds } from './Auth'
@@ -33,6 +34,7 @@ export type BaileysEventMap<T> = {
 
     'messages.delete': { keys: WAMessageKey[] } | { jid: string, all: true }
     'messages.update': WAMessageUpdate[]
+    'messages.media-update': { key: WAMessageKey, media?: { ciphertext: Uint8Array, iv: Uint8Array }, error?: Boom }[]
     /**
      * add/update the given messages. If they were received while the connection was online,
      * the update will have type: "notify"

src/Utils/messages-media.ts

@@ -10,9 +10,11 @@ import { join } from 'path'
 import type { Logger } from 'pino'
 import { Readable, Transform } from 'stream'
 import { URL } from 'url'
+import { proto } from '../../WAProto'
 import { DEFAULT_ORIGIN, MEDIA_PATH_MAP } from '../Defaults'
-import { CommonSocketConfig, DownloadableMessage, MediaConnInfo, MediaDecryptionKeyInfo, MediaType, MessageType, WAGenericMediaMessage, WAMediaUpload, WAMediaUploadFunction, WAMessageContent, WAProto } from '../Types'
+import { BaileysEventMap, CommonSocketConfig, DownloadableMessage, MediaConnInfo, MediaDecryptionKeyInfo, MediaType, MessageType, WAGenericMediaMessage, WAMediaUpload, WAMediaUploadFunction, WAMessageContent, WAProto } from '../Types'
-import { hkdf } from './crypto'
+import { BinaryNode, getBinaryNodeChild, getBinaryNodeChildBuffer, jidNormalizedUser } from '../WABinary'
+import { aesDecryptGCM, aesEncryptGCM, hkdf } from './crypto'
 import { generateMessageID } from './generics'
 
 const getTmpFilesDirectory = () => tmpdir()
@@ -344,12 +346,14 @@ export type MediaDownloadOptions = {
     endByte?: number
 }
 
+export const getUrlFromDirectPath = (directPath: string) => `https://${DEF_HOST}${directPath}`
+
 export const downloadContentFromMessage = (
 	{ mediaKey, directPath, url }: DownloadableMessage,
 	type: MediaType,
 	opts: MediaDownloadOptions = { }
 ) => {
-	const downloadUrl = url || `https://${DEF_HOST}${directPath}`
+	const downloadUrl = url || getUrlFromDirectPath(directPath)
 	const keys = getMediaKeys(mediaKey, type)
 
 	return downloadEncryptedContent(downloadUrl, keys, opts)
@@ -558,4 +562,98 @@ export const getWAUploadToServer = ({ customUploadHosts, fetchAgent, logger }: C
 
 		return urls
 	}
+}
+
+const GCM_AUTH_TAG_LENGTH: number | undefined = 128 >> 3
+
+const getMediaRetryKey = (mediaKey: Buffer | Uint8Array) => {
+	return hkdf(mediaKey, 32, { info: 'WhatsApp Media Retry Notification' })
+}
+
+/**
+ * Generate a binary node that will request the phone to re-upload the media & return the newly uploaded URL
+ */
+export const encryptMediaRetryRequest = (
+	key: proto.IMessageKey,
+	mediaKey: Buffer | Uint8Array,
+	meId: string
+) => {
+	const recp: proto.IServerErrorReceipt = { stanzaId: key.id }
+	const recpBuffer = proto.ServerErrorReceipt.encode(recp).finish()
+
+	const iv = Crypto.randomBytes(12)
+	const retryKey = getMediaRetryKey(mediaKey)
+	const ciphertext = aesEncryptGCM(recpBuffer, retryKey, iv, Buffer.from(key.id))
+
+	const req: BinaryNode = {
+		tag: 'receipt',
+		attrs: {
+			id: key.id,
+			to: jidNormalizedUser(meId),
+			type: 'server-error'
+		},
+		content: [
+			// this encrypt node is actually pretty useless
+			// the media is returned even without this node
+			// keeping it here to maintain parity with WA Web
+			{
+				tag: 'encrypt',
+				attrs: { },
+				content: [
+					{ tag: 'enc_p', attrs: { }, content: ciphertext },
+					{ tag: 'enc_iv', attrs: { }, content: iv }
+				]
+			},
+			{
+				tag: 'rmr',
+				attrs: {
+					jid: key.remoteJid,
+					from_me: (!!key.fromMe).toString(),
+					participant: key.participant || undefined
+				}
+			}
+		]
+	}
+
+	return req
+}
+
+export const decodeMediaRetryNode = (node: BinaryNode) => {
+	const rmrNode = getBinaryNodeChild(node, 'rmr')
+
+	const event: BaileysEventMap<any>['messages.media-update'][number] = {
+		key: {
+			id: node.attrs.id,
+			remoteJid: rmrNode.attrs.jid,
+			fromMe: rmrNode.attrs.from_me === 'true',
+			participant: rmrNode.attrs.participant
+		}
+	}
+
+	const errorNode = getBinaryNodeChild(node, 'error')
+	if(errorNode) {
+		const errorCode = +errorNode.attrs.code
+		event.error = new Boom(`Failed to re-upload media (${errorCode})`, { data: errorNode.attrs })
+	} else {
+		const encryptedInfoNode = getBinaryNodeChild(node, 'encrypt')
+		const ciphertext = getBinaryNodeChildBuffer(encryptedInfoNode, 'enc_p')
+		const iv = getBinaryNodeChildBuffer(encryptedInfoNode, 'enc_iv')
+		if(ciphertext && iv) {
+			event.media = { ciphertext, iv }
+		} else {
+			event.error = new Boom('Failed to re-upload media (missing ciphertext)', { statusCode: 404 })
+		}
+	}
+
+	return event
+}
+
+export const decryptMediaRetryData = (
+	{ ciphertext, iv }: { ciphertext: Uint8Array, iv: Uint8Array },
+	mediaKey: Uint8Array,
+	msgId: string
+) => {
+	const retryKey = getMediaRetryKey(mediaKey)
+	const plaintext = aesDecryptGCM(ciphertext, retryKey, iv, Buffer.from(msgId))
+	return proto.MediaRetryNotification.decode(plaintext)
 }