mirror of
https://github.com/FranP-code/Baileys.git
synced 2025-10-13 00:32:22 +00:00
fix: memory blow on large files sending through URL
This commit is contained in:
@@ -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',
|
||||
|
||||
@@ -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')
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user