Resumable attachment downloads

This commit is contained in:
Fedor Indutny 2024-08-19 13:05:35 -07:00 committed by GitHub
parent 2c92591b59
commit 38f532cdda
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 401 additions and 89 deletions

View file

@ -1,9 +1,10 @@
// Copyright 2020 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
import { createWriteStream } from 'fs';
import { createWriteStream } from 'node:fs';
import { stat } from 'node:fs/promises';
import { isNumber } from 'lodash';
import type { Readable } from 'stream';
import type { Readable, Writable } from 'stream';
import { Transform } from 'stream';
import { pipeline } from 'stream/promises';
import { ensureFile } from 'fs-extra';
@ -24,7 +25,7 @@ import {
} from '../Crypto';
import {
getAttachmentCiphertextLength,
safeUnlinkSync,
safeUnlink,
splitKeys,
type ReencryptedAttachmentV2,
decryptAndReencryptLocally,
@ -124,6 +125,39 @@ export async function downloadAttachment(
let downloadResult: Awaited<ReturnType<typeof downloadToDisk>>;
let { downloadPath } = attachment;
let downloadOffset = 0;
if (downloadPath) {
const absoluteDownloadPath =
window.Signal.Migrations.getAbsoluteAttachmentPath(downloadPath);
try {
({ size: downloadOffset } = await stat(absoluteDownloadPath));
} catch (error) {
if (error.code !== 'ENOENT') {
log.error(
'downloadAttachment: Failed to get file size for previous download',
Errors.toLogFormat(error)
);
try {
await safeUnlink(downloadPath);
} catch {
downloadPath = undefined;
}
}
}
}
// Start over if we go over the size
if (downloadOffset >= size && downloadPath) {
log.warn('downloadAttachment: went over, retrying');
await safeUnlink(downloadPath);
downloadOffset = 0;
}
if (downloadOffset !== 0) {
log.info(`${logId}: resuming from ${downloadOffset}`);
}
if (mediaTier === MediaTier.STANDARD) {
strictAssert(
options.variant !== AttachmentVariant.ThumbnailFromBackup,
@ -135,9 +169,17 @@ export async function downloadAttachment(
const downloadStream = await server.getAttachment({
cdnKey,
cdnNumber,
options,
options: {
...options,
downloadOffset,
},
});
downloadResult = await downloadToDisk({
downloadStream,
size,
downloadPath,
downloadOffset,
});
downloadResult = await downloadToDisk({ downloadStream, size });
} else {
const mediaId =
options.variant === AttachmentVariant.ThumbnailFromBackup
@ -157,10 +199,15 @@ export async function downloadAttachment(
mediaDir,
headers: cdnCredentials.headers,
cdnNumber,
options,
options: {
...options,
downloadOffset,
},
});
downloadResult = await downloadToDisk({
downloadStream,
downloadPath,
downloadOffset,
size: getAttachmentCiphertextLength(
options.variant === AttachmentVariant.ThumbnailFromBackup
? // be generous, accept downloads of up to twice what we expect for thumbnail
@ -170,10 +217,7 @@ export async function downloadAttachment(
});
}
const { relativePath: downloadedRelativePath, downloadSize } = downloadResult;
const cipherTextAbsolutePath =
window.Signal.Migrations.getAbsoluteAttachmentPath(downloadedRelativePath);
const { absolutePath: cipherTextAbsolutePath, downloadSize } = downloadResult;
try {
switch (options.variant) {
@ -226,23 +270,42 @@ export async function downloadAttachment(
}
}
} finally {
safeUnlinkSync(cipherTextAbsolutePath);
await safeUnlink(cipherTextAbsolutePath);
}
}
async function downloadToDisk({
downloadStream,
downloadPath,
downloadOffset = 0,
size,
}: {
downloadStream: Readable;
downloadPath?: string;
downloadOffset?: number;
size: number;
}): Promise<{ relativePath: string; downloadSize: number }> {
const relativeTargetPath = getRelativePath(createName());
const absoluteTargetPath =
window.Signal.Migrations.getAbsoluteAttachmentPath(relativeTargetPath);
}): Promise<{ absolutePath: string; downloadSize: number }> {
const absoluteTargetPath = downloadPath
? window.Signal.Migrations.getAbsoluteDownloadsPath(downloadPath)
: window.Signal.Migrations.getAbsoluteAttachmentPath(
getRelativePath(createName())
);
await ensureFile(absoluteTargetPath);
const writeStream = createWriteStream(absoluteTargetPath);
const targetSize = getAttachmentCiphertextLength(size);
let writeStream: Writable;
if (downloadPath) {
writeStream = createWriteStream(absoluteTargetPath, {
flags: 'a',
start: downloadOffset,
});
} else {
strictAssert(
!downloadOffset,
'Download cannot be resumed without downloadPath'
);
writeStream = createWriteStream(absoluteTargetPath);
}
const targetSize = getAttachmentCiphertextLength(size) - downloadOffset;
let downloadSize = 0;
try {
@ -255,19 +318,23 @@ async function downloadToDisk({
writeStream
);
} catch (error) {
try {
safeUnlinkSync(absoluteTargetPath);
} catch (cleanupError) {
log.error(
'downloadToDisk: Error while cleaning up',
Errors.toLogFormat(cleanupError)
);
if (downloadPath) {
log.warn(`downloadToDisk: stopping at ${downloadSize}`);
} else {
try {
await safeUnlink(absoluteTargetPath);
} catch (cleanupError) {
log.error(
'downloadToDisk: Error while cleaning up',
Errors.toLogFormat(cleanupError)
);
}
}
throw error;
}
return { relativePath: relativeTargetPath, downloadSize };
return { absolutePath: absoluteTargetPath, downloadSize };
}
// A simple transform that throws if it sees more than maxBytes on the stream.