From c2b8718bae44557e86eab4034a760b6922165429 Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Thu, 4 Feb 2021 08:45:59 -0800
Subject: [PATCH] fix(waitForFunction): process isFunction auto-detection
 (#5312)

---
 src/server/frames.ts                 | 24 ++++++++++----
 src/server/injected/utilityScript.ts |  3 --
 src/server/javascript.ts             | 48 ++++++++++++++++------------
 3 files changed, 46 insertions(+), 29 deletions(-)

diff --git a/src/server/frames.ts b/src/server/frames.ts
index eb6036b518..aa9a2a4a38 100644
--- a/src/server/frames.ts
+++ b/src/server/frames.ts
@@ -1019,13 +1019,25 @@ export class Frame extends EventEmitter {
   async _waitForFunctionExpression<R>(expression: string, isFunction: boolean | undefined, arg: any, options: types.WaitForFunctionOptions = {}): Promise<js.SmartHandle<R>> {
     if (typeof options.pollingInterval === 'number')
       assert(options.pollingInterval > 0, 'Cannot poll with non-positive interval: ' + options.pollingInterval);
-    const predicateBody = isFunction ? 'return (' + expression + ')(arg)' :  'return (' + expression + ')';
-    const task: dom.SchedulableTask<R> = injectedScript => injectedScript.evaluateHandle((injectedScript, { predicateBody, polling, arg }) => {
-      const innerPredicate = new Function('arg', predicateBody) as (arg: any) => R;
+    expression = js.normalizeEvaluationExpression(expression, isFunction);
+    const task: dom.SchedulableTask<R> = injectedScript => injectedScript.evaluateHandle((injectedScript, { expression, isFunction, polling, arg }) => {
+      const predicate = (arg: any): R => {
+        let result = self.eval(expression);
+        if (isFunction === true) {
+          result = result(arg);
+        } else if (isFunction === false) {
+          result = result;
+        } else {
+          // auto detect.
+          if (typeof result === 'function')
+            result = result(arg);
+        }
+        return result;
+      };
       if (typeof polling !== 'number')
-        return injectedScript.pollRaf((progress, continuePolling) => innerPredicate(arg) || continuePolling);
-      return injectedScript.pollInterval(polling, (progress, continuePolling) => innerPredicate(arg) || continuePolling);
-    }, { predicateBody, polling: options.pollingInterval, arg });
+        return injectedScript.pollRaf((progress, continuePolling) => predicate(arg) || continuePolling);
+      return injectedScript.pollInterval(polling, (progress, continuePolling) => predicate(arg) || continuePolling);
+    }, { expression, isFunction, polling: options.pollingInterval, arg });
     return runAbortableTask(
         progress => this._scheduleRerunnableHandleTask(progress, 'main', task),
         this._page._timeoutSettings.timeout(options));
diff --git a/src/server/injected/utilityScript.ts b/src/server/injected/utilityScript.ts
index 3d8ac8b263..9a2089b820 100644
--- a/src/server/injected/utilityScript.ts
+++ b/src/server/injected/utilityScript.ts
@@ -21,9 +21,6 @@ export default class UtilityScript {
     const args = argsAndHandles.slice(0, argCount);
     const handles = argsAndHandles.slice(argCount);
     const parameters = args.map(a => parseEvaluationResultValue(a, handles));
-    expression = expression.trim();
-    if (/^(async)?\s*function(\s|\()/.test(expression))
-      expression = '(' + expression + ')';
     let result = global.eval(expression);
     if (isFunction === true) {
       result = result(...parameters);
diff --git a/src/server/javascript.ts b/src/server/javascript.ts
index df51a3d8e9..7f92670d71 100644
--- a/src/server/javascript.ts
+++ b/src/server/javascript.ts
@@ -175,26 +175,7 @@ export async function evaluate(context: ExecutionContext, returnByValue: boolean
 
 export async function evaluateExpression(context: ExecutionContext, returnByValue: boolean, expression: string, isFunction: boolean | undefined, ...args: any[]): Promise<any> {
   const utilityScript = await context.utilityScript();
-
-  if (isFunction) {
-    try {
-      new Function('(' + expression + ')');
-    } catch (e1) {
-      // This means we might have a function shorthand. Try another
-      // time prefixing 'function '.
-      if (expression.startsWith('async '))
-        expression = 'async function ' + expression.substring('async '.length);
-      else
-        expression = 'function ' + expression;
-      try {
-        new Function('(' + expression  + ')');
-      } catch (e2) {
-        // We tried hard to serialize, but there's a weird beast here.
-        throw new Error('Passed function is not well-serializable!');
-      }
-    }
-  }
-
+  expression = normalizeEvaluationExpression(expression, isFunction);
   const handles: (Promise<JSHandle>)[] = [];
   const toDispose: Promise<JSHandle>[] = [];
   const pushHandle = (handle: Promise<JSHandle>): number => {
@@ -245,3 +226,30 @@ export function parseUnserializableValue(unserializableValue: string): any {
   if (unserializableValue === '-0')
     return -0;
 }
+
+export function normalizeEvaluationExpression(expression: string, isFunction: boolean | undefined): string {
+  expression = expression.trim();
+
+  if (isFunction) {
+    try {
+      new Function('(' + expression + ')');
+    } catch (e1) {
+      // This means we might have a function shorthand. Try another
+      // time prefixing 'function '.
+      if (expression.startsWith('async '))
+        expression = 'async function ' + expression.substring('async '.length);
+      else
+        expression = 'function ' + expression;
+      try {
+        new Function('(' + expression  + ')');
+      } catch (e2) {
+        // We tried hard to serialize, but there's a weird beast here.
+        throw new Error('Passed function is not well-serializable!');
+      }
+    }
+  }
+
+  if (/^(async)?\s*function(\s|\()/.test(expression))
+    expression = '(' + expression + ')';
+  return expression;
+}