From 3e5474104213ba6e37a6951cbf2cb21b325e11e1 Mon Sep 17 00:00:00 2001
From: Adhiraj Singh <adhirajsingh1001@gmail.com>
Date: Wed, 10 Nov 2021 19:22:00 +0530
Subject: [PATCH] feat: cache media uploads

---
 src/Socket/messages-send.ts | 10 ++++++----
 src/Types/Message.ts        |  9 +++++++--
 src/Types/index.ts          |  2 ++
 src/Utils/messages.ts       | 27 +++++++++++++++++++--------
 4 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/src/Socket/messages-send.ts b/src/Socket/messages-send.ts
index cd97fcc..5a547ba 100644
--- a/src/Socket/messages-send.ts
+++ b/src/Socket/messages-send.ts
@@ -379,7 +379,7 @@ export const makeMessagesSocket = (config: SocketConfig) => {
         return msgId
     } 
 
-    const waUploadToServer: WAMediaUploadFunction = async(stream, { mediaType, fileEncSha256B64 }) => {
+    const waUploadToServer: WAMediaUploadFunction = async(stream, { mediaType, fileEncSha256B64, timeoutMs }) => {
 		// send a query JSON to obtain the url & auth token to upload our media
 		let uploadInfo = await refreshMediaConn(false)
 
@@ -399,7 +399,8 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 						agent: {
 							https: config.agent
 						},
-						body: stream
+						body: stream,
+                        timeout: timeoutMs
 					}
 				)
 				const result = JSON.parse(responseText)
@@ -456,10 +457,11 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 					{
 						...options,
 						logger,
-						userJid: userJid,
+						userJid,
                         // multi-device does not have this yet
 						//getUrlInfo: generateUrlInfo,
-						upload: waUploadToServer
+						upload: waUploadToServer,
+                        mediaCache: config.mediaCache,
 					}
 				)
                 const isDeleteMsg = 'delete' in content && !!content.delete
diff --git a/src/Types/Message.ts b/src/Types/Message.ts
index 335c9f0..fb768f2 100644
--- a/src/Types/Message.ts
+++ b/src/Types/Message.ts
@@ -1,6 +1,7 @@
 import type { ReadStream } from "fs"
 import type { Logger } from "pino"
 import type { URL } from "url"
+import type NodeCache from "node-cache"
 import type { GroupMetadata } from "./GroupMetadata"
 import { proto } from '../../WAProto'
 
@@ -121,18 +122,22 @@ export type MiscMessageGenerationOptions = {
 	quoted?: WAMessage
     /** disappearing messages settings */
     ephemeralExpiration?: number | string
+
+    mediaUploadTimeoutMs?: number
 }
 export type MessageGenerationOptionsFromContent = MiscMessageGenerationOptions & {
 	userJid: string
 }
 
-export type WAMediaUploadFunction = (readStream: ReadStream, opts: { fileEncSha256B64: string, mediaType: MediaType }) => Promise<{ mediaUrl: string }>
+export type WAMediaUploadFunction = (readStream: ReadStream, opts: { fileEncSha256B64: string, mediaType: MediaType, timeoutMs?: number }) => Promise<{ mediaUrl: string }>
 
 export type MediaGenerationOptions = {
 	logger?: Logger
     upload: WAMediaUploadFunction
     /** cache media so it does not have to be uploaded again */
-    mediaCache?: (url: string) => Promise<WAGenericMediaMessage> | WAGenericMediaMessage
+    mediaCache?: NodeCache
+
+    mediaUploadTimeoutMs?: number
 }
 export type MessageContentGenerationOptions = MediaGenerationOptions & {
 	getUrlInfo?: (text: string) => Promise<WAUrlInfo>
diff --git a/src/Types/index.ts b/src/Types/index.ts
index bcc7b1f..6816d21 100644
--- a/src/Types/index.ts
+++ b/src/Types/index.ts
@@ -48,6 +48,8 @@ export type SocketConfig = {
     emitOwnEvents: boolean
     /** provide a cache to store a user's device list */
     userDevicesCache?: NodeCache
+    /** provide a cache to store media, so does not have to be re-uploaded */
+    mediaCache?: NodeCache
 }
 
 export enum DisconnectReason {
diff --git a/src/Utils/messages.ts b/src/Utils/messages.ts
index 0ceab16..3550bc4 100644
--- a/src/Utils/messages.ts
+++ b/src/Utils/messages.ts
@@ -67,13 +67,19 @@ export const prepareWAMessageMedia = async(
 		[mediaType]: undefined,
 		media: message[mediaType]
 	}
+	// check if cacheable + generate cache key
+	const cacheableKey = typeof uploadData.media === 'object' && 
+			('url' in uploadData.media) && 
+			!!uploadData.media.url && 
+			!!options.mediaCache && (
+				// generate the key
+				mediaType + ':' + uploadData.media.url!.toString()
+			)
 	// check for cache hit
-	if(typeof uploadData.media === 'object' && 'url' in uploadData.media) {
-		const result = !!options.mediaCache && await options.mediaCache!(uploadData.media.url?.toString())
-		if(result) {
-			return WAProto.Message.fromObject({
-				[`${mediaType}Message`]: result
-			})
+	if(cacheableKey) {
+		const mediaBuff: Buffer = options.mediaCache!.get(cacheableKey)
+		if(mediaBuff) {
+			return WAProto.Message.decode(mediaBuff)
 		}
 	}
 	if(mediaType === 'document' && !uploadData.fileName) {
@@ -114,7 +120,7 @@ export const prepareWAMessageMedia = async(
 	}
 	const {mediaUrl} = await options.upload(
 		createReadStream(encBodyPath),
-		{ fileEncSha256B64, mediaType }
+		{ fileEncSha256B64, mediaType, timeoutMs: options.mediaUploadTimeoutMs }
 	)
 	// remove tmp files
 	await Promise.all(
@@ -138,7 +144,12 @@ export const prepareWAMessageMedia = async(
 			}
 		)
 	}
-	return WAProto.Message.fromObject(content)
+	const obj = WAProto.Message.fromObject(content)
+	if(cacheableKey) {
+		options.mediaCache!.set(cacheableKey, WAProto.Message.encode(obj))
+	}
+
+	return obj
 }
 export const prepareDisappearingMessageSettingContent = (ephemeralExpiration?: number) => {
 	ephemeralExpiration = ephemeralExpiration || 0