From 0945dbe409acf868ea428f7483dbfb0ef09c8894 Mon Sep 17 00:00:00 2001
From: Thorsten Luenborg <t.luenborg@googlemail.com>
Date: Sat, 4 Dec 2021 12:45:37 +0100
Subject: [PATCH] feat(runtime-core): allow inject() in custom directives

Set currentInstance before hook calls

close: #5002
---
 .../runtime-core/__tests__/directives.spec.ts | 30 ++++++++++++++++++-
 packages/runtime-core/src/directives.ts       |  9 +++++-
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/packages/runtime-core/__tests__/directives.spec.ts b/packages/runtime-core/__tests__/directives.spec.ts
index 02d2471180..ae0ed54d95 100644
--- a/packages/runtime-core/__tests__/directives.spec.ts
+++ b/packages/runtime-core/__tests__/directives.spec.ts
@@ -7,7 +7,10 @@ import {
   DirectiveHook,
   VNode,
   DirectiveBinding,
-  nextTick
+  nextTick,
+  Directive,
+  ComponentPublicInstance,
+  getCurrentInstance
 } from '@vue/runtime-test'
 import { currentInstance, ComponentInternalInstance } from '../src/component'
 
@@ -395,4 +398,29 @@ describe('directives', () => {
     expect(beforeUpdate).toHaveBeenCalledTimes(1)
     expect(count.value).toBe(1)
   })
+
+  test('should have currentInstance available', async () => {
+    let instance: ComponentInternalInstance | null
+    let bindingInstance: ComponentPublicInstance | null
+    const beforeMount: Directive = (_, binding) => {
+      instance = getCurrentInstance()
+      bindingInstance = binding.instance
+    }
+    const App = {
+      render() {
+        return withDirectives(h('p', 'Test'), [
+          [
+            {
+              beforeMount
+            }
+          ]
+        ])
+      }
+    }
+
+    const root = nodeOps.createElement('div')
+    render(h(App), root)
+    expect(instance!).not.toBe(null)
+    expect(instance!.proxy).toBe(bindingInstance!)
+  })
 })
diff --git a/packages/runtime-core/src/directives.ts b/packages/runtime-core/src/directives.ts
index ec13e951bc..538dcd3612 100644
--- a/packages/runtime-core/src/directives.ts
+++ b/packages/runtime-core/src/directives.ts
@@ -14,7 +14,12 @@ return withDirectives(h(comp), [
 import { VNode } from './vnode'
 import { isFunction, EMPTY_OBJ, makeMap } from '@vue/shared'
 import { warn } from './warning'
-import { ComponentInternalInstance, Data } from './component'
+import {
+  ComponentInternalInstance,
+  Data,
+  setCurrentInstance,
+  unsetCurrentInstance
+} from './component'
 import { currentRenderingInstance } from './componentRenderContext'
 import { callWithAsyncErrorHandling, ErrorCodes } from './errorHandling'
 import { ComponentPublicInstance } from './componentPublicInstance'
@@ -139,12 +144,14 @@ export function invokeDirectiveHook(
       // disable tracking inside all lifecycle hooks
       // since they can potentially be called inside effects.
       pauseTracking()
+      instance && setCurrentInstance(instance)
       callWithAsyncErrorHandling(hook, instance, ErrorCodes.DIRECTIVE_HOOK, [
         vnode.el,
         binding,
         vnode,
         prevVNode
       ])
+      unsetCurrentInstance()
       resetTracking()
     }
   }
-- 
2.47.2