From 43d178753255986b6a481c1b55cc85c2b88f3d9d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas=20de=20Oliveira=20Lopes?= <55464917+jlucaso1@users.noreply.github.com> Date: Fri, 5 Dec 2025 22:06:24 -0300 Subject: [PATCH] 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 | 234 ++++++++++-- src/__tests__/Utils/messages-media.test.ts | 405 +++++++++++++++++++++ 2 files changed, 604 insertions(+), 35 deletions(-) create mode 100644 src/__tests__/Utils/messages-media.test.ts diff --git a/src/Utils/messages-media.ts b/src/Utils/messages-media.ts index 9dd82d72..d2cbe7d3 100644 --- a/src/Utils/messages-media.ts +++ b/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 + timeoutMs?: number + agent?: Agent +} + +export const uploadWithNodeHttp = async ( + { url, filePath, headers, timeoutMs, agent }: UploadParams, + redirectCount = 0 +): Promise => { + 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 => { + // 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 => { + 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 @@ -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) + })() + + 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) - })(), - '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 diff --git a/src/__tests__/Utils/messages-media.test.ts b/src/__tests__/Utils/messages-media.test.ts new file mode 100644 index 00000000..337ef5f5 --- /dev/null +++ b/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 => { + 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 => { + 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 => { + 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('Not JSON') + }) + }) + + 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]) + } + }) +})