From 86b4da1ec2521ad1473f780ceae0a44d586b1842 Mon Sep 17 00:00:00 2001
From: Fedor Indutny <79877362+indutny-signal@users.noreply.github.com>
Date: Fri, 19 Jul 2024 19:17:02 -0700
Subject: [PATCH] Make finalization part of the stream

---
 ts/AttachmentCrypto.ts                | 115 ++++++++++++++------------
 ts/test-node/util/finalStream_test.ts |  40 +++++++++
 ts/util/finalStream.ts                |  22 +++++
 3 files changed, 122 insertions(+), 55 deletions(-)
 create mode 100644 ts/test-node/util/finalStream_test.ts
 create mode 100644 ts/util/finalStream.ts

diff --git a/ts/AttachmentCrypto.ts b/ts/AttachmentCrypto.ts
index 2d63e43c62c..ee84e7f316e 100644
--- a/ts/AttachmentCrypto.ts
+++ b/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,6 +379,65 @@ export async function decryptAttachmentV2ToSink(
         }),
         trimPadding(options.size),
         peekAndUpdateHash(plaintextHash),
+        finalStream(() => {
+          const ourMac = hmac.digest();
+          const ourDigest = digest.digest();
+
+          strictAssert(
+            ourMac.byteLength === ATTACHMENT_MAC_LENGTH,
+            `${logId}: Failed to generate ourMac!`
+          );
+          strictAssert(
+            theirMac != null && theirMac.byteLength === ATTACHMENT_MAC_LENGTH,
+            `${logId}: Failed to find theirMac!`
+          );
+          strictAssert(
+            ourDigest.byteLength === DIGEST_LENGTH,
+            `${logId}: Failed to generate ourDigest!`
+          );
+
+          if (!constantTimeEqual(ourMac, theirMac)) {
+            throw new Error(`${logId}: Bad MAC`);
+          }
+
+          const { type } = options;
+          switch (type) {
+            case 'local':
+            case 'backupThumbnail':
+              log.info(
+                `${logId}: skipping digest check since this is a ${type} attachment`
+              );
+              break;
+            case 'standard':
+              if (!constantTimeEqual(ourDigest, options.theirDigest)) {
+                throw new Error(`${logId}: Bad digest`);
+              }
+              break;
+            default:
+              throw missingCaseError(type);
+          }
+
+          if (!outerEncryption) {
+            return;
+          }
+
+          strictAssert(outerHmac, 'outerHmac must exist');
+
+          const ourOuterMac = outerHmac.digest();
+          strictAssert(
+            ourOuterMac.byteLength === ATTACHMENT_MAC_LENGTH,
+            `${logId}: Failed to generate ourOuterMac!`
+          );
+          strictAssert(
+            theirOuterMac != null &&
+              theirOuterMac.byteLength === ATTACHMENT_MAC_LENGTH,
+            `${logId}: Failed to find theirOuterMac!`
+          );
+
+          if (!constantTimeEqual(ourOuterMac, theirOuterMac)) {
+            throw new Error(`${logId}: Bad outer encryption MAC`);
+          }
+        }),
         sink,
       ].filter(isNotNil)
     );
@@ -397,72 +457,17 @@ export async function decryptAttachmentV2ToSink(
     await readFd?.close();
   }
 
-  const ourMac = hmac.digest();
-  const ourDigest = digest.digest();
   const ourPlaintextHash = plaintextHash.digest('hex');
-
-  strictAssert(
-    ourMac.byteLength === ATTACHMENT_MAC_LENGTH,
-    `${logId}: Failed to generate ourMac!`
-  );
-  strictAssert(
-    theirMac != null && theirMac.byteLength === ATTACHMENT_MAC_LENGTH,
-    `${logId}: Failed to find theirMac!`
-  );
-  strictAssert(
-    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`);
-  }
-
-  const { type } = options;
-  switch (type) {
-    case 'local':
-    case 'backupThumbnail':
-      log.info(
-        `${logId}: skipping digest check since this is a ${type} attachment`
-      );
-      break;
-    case 'standard':
-      if (!constantTimeEqual(ourDigest, options.theirDigest)) {
-        throw new Error(`${logId}: Bad digest`);
-      }
-      break;
-    default:
-      throw missingCaseError(type);
-  }
-
   strictAssert(
     iv != null && iv.byteLength === IV_LENGTH,
     `${logId}: failed to find their iv`
   );
 
-  if (outerEncryption) {
-    strictAssert(outerHmac, 'outerHmac must exist');
-
-    const ourOuterMac = outerHmac.digest();
-    strictAssert(
-      ourOuterMac.byteLength === ATTACHMENT_MAC_LENGTH,
-      `${logId}: Failed to generate ourOuterMac!`
-    );
-    strictAssert(
-      theirOuterMac != null &&
-        theirOuterMac.byteLength === ATTACHMENT_MAC_LENGTH,
-      `${logId}: Failed to find theirOuterMac!`
-    );
-
-    if (!constantTimeEqual(ourOuterMac, theirOuterMac)) {
-      throw new Error(`${logId}: Bad outer encryption MAC`);
-    }
-  }
-
   return {
     iv,
     plaintextHash: ourPlaintextHash,
diff --git a/ts/test-node/util/finalStream_test.ts b/ts/test-node/util/finalStream_test.ts
new file mode 100644
index 00000000000..2c0035c7269
--- /dev/null
+++ b/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'
+    );
+  });
+});
diff --git a/ts/util/finalStream.ts b/ts/util/finalStream.ts
new file mode 100644
index 00000000000..e89e21b9038
--- /dev/null
+++ b/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);
+    },
+  });
+}