From b432c629621405b3c7729c253b5920e7b90554dd Mon Sep 17 00:00:00 2001
From: Eric Jizba <erijiz@microsoft.com>
Date: Mon, 28 Mar 2022 18:49:10 -0700
Subject: [PATCH] Add functionCallback to preInvocationContext

---
 src/eventHandlers/invocationRequest.ts       | 11 +++++++++--
 test/eventHandlers/invocationRequest.test.ts | 19 +++++++++++++++++++
 types-core/index.d.ts                        |  7 ++++++-
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/eventHandlers/invocationRequest.ts b/src/eventHandlers/invocationRequest.ts
index 3f1d7b22..e64fe925 100644
--- a/src/eventHandlers/invocationRequest.ts
+++ b/src/eventHandlers/invocationRequest.ts
@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License.
 
+import { AzureFunction } from '@azure/functions';
 import { HookData, PostInvocationContext, PreInvocationContext } from '@azure/functions-core';
 import { format } from 'util';
 import { AzureFunctionsRpcMessages as rpc } from '../../azure-functions-language-worker-protobuf/src/rpc';
@@ -82,11 +83,17 @@ export async function invocationRequest(channel: WorkerChannel, requestId: strin
         });
 
         const hookData: HookData = {};
-        const userFunction = channel.functionLoader.getFunc(nonNullProp(msg, 'functionId'));
-        const preInvocContext: PreInvocationContext = { hookData, invocationContext: context, inputs };
+        let userFunction = channel.functionLoader.getFunc(nonNullProp(msg, 'functionId'));
+        const preInvocContext: PreInvocationContext = {
+            hookData,
+            invocationContext: context,
+            functionCallback: <AzureFunction>userFunction,
+            inputs,
+        };
 
         await channel.executeHooks('preInvocation', preInvocContext);
         inputs = preInvocContext.inputs;
+        userFunction = preInvocContext.functionCallback;
 
         let rawResult = userFunction(context, ...inputs);
         resultIsPromise = rawResult && typeof rawResult.then === 'function';
diff --git a/test/eventHandlers/invocationRequest.test.ts b/test/eventHandlers/invocationRequest.test.ts
index 8a67f713..fcf87a5e 100644
--- a/test/eventHandlers/invocationRequest.test.ts
+++ b/test/eventHandlers/invocationRequest.test.ts
@@ -576,6 +576,25 @@ describe('invocationRequest', () => {
         });
     }
 
+    it('preInvocationHook respects change to functionCallback', async () => {
+        loader.getFunc.returns(async (invocContext: Context) => {
+            invocContext.log('old function');
+        });
+        loader.getInfo.returns(new FunctionInfo(Binding.queue));
+
+        testDisposables.push(
+            coreApi.registerHook('preInvocation', (context: coreTypes.PreInvocationContext) => {
+                expect(context.functionCallback).to.be.a('function');
+                context.functionCallback = async (invocContext: Context) => {
+                    invocContext.log('new function');
+                };
+            })
+        );
+
+        sendInvokeMessage([InputData.string]);
+        await stream.assertCalledWith(Msg.receivedInvocLog(), Msg.userTestLog('new function'), Msg.invocResponse([]));
+    });
+
     for (const [func, suffix] of TestFunc.logHookData) {
         it('postInvocationHook' + suffix, async () => {
             loader.getFunc.returns(func);
diff --git a/types-core/index.d.ts b/types-core/index.d.ts
index 05be85ac..5cccfbe0 100644
--- a/types-core/index.d.ts
+++ b/types-core/index.d.ts
@@ -1,7 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License.
 
-import { Context } from '@azure/functions';
+import { AzureFunction, Context } from '@azure/functions';
 
 /**
  * This module is shipped as a built-in part of the Azure Functions Node.js worker and is available at runtime
@@ -45,6 +45,11 @@ declare module '@azure/functions-core' {
          * The input values for this specific invocation. Changes to this array _will_ affect the inputs passed to your function
          */
         inputs: any[];
+
+        /**
+         * The function callback for this specific invocation. Changes to this value _will_ affect the function itself
+         */
+        functionCallback: AzureFunction;
     }
 
     /**