fix: improve media upload handling and reduce memory pressure (#2128)

* fix: improve media upload handling and reduce memory pressure

* fix: enhance media upload functionality and add comprehensive tests

src/Utils/messages-media.ts

@@ -3,6 +3,7 @@ import { exec } from 'child_process'
 import * as Crypto from 'crypto'
 import { once } from 'events'
 import { createReadStream, createWriteStream, promises as fs, WriteStream } from 'fs'
+import type { Agent } from 'https'
 import type { IAudioMetadata } from 'music-metadata'
 import { tmpdir } from 'os'
 import { join } from 'path'
@@ -410,10 +411,13 @@ export const encryptedStream = async (
 	const sha256Plain = Crypto.createHash('sha256')
 	const sha256Enc = Crypto.createHash('sha256')
 
-	const onChunk = (buff: Buffer) => {
+	const onChunk = async (buff: Buffer) => {
 		sha256Enc.update(buff)
 		hmac.update(buff)
-		encFileWriteStream.write(buff)
+		// Handle backpressure: if write returns false, wait for drain
+		if (!encFileWriteStream.write(buff)) {
+			await once(encFileWriteStream, 'drain')
+		}
 	}
 
 	try {
@@ -437,10 +441,10 @@ export const encryptedStream = async (
 			}
 
 			sha256Plain.update(data)
-			onChunk(aes.update(data))
+			await onChunk(aes.update(data))
 		}
 
-		onChunk(aes.final())
+		await onChunk(aes.final())
 
 		const mac = hmac.digest().slice(0, 10)
 		sha256Enc.update(mac)
@@ -450,10 +454,18 @@ export const encryptedStream = async (
 
 		encFileWriteStream.write(mac)
 
+		const encFinishPromise = once(encFileWriteStream, 'finish')
+		const originalFinishPromise = originalFileStream ? once(originalFileStream, 'finish') : Promise.resolve()
+
 		encFileWriteStream.end()
 		originalFileStream?.end?.()
 		stream.destroy()
 
+		// Wait for write streams to fully flush to disk
+		// This helps reduce memory pressure by allowing OS to release buffers
+		await encFinishPromise
+		await originalFinishPromise
+
 		logger?.debug('encrypted data successfully')
 
 		return {
@@ -639,6 +651,161 @@ export function extensionForMediaMessage(message: WAMessageContent) {
 	return extension
 }
 
+const isNodeRuntime = (): boolean => {
+	return (
+		typeof process !== 'undefined' &&
+		process.versions?.node !== null &&
+		typeof process.versions.bun === 'undefined' &&
+		typeof (globalThis as any).Deno === 'undefined'
+	)
+}
+
+type MediaUploadResult = {
+	url?: string
+	direct_path?: string
+	meta_hmac?: string
+	ts?: number
+	fbid?: number
+}
+
+export type UploadParams = {
+	url: string
+	filePath: string
+	headers: Record<string, string>
+	timeoutMs?: number
+	agent?: Agent
+}
+
+export const uploadWithNodeHttp = async (
+	{ url, filePath, headers, timeoutMs, agent }: UploadParams,
+	redirectCount = 0
+): Promise<MediaUploadResult | undefined> => {
+	if (redirectCount > 5) {
+		throw new Error('Too many redirects')
+	}
+
+	const parsedUrl = new URL(url)
+	const httpModule = parsedUrl.protocol === 'https:' ? await import('https') : await import('http')
+
+	// Get file size for Content-Length header (required for Node.js streaming)
+	const fileStats = await fs.stat(filePath)
+	const fileSize = fileStats.size
+
+	return new Promise((resolve, reject) => {
+		const req = httpModule.request(
+			{
+				hostname: parsedUrl.hostname,
+				port: parsedUrl.port || (parsedUrl.protocol === 'https:' ? 443 : 80),
+				path: parsedUrl.pathname + parsedUrl.search,
+				method: 'POST',
+				headers: {
+					...headers,
+					'Content-Length': fileSize
+				},
+				agent,
+				timeout: timeoutMs
+			},
+			res => {
+				// Handle redirects (3xx)
+				if (res.statusCode && res.statusCode >= 300 && res.statusCode < 400 && res.headers.location) {
+					res.resume() // Consume response to free resources
+					const newUrl = new URL(res.headers.location, url).toString()
+					resolve(
+						uploadWithNodeHttp(
+							{
+								url: newUrl,
+								filePath,
+								headers,
+								timeoutMs,
+								agent
+							},
+							redirectCount + 1
+						)
+					)
+					return
+				}
+
+				let body = ''
+				res.on('data', chunk => (body += chunk))
+				res.on('end', () => {
+					try {
+						resolve(JSON.parse(body))
+					} catch {
+						resolve(undefined)
+					}
+				})
+			}
+		)
+
+		req.on('error', reject)
+		req.on('timeout', () => {
+			req.destroy()
+			reject(new Error('Upload timeout'))
+		})
+
+		const stream = createReadStream(filePath)
+		stream.pipe(req)
+		stream.on('error', err => {
+			req.destroy()
+			reject(err)
+		})
+	})
+}
+
+const uploadWithFetch = async ({
+	url,
+	filePath,
+	headers,
+	timeoutMs,
+	agent
+}: UploadParams): Promise<MediaUploadResult | undefined> => {
+	// Convert Node.js Readable to Web ReadableStream
+	const nodeStream = createReadStream(filePath)
+	const webStream = Readable.toWeb(nodeStream) as ReadableStream
+
+	const response = await fetch(url, {
+		dispatcher: agent,
+		method: 'POST',
+		body: webStream,
+		headers,
+		duplex: 'half',
+		signal: timeoutMs ? AbortSignal.timeout(timeoutMs) : undefined
+	})
+
+	try {
+		return (await response.json()) as MediaUploadResult
+	} catch {
+		return undefined
+	}
+}
+
+/**
+ * Uploads media to WhatsApp servers.
+ *
+ * ## Why we have two upload implementations:
+ *
+ * Node.js's native `fetch` (powered by undici) has a known bug where it buffers
+ * the entire request body in memory before sending, even when using streams.
+ * This causes memory issues with large files (e.g., 1GB file = 1GB+ memory usage).
+ * See: https://github.com/nodejs/undici/issues/4058
+ *
+ * Other runtimes (Bun, Deno, browsers) correctly stream the request body without
+ * buffering, so we can use the web-standard Fetch API there.
+ *
+ * ## Future considerations:
+ * Once the undici bug is fixed, we can simplify this to use only the Fetch API
+ * across all runtimes. Monitor the GitHub issue for updates.
+ */
+const uploadMedia = async (params: UploadParams, logger?: ILogger): Promise<MediaUploadResult | undefined> => {
+	if (isNodeRuntime()) {
+		logger?.debug('Using Node.js https module for upload (avoids undici buffering bug)')
+		return uploadWithNodeHttp(params)
+	} else {
+		logger?.debug('Using web-standard Fetch API for upload')
+		return uploadWithFetch(params)
+	}
+}
+
 export const getWAUploadToServer = (
 	{ customUploadHosts, fetchAgent, logger, options }: SocketConfig,
 	refreshMediaConn: (force: boolean) => Promise<MediaConnInfo>
@@ -652,45 +819,42 @@ export const getWAUploadToServer = (
 
 		fileEncSha256B64 = encodeBase64EncodedStringForUpload(fileEncSha256B64)
 
+		// Prepare common headers
+		const customHeaders = (() => {
+			const hdrs = options?.headers
+			if (!hdrs) return {}
+			return Array.isArray(hdrs) ? Object.fromEntries(hdrs) : (hdrs as Record<string, string>)
+		})()
+
+		const headers = {
+			...customHeaders,
+			'Content-Type': 'application/octet-stream',
+			Origin: DEFAULT_ORIGIN
+		}
+
 		for (const { hostname } of hosts) {
 			logger.debug(`uploading to "${hostname}"`)
 
-			const auth = encodeURIComponent(uploadInfo.auth) // the auth token
+			const auth = encodeURIComponent(uploadInfo.auth)
 			const url = `https://${hostname}${MEDIA_PATH_MAP[mediaType]}/${fileEncSha256B64}?auth=${auth}&token=${fileEncSha256B64}`
-			// eslint-disable-next-line @typescript-eslint/no-explicit-any
-			let result: any
+
+			let result: MediaUploadResult | undefined
 			try {
-				const stream = createReadStream(filePath)
-				const response = await fetch(url, {
-					dispatcher: fetchAgent,
-					method: 'POST',
-					body: stream as any,
-					headers: {
-						...(() => {
-							const hdrs = options?.headers
-							if (!hdrs) return {}
-							return Array.isArray(hdrs) ? Object.fromEntries(hdrs) : (hdrs as Record<string, string>)
-						})(),
-						'Content-Type': 'application/octet-stream',
-						Origin: DEFAULT_ORIGIN
+				result = await uploadMedia(
+					{
+						url,
+						filePath,
+						headers,
+						timeoutMs,
+						agent: fetchAgent
 					},
-					duplex: 'half',
-					// Note: custom agents/proxy require undici Agent; omitted here.
-					signal: timeoutMs ? AbortSignal.timeout(timeoutMs) : undefined
-				})
-				let parsed: any = undefined
-				try {
-					parsed = await response.json()
-				} catch {
-					parsed = undefined
-				}
+					logger
+				)
 
-				result = parsed
-
-				if (result?.url || result?.directPath) {
+				if (result?.url || result?.direct_path) {
 					urls = {
-						mediaUrl: result.url,
-						directPath: result.direct_path,
+						mediaUrl: result.url!,
+						directPath: result.direct_path!,
 						meta_hmac: result.meta_hmac,
 						fbid: result.fbid,
 						ts: result.ts

src/__tests__/Utils/messages-media.test.ts

@@ -0,0 +1,405 @@
+import * as fs from 'fs'
+import * as http from 'http'
+import * as os from 'os'
+import * as path from 'path'
+import { Readable } from 'stream'
+import { encryptedStream, type UploadParams, uploadWithNodeHttp } from '../../Utils/messages-media'
+
+const createTempFile = async (content: string): Promise<string> => {
+	const filePath = path.join(os.tmpdir(), `test-upload-${Date.now()}.txt`)
+	await fs.promises.writeFile(filePath, content)
+	return filePath
+}
+
+const cleanupTempFile = async (filePath: string): Promise<void> => {
+	try {
+		await fs.promises.unlink(filePath)
+	} catch {}
+}
+
+describe('uploadWithNodeHttp', () => {
+	let server: http.Server
+	let serverPort: number
+	let tempFilePath: string
+	const testFileContent = 'Hello, this is test content for upload!'
+
+	beforeAll(async () => {
+		tempFilePath = await createTempFile(testFileContent)
+	})
+
+	afterAll(async () => {
+		await cleanupTempFile(tempFilePath)
+	})
+
+	afterEach(() => {
+		if (server) {
+			server.close()
+		}
+	})
+
+	const startServer = (handler: http.RequestListener): Promise<number> => {
+		return new Promise(resolve => {
+			server = http.createServer(handler)
+			server.listen(0, () => {
+				const address = server.address()
+				if (address && typeof address === 'object') {
+					serverPort = address.port
+					resolve(serverPort)
+				}
+			})
+		})
+	}
+
+	it('should successfully upload a file and receive JSON response', async () => {
+		const expectedResponse = { url: 'https://example.com/media/123', direct_path: '/media/123' }
+		let receivedBody = ''
+
+		await startServer((req, res) => {
+			req.on('data', chunk => {
+				receivedBody += chunk
+			})
+			req.on('end', () => {
+				res.writeHead(200, { 'Content-Type': 'application/json' })
+				res.end(JSON.stringify(expectedResponse))
+			})
+		})
+
+		const params: UploadParams = {
+			url: `http://localhost:${serverPort}/upload`,
+			filePath: tempFilePath,
+			headers: { 'Content-Type': 'application/octet-stream' }
+		}
+
+		const result = await uploadWithNodeHttp(params)
+
+		expect(result).toEqual(expectedResponse)
+		expect(receivedBody).toBe(testFileContent)
+	})
+
+	it('should follow a single redirect (302)', async () => {
+		const expectedResponse = { url: 'https://example.com/media/456', direct_path: '/media/456' }
+		let requestCount = 0
+
+		await startServer((req, res) => {
+			requestCount++
+			if (req.url === '/upload') {
+				res.writeHead(302, { Location: `http://localhost:${serverPort}/final` })
+				res.end()
+			} else if (req.url === '/final') {
+				let body = ''
+				req.on('data', chunk => (body += chunk))
+				req.on('end', () => {
+					res.writeHead(200, { 'Content-Type': 'application/json' })
+					res.end(JSON.stringify(expectedResponse))
+				})
+			}
+		})
+
+		const params: UploadParams = {
+			url: `http://localhost:${serverPort}/upload`,
+			filePath: tempFilePath,
+			headers: { 'Content-Type': 'application/octet-stream' }
+		}
+
+		const result = await uploadWithNodeHttp(params)
+
+		expect(result).toEqual(expectedResponse)
+		expect(requestCount).toBe(2)
+	})
+
+	it('should follow multiple redirects (301 -> 302 -> 200)', async () => {
+		const expectedResponse = { url: 'https://example.com/media/789', direct_path: '/media/789' }
+		let requestCount = 0
+
+		await startServer((req, res) => {
+			requestCount++
+			if (req.url === '/upload') {
+				res.writeHead(301, { Location: `http://localhost:${serverPort}/redirect1` })
+				res.end()
+			} else if (req.url === '/redirect1') {
+				res.writeHead(302, { Location: `http://localhost:${serverPort}/redirect2` })
+				res.end()
+			} else if (req.url === '/redirect2') {
+				res.writeHead(307, { Location: `http://localhost:${serverPort}/final` })
+				res.end()
+			} else if (req.url === '/final') {
+				let body = ''
+				req.on('data', chunk => (body += chunk))
+				req.on('end', () => {
+					res.writeHead(200, { 'Content-Type': 'application/json' })
+					res.end(JSON.stringify(expectedResponse))
+				})
+			}
+		})
+
+		const params: UploadParams = {
+			url: `http://localhost:${serverPort}/upload`,
+			filePath: tempFilePath,
+			headers: { 'Content-Type': 'application/octet-stream' }
+		}
+
+		const result = await uploadWithNodeHttp(params)
+
+		expect(result).toEqual(expectedResponse)
+		expect(requestCount).toBe(4)
+	})
+
+	it('should throw error on too many redirects (more than 5)', async () => {
+		await startServer((req, res) => {
+			const currentNum = parseInt(req.url?.replace('/redirect', '') || '0')
+			res.writeHead(302, { Location: `http://localhost:${serverPort}/redirect${currentNum + 1}` })
+			res.end()
+		})
+
+		const params: UploadParams = {
+			url: `http://localhost:${serverPort}/redirect0`,
+			filePath: tempFilePath,
+			headers: { 'Content-Type': 'application/octet-stream' }
+		}
+
+		await expect(uploadWithNodeHttp(params)).rejects.toThrow('Too many redirects')
+	})
+
+	it('should return undefined for non-JSON response', async () => {
+		await startServer((req, res) => {
+			let body = ''
+			req.on('data', chunk => (body += chunk))
+			req.on('end', () => {
+				res.writeHead(200, { 'Content-Type': 'text/html' })
+				res.end('<html>Not JSON</html>')
+			})
+		})
+
+		const params: UploadParams = {
+			url: `http://localhost:${serverPort}/upload`,
+			filePath: tempFilePath,
+			headers: { 'Content-Type': 'application/octet-stream' }
+		}
+
+		const result = await uploadWithNodeHttp(params)
+
+		expect(result).toBeUndefined()
+	})
+
+	it('should handle relative redirect URLs', async () => {
+		const expectedResponse = { url: 'https://example.com/media/rel', direct_path: '/media/rel' }
+		let requestCount = 0
+
+		await startServer((req, res) => {
+			requestCount++
+			if (req.url === '/upload') {
+				res.writeHead(302, { Location: '/final' })
+				res.end()
+			} else if (req.url === '/final') {
+				let body = ''
+				req.on('data', chunk => (body += chunk))
+				req.on('end', () => {
+					res.writeHead(200, { 'Content-Type': 'application/json' })
+					res.end(JSON.stringify(expectedResponse))
+				})
+			}
+		})
+
+		const params: UploadParams = {
+			url: `http://localhost:${serverPort}/upload`,
+			filePath: tempFilePath,
+			headers: { 'Content-Type': 'application/octet-stream' }
+		}
+
+		const result = await uploadWithNodeHttp(params)
+
+		expect(result).toEqual(expectedResponse)
+		expect(requestCount).toBe(2)
+	})
+
+	it('should preserve headers on redirect', async () => {
+		const expectedResponse = { success: true }
+		let capturedHeaders: http.IncomingHttpHeaders | undefined
+
+		await startServer((req, res) => {
+			if (req.url === '/upload') {
+				res.writeHead(302, { Location: `http://localhost:${serverPort}/final` })
+				res.end()
+			} else if (req.url === '/final') {
+				capturedHeaders = req.headers
+				let body = ''
+				req.on('data', chunk => (body += chunk))
+				req.on('end', () => {
+					res.writeHead(200, { 'Content-Type': 'application/json' })
+					res.end(JSON.stringify(expectedResponse))
+				})
+			}
+		})
+
+		const customHeaders = {
+			'Content-Type': 'application/octet-stream',
+			'X-Custom-Header': 'test-value',
+			Authorization: 'Bearer token123'
+		}
+
+		const params: UploadParams = {
+			url: `http://localhost:${serverPort}/upload`,
+			filePath: tempFilePath,
+			headers: customHeaders
+		}
+
+		const result = await uploadWithNodeHttp(params)
+
+		expect(result).toEqual(expectedResponse)
+		expect(capturedHeaders?.['x-custom-header']).toBe('test-value')
+		expect(capturedHeaders?.['authorization']).toBe('Bearer token123')
+	})
+
+	it('should re-stream file content on redirect', async () => {
+		const expectedResponse = { success: true }
+		let finalReceivedBody = ''
+
+		await startServer((req, res) => {
+			if (req.url === '/upload') {
+				req.on('data', () => {})
+				req.on('end', () => {
+					res.writeHead(302, { Location: `http://localhost:${serverPort}/final` })
+					res.end()
+				})
+			} else if (req.url === '/final') {
+				req.on('data', chunk => {
+					finalReceivedBody += chunk
+				})
+				req.on('end', () => {
+					res.writeHead(200, { 'Content-Type': 'application/json' })
+					res.end(JSON.stringify(expectedResponse))
+				})
+			}
+		})
+
+		const params: UploadParams = {
+			url: `http://localhost:${serverPort}/upload`,
+			filePath: tempFilePath,
+			headers: { 'Content-Type': 'application/octet-stream' }
+		}
+
+		const result = await uploadWithNodeHttp(params)
+
+		expect(result).toEqual(expectedResponse)
+		expect(finalReceivedBody).toBe(testFileContent)
+	})
+})
+
+describe('encryptedStream', () => {
+	const cleanupFiles = async (files: (string | undefined)[]) => {
+		for (const file of files) {
+			if (file) {
+				try {
+					await fs.promises.unlink(file)
+				} catch {}
+			}
+		}
+	}
+
+	it('should encrypt a buffer and return valid result without hanging', async () => {
+		const testData = Buffer.from('Hello, this is test content for encryption!')
+
+		const result = await encryptedStream(testData, 'image')
+
+		expect(result).toBeDefined()
+		expect(result.mediaKey).toBeDefined()
+		expect(result.mediaKey.length).toBe(32)
+		expect(result.encFilePath).toBeDefined()
+		expect(result.fileSha256).toBeDefined()
+		expect(result.fileEncSha256).toBeDefined()
+		expect(result.mac).toBeDefined()
+		expect(result.mac.length).toBe(10)
+		expect(result.fileLength).toBe(testData.length)
+
+		const encFileExists = await fs.promises
+			.access(result.encFilePath)
+			.then(() => true)
+			.catch(() => false)
+		expect(encFileExists).toBe(true)
+
+		await cleanupFiles([result.encFilePath, result.originalFilePath])
+	})
+
+	it('should encrypt a stream and complete without race condition', async () => {
+		const chunks = ['chunk1', 'chunk2', 'chunk3', 'chunk4', 'chunk5']
+		const testStream = Readable.from(chunks.map(c => Buffer.from(c)))
+
+		const result = await encryptedStream({ stream: testStream }, 'document')
+
+		expect(result).toBeDefined()
+		expect(result.mediaKey).toBeDefined()
+		expect(result.encFilePath).toBeDefined()
+		expect(result.fileLength).toBe(chunks.join('').length)
+
+		await cleanupFiles([result.encFilePath, result.originalFilePath])
+	})
+
+	it('should save original file when saveOriginalFileIfRequired is true', async () => {
+		const testData = Buffer.from('Original file content to save')
+
+		const result = await encryptedStream(testData, 'audio', {
+			saveOriginalFileIfRequired: true
+		})
+
+		expect(result).toBeDefined()
+		expect(result.originalFilePath).toBeDefined()
+
+		const originalContent = await fs.promises.readFile(result.originalFilePath!)
+		expect(originalContent.toString()).toBe(testData.toString())
+
+		await cleanupFiles([result.encFilePath, result.originalFilePath])
+	})
+
+	it('should complete encryption for various media types', async () => {
+		const mediaTypes = ['image', 'video', 'audio', 'document', 'sticker'] as const
+		const testData = Buffer.from('Test data for different media types')
+
+		for (const mediaType of mediaTypes) {
+			const result = await encryptedStream(testData, mediaType)
+
+			expect(result).toBeDefined()
+			expect(result.mediaKey).toBeDefined()
+			expect(result.encFilePath).toBeDefined()
+
+			await cleanupFiles([result.encFilePath, result.originalFilePath])
+		}
+	})
+
+	it('should handle empty buffer without hanging', async () => {
+		const emptyData = Buffer.from('')
+
+		const result = await encryptedStream(emptyData, 'image')
+
+		expect(result).toBeDefined()
+		expect(result.fileLength).toBe(0)
+		expect(result.encFilePath).toBeDefined()
+
+		await cleanupFiles([result.encFilePath, result.originalFilePath])
+	})
+
+	it('should handle small content that finishes quickly', async () => {
+		const smallData = Buffer.from('x')
+
+		const result = await encryptedStream(smallData, 'image')
+
+		expect(result).toBeDefined()
+		expect(result.fileLength).toBe(1)
+
+		await cleanupFiles([result.encFilePath, result.originalFilePath])
+	})
+
+	it('should complete multiple concurrent encryptions without deadlock', async () => {
+		const testData = Buffer.from('Concurrent encryption test')
+
+		const promises = Array.from({ length: 5 }, () => encryptedStream(testData, 'image'))
+
+		const results = await Promise.all(promises)
+
+		expect(results.length).toBe(5)
+		for (const result of results) {
+			expect(result).toBeDefined()
+			expect(result.mediaKey).toBeDefined()
+			await cleanupFiles([result.encFilePath, result.originalFilePath])
+		}
+	})
+})