fix: memory blow on large files sending through URL

This commit is contained in:
canove
2025-05-05 21:11:40 -03:00
parent a85251250d
commit bff86ed4c1
2 changed files with 68 additions and 67 deletions

View File

@@ -345,24 +345,35 @@ export const encryptedStream = async(
const mediaKey = Crypto.randomBytes(32)
const { cipherKey, iv, macKey } = await getMediaKeys(mediaKey, mediaType)
const encWriteStream = new Readable({ read: () => {} })
const encFilePath = join(
getTmpFilesDirectory(),
mediaType + generateMessageIDV2() + "-enc"
);
const encFileWriteStream = createWriteStream(encFilePath);
let bodyPath: string | undefined
let writeStream: WriteStream | undefined
let didSaveToTmpPath = false
if(type === 'file') {
bodyPath = (media as WAMediaPayloadURL).url.toString()
} else if(saveOriginalFileIfRequired) {
bodyPath = join(getTmpFilesDirectory(), mediaType + generateMessageIDV2())
writeStream = createWriteStream(bodyPath)
didSaveToTmpPath = true
}
let originalFileStream: WriteStream | undefined;
let originalFilePath: string | undefined;
if (saveOriginalFileIfRequired) {
originalFilePath = join(
getTmpFilesDirectory(),
mediaType + generateMessageIDV2() + "-original"
);
originalFileStream = createWriteStream(originalFilePath);
}
let fileLength = 0
const aes = Crypto.createCipheriv('aes-256-cbc', cipherKey, iv)
let hmac = Crypto.createHmac('sha256', macKey!).update(iv)
let sha256Plain = Crypto.createHash('sha256')
let sha256Enc = Crypto.createHash('sha256')
const hmac = Crypto.createHmac('sha256', macKey!).update(iv)
const sha256Plain = Crypto.createHash('sha256')
const sha256Enc = Crypto.createHash('sha256')
const onChunk = (buff: Buffer) => {
sha256Enc.update(buff);
hmac.update(buff);
encFileWriteStream.write(buff);
};
try {
for await (const data of stream) {
@@ -379,68 +390,62 @@ export const encryptedStream = async(
data: { media, type }
}
)
}
sha256Plain = sha256Plain.update(data)
if(writeStream && !writeStream.write(data)) {
await once(writeStream, 'drain')
}
}
if (originalFileStream) {
if (!originalFileStream.write(data)) {
await once(originalFileStream, "drain");
}
}
sha256Plain.update(data)
onChunk(aes.update(data))
}
onChunk(aes.final())
const mac = hmac.digest().slice(0, 10)
sha256Enc = sha256Enc.update(mac)
sha256Enc.update(mac)
const fileSha256 = sha256Plain.digest()
const fileEncSha256 = sha256Enc.digest()
encWriteStream.push(mac)
encWriteStream.push(null)
encFileWriteStream.write(mac);
writeStream?.end()
stream.destroy()
encFileWriteStream.end();
originalFileStream?.end?.();
stream.destroy();
logger?.debug('encrypted data successfully')
return {
mediaKey,
encWriteStream,
bodyPath,
originalFilePath,
encFilePath,
mac,
fileEncSha256,
fileSha256,
fileLength,
didSaveToTmpPath
fileLength
}
} catch(error) {
// destroy all streams with error
encWriteStream.destroy()
writeStream?.destroy()
encFileWriteStream.destroy()
originalFileStream?.destroy?.()
aes.destroy()
hmac.destroy()
sha256Plain.destroy()
sha256Enc.destroy()
stream.destroy()
if(didSaveToTmpPath) {
try {
await fs.unlink(bodyPath!)
} catch(err) {
logger?.error({ err }, 'failed to save to tmp path')
}
}
try {
await fs.unlink(encFilePath)
if (originalFilePath) await fs.unlink(originalFilePath)
} catch(err) {
logger?.error({ err }, 'failed deleting tmp files')
}
throw error
}
function onChunk(buff: Buffer) {
sha256Enc = sha256Enc.update(buff)
hmac = hmac.update(buff)
encWriteStream.push(buff)
}
}
const DEF_HOST = 'mmg.whatsapp.net'
@@ -620,7 +625,8 @@ export const getWAUploadToServer = (
url,
stream,
{
...options,
...options,
maxRedirects: 0,
headers: {
...options.headers || { },
'Content-Type': 'application/octet-stream',

View File

@@ -1,7 +1,7 @@
import { Boom } from '@hapi/boom'
import axios from 'axios'
import { randomBytes } from 'crypto'
import { promises as fs } from 'fs'
import { createReadStream, promises as fs } from 'fs'
import { type Transform } from 'stream'
import { proto } from '../../WAProto'
import { MEDIA_KEYS, URL_REGEX, WA_DEFAULT_EPHEMERAL } from '../Defaults'
@@ -155,15 +155,14 @@ export const prepareWAMessageMedia = async(
(typeof uploadData['jpegThumbnail'] === 'undefined')
const requiresWaveformProcessing = mediaType === 'audio' && uploadData.ptt === true
const requiresAudioBackground = options.backgroundColor && mediaType === 'audio' && uploadData.ptt === true
const requiresOriginalForSomeProcessing = requiresDurationComputation || requiresThumbnailComputation
const requiresOriginalForSomeProcessing = requiresDurationComputation || requiresThumbnailComputation
const {
mediaKey,
encWriteStream,
bodyPath,
encFilePath,
originalFilePath,
fileEncSha256,
fileSha256,
fileLength,
didSaveToTmpPath
fileLength
} = await encryptedStream(
uploadData.media,
options.mediaTypeOverride || mediaType,
@@ -178,7 +177,7 @@ export const prepareWAMessageMedia = async(
const [{ mediaUrl, directPath }] = await Promise.all([
(async() => {
const result = await options.upload(
encWriteStream,
createReadStream(encFilePath),
{ fileEncSha256B64, mediaType, timeoutMs: options.mediaUploadTimeoutMs }
)
logger?.debug({ mediaType, cacheableKey }, 'uploaded media')
@@ -190,7 +189,7 @@ export const prepareWAMessageMedia = async(
const {
thumbnail,
originalImageDimensions
} = await generateThumbnail(bodyPath!, mediaType as 'image' | 'video', options)
} = await generateThumbnail(originalFilePath!, mediaType as 'image' | 'video', options)
uploadData.jpegThumbnail = thumbnail
if(!uploadData.width && originalImageDimensions) {
uploadData.width = originalImageDimensions.width
@@ -202,12 +201,12 @@ export const prepareWAMessageMedia = async(
}
if(requiresDurationComputation) {
uploadData.seconds = await getAudioDuration(bodyPath!)
uploadData.seconds = await getAudioDuration(originalFilePath!)
logger?.debug('computed audio duration')
}
if(requiresWaveformProcessing) {
uploadData.waveform = await getAudioWaveform(bodyPath!, logger)
uploadData.waveform = await getAudioWaveform(originalFilePath!, logger)
logger?.debug('processed waveform')
}
@@ -222,17 +221,13 @@ export const prepareWAMessageMedia = async(
])
.finally(
async() => {
encWriteStream.destroy()
// remove tmp files
if(didSaveToTmpPath && bodyPath) {
try {
await fs.access(bodyPath)
await fs.unlink(bodyPath)
logger?.debug('removed tmp file')
} catch(error) {
logger?.warn('failed to remove tmp file')
}
}
try {
await fs.unlink(encFilePath)
if (originalFilePath) await fs.unlink(originalFilePath)
logger?.debug('removed tmp files')
} catch(error) {
logger?.warn('failed to remove tmp file')
}
}
)