feat: support big file downloads (read from WA cache + streaming) (#201697)
* feat: support big file downloads (read from WA cache + streaming) * fix: simplify downloadMediaStream to use a single handle instead of two handles (resultHandle + blobHandle) with disposal gaps, get the blob handle directly from evaluateHandle and read metadata via blobHandle.evaluate. This eliminates the intermediate resultHandle entirely. blobhandle is disposed in all exit paths: null metadata, evaluate error, and generator finally.
This commit is contained in:
Adi1231234
GitHub
parent
2dc9466fac
commit
4df15e58b9
2 changed files with 175 additions and 76 deletions
|
|
@ -1,5 +1,6 @@
|
|||
'use strict';
|
||||
|
||||
const { Readable } = require('stream');
|
||||
const Base = require('./Base');
|
||||
const MessageMedia = require('./MessageMedia');
|
||||
const Location = require('./Location');
|
||||
|
|
@ -507,84 +508,25 @@ class Message extends Base {
|
|||
}
|
||||
|
||||
/**
|
||||
* Downloads and returns the attatched message media
|
||||
* @returns {Promise<MessageMedia>}
|
||||
* Downloads and returns the attached message media
|
||||
* @returns {Promise<MessageMedia|undefined>}
|
||||
*/
|
||||
async downloadMedia() {
|
||||
if (!this.hasMedia) {
|
||||
return undefined;
|
||||
}
|
||||
if (!this.hasMedia) return undefined;
|
||||
|
||||
const result = await this.client.pupPage.evaluate(async (msgId) => {
|
||||
const msg =
|
||||
window.require('WAWebCollections').Msg.get(msgId) ||
|
||||
(
|
||||
await window
|
||||
.require('WAWebCollections')
|
||||
.Msg.getMessagesById([msgId])
|
||||
)?.messages?.[0];
|
||||
const resolved = await window.WWebJS.resolveMediaBlob(msgId);
|
||||
if (!resolved) return null;
|
||||
|
||||
// REUPLOADING mediaStage means the media is expired and the download button is spinning, cannot be downloaded now
|
||||
if (
|
||||
!msg ||
|
||||
!msg.mediaData ||
|
||||
msg.mediaData.mediaStage === 'REUPLOADING'
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
if (msg.mediaData.mediaStage != 'RESOLVED') {
|
||||
// try to resolve media
|
||||
await msg.downloadMedia({
|
||||
downloadEvenIfExpensive: true,
|
||||
rmrReason: 1,
|
||||
});
|
||||
}
|
||||
|
||||
if (
|
||||
msg.mediaData.mediaStage.includes('ERROR') ||
|
||||
msg.mediaData.mediaStage === 'FETCHING'
|
||||
) {
|
||||
// media could not be downloaded
|
||||
return undefined;
|
||||
}
|
||||
|
||||
try {
|
||||
const mockQpl = {
|
||||
addAnnotations: function () {
|
||||
return this;
|
||||
},
|
||||
addPoint: function () {
|
||||
return this;
|
||||
},
|
||||
};
|
||||
const decryptedMedia = await window
|
||||
.require('WAWebDownloadManager')
|
||||
.downloadManager.downloadAndMaybeDecrypt({
|
||||
directPath: msg.directPath,
|
||||
encFilehash: msg.encFilehash,
|
||||
filehash: msg.filehash,
|
||||
mediaKey: msg.mediaKey,
|
||||
mediaKeyTimestamp: msg.mediaKeyTimestamp,
|
||||
type: msg.type,
|
||||
signal: new AbortController().signal,
|
||||
downloadQpl: mockQpl,
|
||||
});
|
||||
|
||||
const data =
|
||||
await window.WWebJS.arrayBufferToBase64Async(
|
||||
decryptedMedia,
|
||||
);
|
||||
|
||||
return {
|
||||
data,
|
||||
mimetype: msg.mimetype,
|
||||
filename: msg.filename,
|
||||
filesize: msg.size,
|
||||
};
|
||||
} catch (e) {
|
||||
if (e.status && e.status === 404) return undefined;
|
||||
throw e;
|
||||
}
|
||||
const data = await window.WWebJS.arrayBufferToBase64Async(
|
||||
await resolved.blob.arrayBuffer(),
|
||||
);
|
||||
return {
|
||||
data,
|
||||
mimetype: resolved.mimetype,
|
||||
filename: resolved.filename,
|
||||
filesize: resolved.filesize,
|
||||
};
|
||||
}, this.id._serialized);
|
||||
|
||||
if (!result) return undefined;
|
||||
|
|
@ -596,6 +538,67 @@ class Message extends Base {
|
|||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Like downloadMedia(), but returns a Readable stream instead of loading the entire file into memory.
|
||||
* @param {Object} [options]
|
||||
* @param {number} [options.chunkSize=10485760] Size in bytes of each chunk read from the browser (default 10MB)
|
||||
* @returns {Promise<MessageMediaStream|undefined>} undefined if media is unavailable
|
||||
*/
|
||||
async downloadMediaStream({ chunkSize = 10 * 1024 * 1024 } = {}) {
|
||||
if (!this.hasMedia) return undefined;
|
||||
|
||||
const blobHandle = await this.client.pupPage.evaluateHandle(
|
||||
async (msgId) => {
|
||||
const result = await window.WWebJS.resolveMediaBlob(msgId);
|
||||
return result?.blob ?? null;
|
||||
},
|
||||
this.id._serialized,
|
||||
);
|
||||
|
||||
let metadata;
|
||||
try {
|
||||
metadata = await blobHandle.evaluate((blob, msgId) => {
|
||||
if (!blob) return null;
|
||||
const msg = window.require('WAWebCollections').Msg.get(msgId);
|
||||
return {
|
||||
blobSize: blob.size,
|
||||
mimetype: msg?.mimetype,
|
||||
filename: msg?.filename,
|
||||
filesize: msg?.size,
|
||||
};
|
||||
}, this.id._serialized);
|
||||
} catch (err) {
|
||||
await blobHandle.dispose().catch(() => {});
|
||||
throw err;
|
||||
}
|
||||
if (!metadata) {
|
||||
await blobHandle.dispose().catch(() => {});
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const { blobSize, ...rest } = metadata;
|
||||
|
||||
async function* readChunks() {
|
||||
try {
|
||||
for (let offset = 0; offset < blobSize; offset += chunkSize) {
|
||||
const base64 = await blobHandle.evaluate(
|
||||
async (blob, s, e) =>
|
||||
window.WWebJS.arrayBufferToBase64Async(
|
||||
await blob.slice(s, e).arrayBuffer(),
|
||||
),
|
||||
offset,
|
||||
offset + chunkSize,
|
||||
);
|
||||
yield Buffer.from(base64, 'base64');
|
||||
}
|
||||
} finally {
|
||||
await blobHandle.dispose().catch(() => {});
|
||||
}
|
||||
}
|
||||
|
||||
return { stream: Readable.from(readChunks()), ...rest };
|
||||
}
|
||||
|
||||
/**
|
||||
* Deletes a message from the chat
|
||||
* @param {?boolean} everyone If true and the message is sent by the current user or the user is an admin, will delete it for everyone in the chat.
|
||||
|
|
|
|||
|
|
@ -1105,6 +1105,63 @@ exports.LoadUtils = () => {
|
|||
});
|
||||
};
|
||||
|
||||
/**
|
||||
* Resolves the media blob and metadata for a message.
|
||||
* Shared by downloadMedia and downloadMediaStream.
|
||||
* @param {string} msgId
|
||||
* @returns {Promise<{blob: Blob, mimetype: string, filename: string, filesize: number}|null>}
|
||||
*/
|
||||
window.WWebJS.resolveMediaBlob = async (msgId) => {
|
||||
const { Msg } = window.require('WAWebCollections');
|
||||
const msg =
|
||||
Msg.get(msgId) ||
|
||||
(await Msg.getMessagesById([msgId]))?.messages?.[0];
|
||||
|
||||
if (
|
||||
!msg ||
|
||||
!msg.mediaData ||
|
||||
msg.mediaData.mediaStage === 'REUPLOADING'
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// Always call internal downloadMedia - never skip based on
|
||||
// mediaStage, because cache eviction can leave stage=RESOLVED
|
||||
// with empty InMemoryMediaBlobCache.
|
||||
await msg.downloadMedia({
|
||||
downloadEvenIfExpensive: true,
|
||||
rmrReason: 1,
|
||||
isUserInitiated: true,
|
||||
});
|
||||
|
||||
if (
|
||||
msg.mediaData.mediaStage.includes('ERROR') ||
|
||||
msg.mediaData.mediaStage === 'FETCHING'
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const cached = window
|
||||
.require('WAWebMediaInMemoryBlobCache')
|
||||
.InMemoryMediaBlobCache.get(msg.mediaObject?.filehash);
|
||||
|
||||
let blob;
|
||||
if (cached) {
|
||||
blob = cached;
|
||||
} else if (msg.mediaObject?.mediaBlob) {
|
||||
blob = msg.mediaObject.mediaBlob.forceToBlob();
|
||||
}
|
||||
|
||||
if (!blob) return null;
|
||||
|
||||
return {
|
||||
blob,
|
||||
mimetype: msg.mimetype,
|
||||
filename: msg.filename,
|
||||
filesize: msg.size,
|
||||
};
|
||||
};
|
||||
|
||||
window.WWebJS.arrayBufferToBase64 = (arrayBuffer) => {
|
||||
let binary = '';
|
||||
const bytes = new Uint8Array(arrayBuffer);
|
||||
|
|
|
|||
Loading…
Reference in a new issue