transformWithCancel: fix race condition in cancellation (#43)

Relevant for e.g.passiveClones.
openpgpjs · Jun 18, 2024 · 92f34fd · 92f34fd
1 parent 0ce871e
commit 92f34fd
Showing 1 changed file with 18 additions and 4 deletions.
diff --git a/lib/streams.js b/lib/streams.js
@@ -165,9 +165,10 @@ function transformRaw(input, options) {
  * @param {Function} cancel
  * @returns {TransformStream}
  */
-function transformWithCancel(cancel) {
+function transformWithCancel(customCancel) {
   let pulled = false;
-  let backpressureChangePromiseResolve;
+  let cancelled = false;
+  let backpressureChangePromiseResolve, backpressureChangePromiseReject;
   let outputController;
   return {
     readable: new ReadableStream({
@@ -181,16 +182,29 @@ function transformWithCancel(cancel) {
           pulled = true;
         }
       },
-      cancel
+      async cancel(reason) {
+        cancelled = true;
+        if (customCancel) {
+          await customCancel(reason);
+        }
+        if (backpressureChangePromiseReject) {
+          backpressureChangePromiseReject(reason);
+        }
+      }
     }, {highWaterMark: 0}),
     writable: new WritableStream({
       write: async function(chunk) {
+        if (cancelled) {
+          throw new Error('Stream is cancelled');
+        }
         outputController.enqueue(chunk);
         if (!pulled) {
-          await new Promise(resolve => {
+          await new Promise((resolve, reject) => {
             backpressureChangePromiseResolve = resolve;
+            backpressureChangePromiseReject = reject;
           });
           backpressureChangePromiseResolve = null;
+          backpressureChangePromiseReject = null;
         } else {
           pulled = false;
         }