Make finalization part of the stream

ts/AttachmentCrypto.ts

@@ -21,6 +21,7 @@ import { createName, getRelativePath } from './util/attachmentPath';
 import { appendPaddingStream, logPadSize } from './util/logPadding';
 import { prependStream } from './util/prependStream';
 import { appendMacStream } from './util/appendMacStream';
+import { finalStream } from './util/finalStream';
 import { getIvAndDecipher } from './util/getIvAndDecipher';
 import { getMacAndUpdateHmac } from './util/getMacAndUpdateHmac';
 import { trimPadding } from './util/trimPadding';
@@ -378,28 +379,9 @@ export async function decryptAttachmentV2ToSink(
         }),
         trimPadding(options.size),
         peekAndUpdateHash(plaintextHash),
-        sink,
-      ].filter(isNotNil)
-    );
-  } catch (error) {
-    // These errors happen when canceling fetch from `attachment://` urls,
-    // ignore them to avoid noise in the logs.
-    if (error.name === 'AbortError') {
-      throw error;
-    }
-
-    log.error(
-      `${logId}: Failed to decrypt attachment`,
-      Errors.toLogFormat(error)
-    );
-    throw error;
-  } finally {
-    await readFd?.close();
-  }
-
+        finalStream(() => {
           const ourMac = hmac.digest();
           const ourDigest = digest.digest();
-  const ourPlaintextHash = plaintextHash.digest('hex');
 
           strictAssert(
             ourMac.byteLength === ATTACHMENT_MAC_LENGTH,
@@ -413,10 +395,6 @@ export async function decryptAttachmentV2ToSink(
             ourDigest.byteLength === DIGEST_LENGTH,
             `${logId}: Failed to generate ourDigest!`
           );
-  strictAssert(
-    ourPlaintextHash.length === HEX_DIGEST_LENGTH,
-    `${logId}: Failed to generate file hash!`
-  );
 
           if (!constantTimeEqual(ourMac, theirMac)) {
             throw new Error(`${logId}: Bad MAC`);
@@ -439,12 +417,10 @@ export async function decryptAttachmentV2ToSink(
               throw missingCaseError(type);
           }
 
-  strictAssert(
-    iv != null && iv.byteLength === IV_LENGTH,
-    `${logId}: failed to find their iv`
-  );
+          if (!outerEncryption) {
+            return;
+          }
 
-  if (outerEncryption) {
           strictAssert(outerHmac, 'outerHmac must exist');
 
           const ourOuterMac = outerHmac.digest();
@@ -461,8 +437,37 @@ export async function decryptAttachmentV2ToSink(
           if (!constantTimeEqual(ourOuterMac, theirOuterMac)) {
             throw new Error(`${logId}: Bad outer encryption MAC`);
           }
+        }),
+        sink,
+      ].filter(isNotNil)
+    );
+  } catch (error) {
+    // These errors happen when canceling fetch from `attachment://` urls,
+    // ignore them to avoid noise in the logs.
+    if (error.name === 'AbortError') {
+      throw error;
     }
 
+    log.error(
+      `${logId}: Failed to decrypt attachment`,
+      Errors.toLogFormat(error)
+    );
+    throw error;
+  } finally {
+    await readFd?.close();
+  }
+
+  const ourPlaintextHash = plaintextHash.digest('hex');
+  strictAssert(
+    ourPlaintextHash.length === HEX_DIGEST_LENGTH,
+    `${logId}: Failed to generate file hash!`
+  );
+
+  strictAssert(
+    iv != null && iv.byteLength === IV_LENGTH,
+    `${logId}: failed to find their iv`
+  );
+
   return {
     iv,
     plaintextHash: ourPlaintextHash,

ts/test-node/util/finalStream_test.ts

@@ -0,0 +1,40 @@
+// Copyright 2024 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import { assert } from 'chai';
+import { Readable } from 'node:stream';
+import { pipeline } from 'node:stream/promises';
+
+import { finalStream } from '../../util/finalStream';
+
+describe('finalStream', () => {
+  it('should invoke callback before pipeline resolves', async () => {
+    let called = false;
+    await pipeline(
+      Readable.from(['abc']),
+      finalStream(async () => {
+        // Forcing next tick
+        await Promise.resolve();
+
+        called = true;
+      })
+    );
+
+    assert.isTrue(called);
+  });
+
+  it('should propagate errors from callback', async () => {
+    await assert.isRejected(
+      pipeline(
+        Readable.from(['abc']),
+        finalStream(async () => {
+          // Forcing next tick
+          await Promise.resolve();
+
+          throw new Error('failure');
+        })
+      ),
+      'failure'
+    );
+  });
+});

ts/util/finalStream.ts

@@ -0,0 +1,22 @@
+// Copyright 2024 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import { Transform } from 'node:stream';
+
+export function finalStream(finalizer: () => Promise<void> | void): Transform {
+  return new Transform({
+    transform(data, enc, callback) {
+      this.push(Buffer.from(data, enc));
+      callback();
+    },
+    async final(callback) {
+      try {
+        await finalizer();
+      } catch (error) {
+        callback(error);
+        return;
+      }
+      callback(null);
+    },
+  });
+}