From 4063e0bf1322485264e4d117ce6f0f42c237e284 Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna@addaleax.net>
Date: Wed, 1 Dec 2021 15:04:14 +0100
Subject: [PATCH] fix(async-rewriter): ensure finally blocks always run
 MONGOSH-1076

Ensure that `finally` blocks are run even if the `try` block did
not result in an exception.
---
 packages/async-rewriter2/README.md            | 11 +++-
 .../src/async-writer-babel.spec.ts            | 57 +++++++++++++++++++
 .../src/stages/uncatchable-exceptions.ts      | 25 +++++++-
 3 files changed, 89 insertions(+), 4 deletions(-)

diff --git a/packages/async-rewriter2/README.md b/packages/async-rewriter2/README.md
index f83ea9db4e..22adfe11b9 100644
--- a/packages/async-rewriter2/README.md
+++ b/packages/async-rewriter2/README.md
@@ -107,14 +107,21 @@ try {
 into
 
 ```js
-let _isCatchable;
+let _isCatchable = true;
 
 try {
   foo1();
 } catch (err) {
   _isCatchable = !err || !err[Symbol.for('@@mongosh.uncatchable')];
 
-  if (_isCatchable) bar1(err); else throw err;
+  if (_isCatchable) {
+    try {
+      bar1(err);
+    } catch (innerErr) {
+      _isCatchable = !innerErr || !innerErr[Symbol.for('@@mongosh.uncatchable')];
+      throw innerErr;
+    }
+  } else throw err;
 } finally {
   if (_isCatchable) baz();
 }
diff --git a/packages/async-rewriter2/src/async-writer-babel.spec.ts b/packages/async-rewriter2/src/async-writer-babel.spec.ts
index 02e7e259a1..fb0a99343b 100644
--- a/packages/async-rewriter2/src/async-writer-babel.spec.ts
+++ b/packages/async-rewriter2/src/async-writer-babel.spec.ts
@@ -925,6 +925,33 @@ describe('AsyncWriter', () => {
       expect(result).to.equal('finally');
     });
 
+    it('allows returning from finally when exception was thrown in catch block', () => {
+      const result = runTranspiledCode(`
+      (() => {
+        try {
+          throw new Error('first error');
+        } catch (err) {
+          throw new Error('second error');
+        } finally {
+          return 'finally';
+        }
+      })();`);
+      expect(result).to.equal('finally');
+    });
+
+    it('allows returning from finally when no exception is thrown', () => {
+      const result = runTranspiledCode(`
+      (() => {
+        try {
+        } catch (err) {
+          return 'catch';
+        } finally {
+          return 'finally';
+        }
+      })();`);
+      expect(result).to.equal('finally');
+    });
+
     it('allows finally without catch', () => {
       const result = runTranspiledCode(`
       (() => {
@@ -937,6 +964,18 @@ describe('AsyncWriter', () => {
       expect(result).to.equal('finally');
     });
 
+    it('allows finally without catch with return from try block', () => {
+      const result = runTranspiledCode(`
+      (() => {
+        try {
+          return 'try';
+        } finally {
+          return 'finally';
+        }
+      })();`);
+      expect(result).to.equal('finally');
+    });
+
     it('allows throwing primitives', () => {
       const result = runTranspiledCode(`
       (() => {
@@ -1020,5 +1059,23 @@ describe('AsyncWriter', () => {
         expect(err.message).to.equal('uncatchable!');
       }
     });
+
+    it('does not catch uncatchable exceptions thrown from catch block', () => {
+      try {
+        runTranspiledCode(`
+        (() => {
+          try {
+            throw new Error('regular error');
+          } catch {
+            throwUncatchable();
+          } finally {
+            return;
+          }
+        })();`);
+        expect.fail('missed exception');
+      } catch (err) {
+        expect(err.message).to.equal('uncatchable!');
+      }
+    });
   });
 });
diff --git a/packages/async-rewriter2/src/stages/uncatchable-exceptions.ts b/packages/async-rewriter2/src/stages/uncatchable-exceptions.ts
index 480795947f..07e0b28bdf 100644
--- a/packages/async-rewriter2/src/stages/uncatchable-exceptions.ts
+++ b/packages/async-rewriter2/src/stages/uncatchable-exceptions.ts
@@ -75,8 +75,10 @@ export default ({ types: t }: { types: typeof BabelTypes }): babel.PluginObj<{}>
           // (i.e. whether the finalizer is allowed to run) outside of the
           // try/catch/finally block.
           const isCatchable = path.scope.generateUidIdentifier('_isCatchable');
+          const exceptionFromCatchIdentifier = path.scope.generateUidIdentifier('_innerExc');
           path.replaceWithMultiple([
-            t.variableDeclaration('let', [t.variableDeclarator(isCatchable)]),
+            // let isCatchable = true /* for the case in which no exception is thrown */
+            t.variableDeclaration('let', [t.variableDeclarator(isCatchable, t.booleanLiteral(true))]),
             Object.assign(
               t.tryStatement(
                 block,
@@ -89,7 +91,26 @@ export default ({ types: t }: { types: typeof BabelTypes }): babel.PluginObj<{}>
                         isCatchable,
                         notUncatchableCheck({ ERR_IDENTIFIER: catchParam }))),
                     // if (isCatchable) { ... } else throw err;
-                    t.ifStatement(isCatchable, handler.body, t.throwStatement(catchParam))
+                    t.ifStatement(
+                      isCatchable,
+                      // try/catch around the catch body so we know whether finally should run here
+                      Object.assign(t.tryStatement(
+                        handler.body,
+                        // catch (err) {
+                        t.catchClause(
+                          exceptionFromCatchIdentifier,
+                          t.blockStatement([
+                            // isCatchable = !err[isUncatchableSymbol]
+                            t.expressionStatement(
+                              t.assignmentExpression('=',
+                                isCatchable,
+                                notUncatchableCheck({ ERR_IDENTIFIER: exceptionFromCatchIdentifier }))),
+                            // throw err;
+                            t.throwStatement(exceptionFromCatchIdentifier)
+                          ])
+                        )
+                      ), { [isGeneratedTryCatch]: true }),
+                      t.throwStatement(catchParam))
                   ]),
                 ),
                 t.blockStatement([