Use streams to download attachments directly to disk

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

ts/test-electron/ContactsParser_test.ts

@@ -0,0 +1,277 @@
+// Copyright 2015 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import { assert } from 'chai';
+import { createReadStream, readFileSync, unlinkSync, writeFileSync } from 'fs';
+import { v4 as generateGuid } from 'uuid';
+import { join } from 'path';
+import { pipeline } from 'stream/promises';
+import { Transform } from 'stream';
+
+import protobuf from '../protobuf/wrap';
+import * as log from '../logging/log';
+import * as Bytes from '../Bytes';
+import * as Errors from '../types/errors';
+import { SignalService as Proto } from '../protobuf';
+import {
+  ParseContactsTransform,
+  parseContactsV2,
+} from '../textsecure/ContactsParser';
+import type { ContactDetailsWithAvatar } from '../textsecure/ContactsParser';
+import { createTempDir, deleteTempDir } from '../updater/common';
+import { strictAssert } from '../util/assert';
+
+const { Writer } = protobuf;
+
+describe('ContactsParser', () => {
+  let tempDir: string;
+
+  beforeEach(async () => {
+    tempDir = await createTempDir();
+  });
+  afterEach(async () => {
+    await deleteTempDir(log, tempDir);
+  });
+
+  describe('parseContactsV2', () => {
+    it('parses an array buffer of contacts', async () => {
+      let absolutePath: string | undefined;
+
+      try {
+        const bytes = getTestBuffer();
+        const fileName = generateGuid();
+        absolutePath = join(tempDir, fileName);
+        writeFileSync(absolutePath, bytes);
+
+        const contacts = await parseContactsV2({ absolutePath });
+        assert.strictEqual(contacts.length, 3);
+
+        contacts.forEach(contact => {
+          verifyContact(contact);
+        });
+      } finally {
+        if (absolutePath) {
+          unlinkSync(absolutePath);
+        }
+      }
+    });
+
+    it('parses an array buffer of contacts with small chunk size', async () => {
+      let absolutePath: string | undefined;
+
+      try {
+        const bytes = getTestBuffer();
+        const fileName = generateGuid();
+        absolutePath = join(tempDir, fileName);
+        writeFileSync(absolutePath, bytes);
+
+        const contacts = await parseContactsWithSmallChunkSize({
+          absolutePath,
+        });
+        assert.strictEqual(contacts.length, 3);
+
+        contacts.forEach(contact => {
+          verifyContact(contact);
+        });
+      } finally {
+        if (absolutePath) {
+          unlinkSync(absolutePath);
+        }
+      }
+    });
+
+    it('parses an array buffer of contacts where one contact has no avatar', async () => {
+      let absolutePath: string | undefined;
+
+      try {
+        const bytes = Bytes.concatenate([
+          generatePrefixedContact(undefined),
+          getTestBuffer(),
+        ]);
+
+        const fileName = generateGuid();
+        absolutePath = join(tempDir, fileName);
+        writeFileSync(absolutePath, bytes);
+
+        const contacts = await parseContactsWithSmallChunkSize({
+          absolutePath,
+        });
+        assert.strictEqual(contacts.length, 4);
+
+        contacts.forEach((contact, index) => {
+          const avatarIsMissing = index === 0;
+          verifyContact(contact, avatarIsMissing);
+        });
+      } finally {
+        if (absolutePath) {
+          unlinkSync(absolutePath);
+        }
+      }
+    });
+
+    it('parses an array buffer of contacts where contacts are dropped due to missing ACI', async () => {
+      let absolutePath: string | undefined;
+
+      try {
+        const avatarBuffer = generateAvatar();
+        const bytes = Bytes.concatenate([
+          generatePrefixedContact(avatarBuffer, 'invalid'),
+          avatarBuffer,
+          generatePrefixedContact(undefined, 'invalid'),
+          getTestBuffer(),
+        ]);
+
+        const fileName = generateGuid();
+        absolutePath = join(tempDir, fileName);
+        writeFileSync(absolutePath, bytes);
+
+        const contacts = await parseContactsWithSmallChunkSize({
+          absolutePath,
+        });
+        assert.strictEqual(contacts.length, 3);
+
+        contacts.forEach(contact => {
+          verifyContact(contact);
+        });
+      } finally {
+        if (absolutePath) {
+          unlinkSync(absolutePath);
+        }
+      }
+    });
+  });
+});
+
+class SmallChunksTransform extends Transform {
+  constructor(private chunkSize: number) {
+    super();
+  }
+
+  override _transform(
+    incomingChunk: Buffer | undefined,
+    _encoding: string,
+    done: (error?: Error) => void
+  ) {
+    if (!incomingChunk || incomingChunk.byteLength === 0) {
+      done();
+      return;
+    }
+
+    try {
+      const totalSize = incomingChunk.byteLength;
+
+      const chunkCount = Math.floor(totalSize / this.chunkSize);
+      const remainder = totalSize % this.chunkSize;
+
+      for (let i = 0; i < chunkCount; i += 1) {
+        const start = i * this.chunkSize;
+        const end = start + this.chunkSize;
+        this.push(incomingChunk.subarray(start, end));
+      }
+      if (remainder > 0) {
+        this.push(incomingChunk.subarray(chunkCount * this.chunkSize));
+      }
+    } catch (error) {
+      done(error);
+      return;
+    }
+
+    done();
+  }
+}
+
+function generateAvatar(): Uint8Array {
+  const result = new Uint8Array(255);
+  for (let i = 0; i < result.length; i += 1) {
+    result[i] = i;
+  }
+  return result;
+}
+
+function getTestBuffer(): Uint8Array {
+  const avatarBuffer = generateAvatar();
+  const prefixedContact = generatePrefixedContact(avatarBuffer);
+
+  const chunks: Array<Uint8Array> = [];
+  for (let i = 0; i < 3; i += 1) {
+    chunks.push(prefixedContact);
+    chunks.push(avatarBuffer);
+  }
+
+  return Bytes.concatenate(chunks);
+}
+
+function generatePrefixedContact(
+  avatarBuffer: Uint8Array | undefined,
+  aci = '7198E1BD-1293-452A-A098-F982FF201902'
+) {
+  const contactInfoBuffer = Proto.ContactDetails.encode({
+    name: 'Zero Cool',
+    number: '+10000000000',
+    aci,
+    avatar: avatarBuffer
+      ? { contentType: 'image/jpeg', length: avatarBuffer.length }
+      : undefined,
+  }).finish();
+
+  const writer = new Writer();
+  writer.bytes(contactInfoBuffer);
+  const prefixedContact = writer.finish();
+  return prefixedContact;
+}
+
+function verifyContact(
+  contact: ContactDetailsWithAvatar,
+  avatarIsMissing?: boolean
+) {
+  assert.strictEqual(contact.name, 'Zero Cool');
+  assert.strictEqual(contact.number, '+10000000000');
+  assert.strictEqual(contact.aci, '7198e1bd-1293-452a-a098-f982ff201902');
+
+  if (avatarIsMissing) {
+    return;
+  }
+
+  const path = contact.avatar?.path;
+  strictAssert(path, 'Avatar needs path');
+
+  const absoluteAttachmentPath =
+    window.Signal.Migrations.getAbsoluteAttachmentPath(path);
+  const avatarBytes = readFileSync(absoluteAttachmentPath);
+  unlinkSync(absoluteAttachmentPath);
+
+  for (let j = 0; j < 255; j += 1) {
+    assert.strictEqual(avatarBytes[j], j);
+  }
+}
+
+async function parseContactsWithSmallChunkSize({
+  absolutePath,
+}: {
+  absolutePath: string;
+}): Promise<ReadonlyArray<ContactDetailsWithAvatar>> {
+  const logId = 'parseContactsWithSmallChunkSize';
+
+  const readStream = createReadStream(absolutePath);
+  const smallChunksTransform = new SmallChunksTransform(32);
+  const parseContactsTransform = new ParseContactsTransform();
+
+  try {
+    await pipeline(readStream, smallChunksTransform, parseContactsTransform);
+  } catch (error) {
+    try {
+      readStream.close();
+    } catch (cleanupError) {
+      log.error(
+        `${logId}: Failed to clean up after error`,
+        Errors.toLogFormat(cleanupError)
+      );
+    }
+
+    throw error;
+  }
+
+  readStream.close();
+
+  return parseContactsTransform.contacts;
+}