Use streams to download attachments directly to disk

Co-authored-by: trevor-signal <131492920+trevor-signal@users.noreply.github.com>

ts/textsecure/downloadAttachment.ts

@@ -1,19 +1,40 @@
 // Copyright 2020 Signal Messenger, LLC
 // SPDX-License-Identifier: AGPL-3.0-only
 
-import { isNumber } from 'lodash';
+import { createWriteStream, existsSync, unlinkSync } from 'fs';
+import { isNumber, omit } from 'lodash';
+import type { Readable } from 'stream';
+import { Transform } from 'stream';
+import { pipeline } from 'stream/promises';
+import { ensureFile } from 'fs-extra';
 
+import * as log from '../logging/log';
+import * as Errors from '../types/errors';
 import { strictAssert } from '../util/assert';
 import { dropNull } from '../util/dropNull';
-import type { DownloadedAttachmentType } from '../types/Attachment';
+import {
+  AttachmentSizeError,
+  type AttachmentType,
+  type DownloadedAttachmentType,
+} from '../types/Attachment';
 import * as MIME from '../types/MIME';
 import * as Bytes from '../Bytes';
-import { getFirstBytes, decryptAttachment } from '../Crypto';
+import {
+  getFirstBytes,
+  decryptAttachmentV1,
+  getAttachmentSizeBucket,
+} from '../Crypto';
+import {
+  decryptAttachmentV2,
+  IV_LENGTH,
+  ATTACHMENT_MAC_LENGTH,
+} from '../AttachmentCrypto';
 
 import type { ProcessedAttachment } from './Types.d';
 import type { WebAPIType } from './WebAPI';
+import { createName, getRelativePath } from '../windows/attachments';
 
-export async function downloadAttachment(
+export async function downloadAttachmentV1(
   server: WebAPIType,
   attachment: ProcessedAttachment,
   options?: {
@@ -28,7 +49,6 @@ export async function downloadAttachment(
     throw new Error('downloadAttachment: Attachment was missing cdnId!');
   }
 
-  strictAssert(cdnId, 'attachment without cdnId');
   const encrypted = await server.getAttachment(
     cdnId,
     dropNull(cdnNumber),
@@ -41,9 +61,8 @@ export async function downloadAttachment(
   }
 
   strictAssert(key, 'attachment has no key');
-  strictAssert(digest, 'attachment has no digest');
 
-  const paddedData = decryptAttachment(
+  const paddedData = decryptAttachmentV1(
     encrypted,
     Bytes.fromBase64(key),
     Bytes.fromBase64(digest)
@@ -67,3 +86,132 @@ export async function downloadAttachment(
     data,
   };
 }
+
+export async function downloadAttachmentV2(
+  server: WebAPIType,
+  attachment: ProcessedAttachment,
+  options?: {
+    disableRetries?: boolean;
+    timeout?: number;
+  }
+): Promise<AttachmentType> {
+  const { cdnId, cdnKey, cdnNumber, contentType, digest, key, size } =
+    attachment;
+
+  const cdn = cdnId || cdnKey;
+  const logId = `downloadAttachmentV2(${cdn}):`;
+
+  strictAssert(cdn, `${logId}: missing cdnId or cdnKey`);
+  strictAssert(digest, `${logId}: missing digest`);
+  strictAssert(key, `${logId}: missing key`);
+  strictAssert(isNumber(size), `${logId}: missing size`);
+
+  const downloadStream = await server.getAttachmentV2(
+    cdn,
+    dropNull(cdnNumber),
+    options
+  );
+
+  const cipherTextRelativePath = await downloadToDisk({ downloadStream, size });
+  const cipherTextAbsolutePath =
+    window.Signal.Migrations.getAbsoluteAttachmentPath(cipherTextRelativePath);
+
+  const relativePath = await decryptAttachmentV2({
+    ciphertextPath: cipherTextAbsolutePath,
+    id: cdn,
+    keys: Bytes.fromBase64(key),
+    size,
+    theirDigest: Bytes.fromBase64(digest),
+  });
+
+  if (existsSync(cipherTextAbsolutePath)) {
+    unlinkSync(cipherTextAbsolutePath);
+  }
+
+  return {
+    ...omit(attachment, 'key'),
+    path: relativePath,
+    size,
+    contentType: contentType
+      ? MIME.stringToMIMEType(contentType)
+      : MIME.APPLICATION_OCTET_STREAM,
+  };
+}
+
+async function downloadToDisk({
+  downloadStream,
+  size,
+}: {
+  downloadStream: Readable;
+  size: number;
+}): Promise<string> {
+  const relativeTargetPath = getRelativePath(createName());
+  const absoluteTargetPath =
+    window.Signal.Migrations.getAbsoluteAttachmentPath(relativeTargetPath);
+  await ensureFile(absoluteTargetPath);
+  const writeStream = createWriteStream(absoluteTargetPath);
+
+  const targetSize =
+    getAttachmentSizeBucket(size) * 1.05 + IV_LENGTH + ATTACHMENT_MAC_LENGTH;
+  const checkSizeTransform = new CheckSizeTransform(targetSize);
+
+  try {
+    await pipeline(downloadStream, checkSizeTransform, writeStream);
+  } catch (error) {
+    try {
+      writeStream.close();
+      if (absoluteTargetPath && existsSync(absoluteTargetPath)) {
+        unlinkSync(absoluteTargetPath);
+      }
+    } catch (cleanupError) {
+      log.error(
+        'downloadToDisk: Error while cleaning up',
+        Errors.toLogFormat(cleanupError)
+      );
+    }
+
+    throw error;
+  }
+
+  return relativeTargetPath;
+}
+
+// A simple transform that throws if it sees more than maxBytes on the stream.
+class CheckSizeTransform extends Transform {
+  private bytesSeen = 0;
+
+  constructor(private maxBytes: number) {
+    super();
+  }
+
+  override _transform(
+    chunk: Buffer | undefined,
+    _encoding: string,
+    done: (error?: Error) => void
+  ) {
+    if (!chunk || chunk.byteLength === 0) {
+      done();
+      return;
+    }
+
+    try {
+      this.bytesSeen += chunk.byteLength;
+
+      if (this.bytesSeen > this.maxBytes) {
+        done(
+          new AttachmentSizeError(
+            `CheckSizeTransform: Saw ${this.bytesSeen} bytes, max is ${this.maxBytes} bytes`
+          )
+        );
+        return;
+      }
+
+      this.push(chunk);
+    } catch (error) {
+      done(error);
+      return;
+    }
+
+    done();
+  }
+}