Add cache layer for attachment streaming

ts/test-node/util/ValidatingPassThrough_test.ts

@@ -0,0 +1,72 @@
+// Copyright 2024 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import { assert } from 'chai';
+import { randomBytes } from 'node:crypto';
+import { Readable } from 'node:stream';
+import { pipeline, finished } from 'node:stream/promises';
+import {
+  inferChunkSize,
+  DigestingWritable,
+} from '@signalapp/libsignal-client/dist/incremental_mac';
+
+import { ValidatingPassThrough } from '../../util/ValidatingPassThrough';
+
+// Use uneven chunk size to trigger buffering
+const CHUNK_SIZE = 13579;
+
+function toChunkedReadable(buffer: Buffer): Readable {
+  const chunked = new Array<Buffer>();
+  for (let i = 0; i < buffer.byteLength; i += CHUNK_SIZE) {
+    chunked.push(buffer.subarray(i, i + CHUNK_SIZE));
+  }
+
+  return Readable.from(chunked);
+}
+
+describe('ValidatingPassThrough', () => {
+  it('should emit whole source stream', async () => {
+    const source = randomBytes(10 * 1024 * 1024);
+    const key = randomBytes(32);
+
+    const chunkSize = inferChunkSize(source.byteLength);
+    const writable = new DigestingWritable(key, chunkSize);
+    await pipeline(Readable.from(source), writable);
+
+    const digest = writable.getFinalDigest();
+    const validator = new ValidatingPassThrough(key, chunkSize, digest);
+
+    const received = new Array<Buffer>();
+    validator.on('data', chunk => received.push(chunk));
+
+    await Promise.all([
+      pipeline(toChunkedReadable(source), validator),
+      finished(validator),
+    ]);
+
+    const actual = Buffer.concat(received);
+    assert.isTrue(actual.equals(source));
+  });
+
+  it('should emit error on digest mismatch', async () => {
+    const source = randomBytes(10 * 1024 * 1024);
+    const key = randomBytes(32);
+
+    const chunkSize = inferChunkSize(source.byteLength);
+    const writable = new DigestingWritable(key, chunkSize);
+    await pipeline(Readable.from(source), writable);
+
+    const digest = writable.getFinalDigest();
+    const wrongKey = randomBytes(32);
+    const validator = new ValidatingPassThrough(wrongKey, chunkSize, digest);
+
+    validator.on('data', () => {
+      throw new Error('Should not be called');
+    });
+
+    await assert.isRejected(
+      pipeline(toChunkedReadable(source), validator),
+      'Corrupted input data'
+    );
+  });
+});

ts/util/ValidatingPassThrough.ts

@@ -0,0 +1,87 @@
+// Copyright 2024 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import { noop } from 'lodash';
+import { Transform } from 'node:stream';
+import {
+  ValidatingWritable,
+  type ChunkSizeChoice,
+} from '@signalapp/libsignal-client/dist/incremental_mac';
+
+export class ValidatingPassThrough extends Transform {
+  private validator: ValidatingWritable;
+  private buffer = new Array<Buffer>();
+
+  constructor(key: Buffer, sizeChoice: ChunkSizeChoice, digest: Buffer) {
+    super();
+    this.validator = new ValidatingWritable(key, sizeChoice, digest);
+
+    // We handle errors coming from write/end
+    this.validator.on('error', noop);
+  }
+
+  public override _transform(
+    data: Buffer,
+    enc: BufferEncoding,
+    callback: (error?: null | Error) => void
+  ): void {
+    const start = this.validator.validatedSize();
+    this.validator.write(data, enc, err => {
+      if (err) {
+        return callback(err);
+      }
+
+      this.buffer.push(data);
+
+      const end = this.validator.validatedSize();
+      const readySize = end - start;
+
+      // Fully buffer
+      if (readySize === 0) {
+        return callback(null);
+      }
+
+      const { buffer } = this;
+      this.buffer = [];
+      let validated = 0;
+      for (const chunk of buffer) {
+        validated += chunk.byteLength;
+
+        // Buffered chunk is fully validated - push it without slicing
+        if (validated <= readySize) {
+          this.push(chunk);
+          continue;
+        }
+
+        // Validation boundary lies within the chunk, split it
+        const notValidated = validated - readySize;
+        this.push(chunk.subarray(0, -notValidated));
+        this.buffer.push(chunk.subarray(-notValidated));
+
+        // Technically this chunk must be the last chunk so we could break,
+        // but for consistency keep looping.
+      }
+      callback(null);
+    });
+  }
+
+  public override _final(callback: (error?: null | Error) => void): void {
+    const start = this.validator.validatedSize();
+    this.validator.end((err?: Error) => {
+      if (err) {
+        return callback(err);
+      }
+
+      const end = this.validator.validatedSize();
+      const readySize = end - start;
+      const buffer = Buffer.concat(this.buffer);
+      this.buffer = [];
+      if (buffer.byteLength !== readySize) {
+        return callback(new Error('Stream not fully processed'));
+      }
+      this.push(buffer);
+
+      callback(null);
+    });
+  }
+}