From 9ee41e14d2d173866300e75758468c6788180277 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Tue, 29 Jun 2021 14:21:31 -0400
Subject: [PATCH] fix: prevent withAsyncContext currentInstance leak in edge
 cases

---
 .../__tests__/apiSetupHelpers.spec.ts         | 185 +++++++++++++++---
 packages/runtime-core/src/apiSetupHelpers.ts  |  13 +-
 packages/runtime-core/src/component.ts        |   5 +
 3 files changed, 170 insertions(+), 33 deletions(-)

diff --git a/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts b/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts
index 17d4e298..cf698d7b 100644
--- a/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts
+++ b/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts
@@ -1,11 +1,13 @@
 import {
   ComponentInternalInstance,
+  createApp,
   defineComponent,
   getCurrentInstance,
   h,
   nodeOps,
   onMounted,
   render,
+  serializeInner,
   SetupContext,
   Suspense
 } from '@vue/runtime-test'
@@ -95,38 +97,161 @@ describe('SFC <script setup> helpers', () => {
     ).toHaveBeenWarned()
   })
 
-  test('withAsyncContext', async () => {
-    const spy = jest.fn()
-
-    let beforeInstance: ComponentInternalInstance | null = null
-    let afterInstance: ComponentInternalInstance | null = null
-    let resolve: (msg: string) => void
-
-    const Comp = defineComponent({
-      async setup() {
-        beforeInstance = getCurrentInstance()
-        const msg = await withAsyncContext(
-          new Promise(r => {
-            resolve = r
-          })
-        )
-        // register the lifecycle after an await statement
-        onMounted(spy)
-        afterInstance = getCurrentInstance()
-        return () => msg
-      }
+  describe('withAsyncContext', () => {
+    // disable options API because applyOptions() also resets currentInstance
+    // and we want to ensure the logic works even with Options API disabled.
+    beforeEach(() => {
+      __FEATURE_OPTIONS_API__ = false
     })
 
-    const root = nodeOps.createElement('div')
-    render(h(() => h(Suspense, () => h(Comp))), root)
+    afterEach(() => {
+      __FEATURE_OPTIONS_API__ = true
+    })
 
-    expect(spy).not.toHaveBeenCalled()
-    resolve!('hello')
-    // wait a macro task tick for all micro ticks to resolve
-    await new Promise(r => setTimeout(r))
-    // mount hook should have been called
-    expect(spy).toHaveBeenCalled()
-    // should retain same instance before/after the await call
-    expect(beforeInstance).toBe(afterInstance)
+    test('basic', async () => {
+      const spy = jest.fn()
+
+      let beforeInstance: ComponentInternalInstance | null = null
+      let afterInstance: ComponentInternalInstance | null = null
+      let resolve: (msg: string) => void
+
+      const Comp = defineComponent({
+        async setup() {
+          beforeInstance = getCurrentInstance()
+          const msg = await withAsyncContext(
+            new Promise(r => {
+              resolve = r
+            })
+          )
+          // register the lifecycle after an await statement
+          onMounted(spy)
+          afterInstance = getCurrentInstance()
+          return () => msg
+        }
+      })
+
+      const root = nodeOps.createElement('div')
+      render(h(() => h(Suspense, () => h(Comp))), root)
+
+      expect(spy).not.toHaveBeenCalled()
+      resolve!('hello')
+      // wait a macro task tick for all micro ticks to resolve
+      await new Promise(r => setTimeout(r))
+      // mount hook should have been called
+      expect(spy).toHaveBeenCalled()
+      // should retain same instance before/after the await call
+      expect(beforeInstance).toBe(afterInstance)
+      expect(serializeInner(root)).toBe('hello')
+    })
+
+    test('error handling', async () => {
+      const spy = jest.fn()
+
+      let beforeInstance: ComponentInternalInstance | null = null
+      let afterInstance: ComponentInternalInstance | null = null
+      let reject: () => void
+
+      const Comp = defineComponent({
+        async setup() {
+          beforeInstance = getCurrentInstance()
+          try {
+            await withAsyncContext(
+              new Promise((r, rj) => {
+                reject = rj
+              })
+            )
+          } catch (e) {
+            // ignore
+          }
+          // register the lifecycle after an await statement
+          onMounted(spy)
+          afterInstance = getCurrentInstance()
+          return () => ''
+        }
+      })
+
+      const root = nodeOps.createElement('div')
+      render(h(() => h(Suspense, () => h(Comp))), root)
+
+      expect(spy).not.toHaveBeenCalled()
+      reject!()
+      // wait a macro task tick for all micro ticks to resolve
+      await new Promise(r => setTimeout(r))
+      // mount hook should have been called
+      expect(spy).toHaveBeenCalled()
+      // should retain same instance before/after the await call
+      expect(beforeInstance).toBe(afterInstance)
+    })
+
+    test('should not leak instance on multiple awaits', async () => {
+      let resolve: (val?: any) => void
+      let beforeInstance: ComponentInternalInstance | null = null
+      let afterInstance: ComponentInternalInstance | null = null
+      let inBandInstance: ComponentInternalInstance | null = null
+      let outOfBandInstance: ComponentInternalInstance | null = null
+
+      const ready = new Promise(r => {
+        resolve = r
+      })
+
+      async function doAsyncWork() {
+        // should still have instance
+        inBandInstance = getCurrentInstance()
+        await Promise.resolve()
+        // should not leak instance
+        outOfBandInstance = getCurrentInstance()
+      }
+
+      const Comp = defineComponent({
+        async setup() {
+          beforeInstance = getCurrentInstance()
+          // first await
+          await withAsyncContext(Promise.resolve())
+          // setup exit, instance set to null, then resumed
+          await withAsyncContext(doAsyncWork())
+          afterInstance = getCurrentInstance()
+          return () => {
+            resolve()
+            return ''
+          }
+        }
+      })
+
+      const root = nodeOps.createElement('div')
+      render(h(() => h(Suspense, () => h(Comp))), root)
+
+      await ready
+      expect(inBandInstance).toBe(beforeInstance)
+      expect(outOfBandInstance).toBeNull()
+      expect(afterInstance).toBe(beforeInstance)
+      expect(getCurrentInstance()).toBeNull()
+    })
+
+    test('should not leak on multiple awaits + error', async () => {
+      let resolve: (val?: any) => void
+      const ready = new Promise(r => {
+        resolve = r
+      })
+
+      const Comp = defineComponent({
+        async setup() {
+          await withAsyncContext(Promise.resolve())
+          await withAsyncContext(Promise.reject())
+        },
+        render() {}
+      })
+
+      const app = createApp(() => h(Suspense, () => h(Comp)))
+      app.config.errorHandler = () => {
+        resolve()
+        return false
+      }
+
+      const root = nodeOps.createElement('div')
+      app.mount(root)
+
+      await ready
+      expect(getCurrentInstance()).toBeNull()
+    })
   })
 })
diff --git a/packages/runtime-core/src/apiSetupHelpers.ts b/packages/runtime-core/src/apiSetupHelpers.ts
index b9674143..c0eb9051 100644
--- a/packages/runtime-core/src/apiSetupHelpers.ts
+++ b/packages/runtime-core/src/apiSetupHelpers.ts
@@ -231,13 +231,20 @@ export function mergeDefaults(
 /**
  * Runtime helper for storing and resuming current instance context in
  * async setup().
- * @internal
  */
 export async function withAsyncContext<T>(
   awaitable: T | Promise<T>
 ): Promise<T> {
   const ctx = getCurrentInstance()
-  const res = await awaitable
-  setCurrentInstance(ctx)
+  setCurrentInstance(null) // unset after storing instance
+  if (__DEV__ && !ctx) {
+    warn(`withAsyncContext() called when there is no active context instance.`)
+  }
+  let res: T
+  try {
+    res = await awaitable
+  } finally {
+    setCurrentInstance(ctx)
+  }
   return res
 }
diff --git a/packages/runtime-core/src/component.ts b/packages/runtime-core/src/component.ts
index 01717af2..172e2433 100644
--- a/packages/runtime-core/src/component.ts
+++ b/packages/runtime-core/src/component.ts
@@ -623,6 +623,11 @@ function setupStatefulComponent(
     currentInstance = null
 
     if (isPromise(setupResult)) {
+      const unsetInstance = () => {
+        currentInstance = null
+      }
+      setupResult.then(unsetInstance, unsetInstance)
+
       if (isSSR) {
         // return the promise so server-renderer can wait on it
         return setupResult