From 1fa9d30992d1713584b0c3e55701f4800c27888d Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Mon, 29 Jun 2020 16:25:52 -0700
Subject: [PATCH] fix(evaluate): awaitPromise when Promise is overwritten
 (#2759)

Firefox and WebKit require native promises to provide awaitPromise
functionality. When the Promise is overwritten, all evaluations
in the main world produce wrong Promise, so we wrap with async
function to get a native promise instead.
---
 src/common/utilityScriptSerializers.ts |  3 --
 src/injected/utilityScript.ts          | 12 +++++--
 test/evaluation.spec.js                | 44 ++++++++++++++++++++++++++
 test/frame.spec.js                     |  2 +-
 4 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/src/common/utilityScriptSerializers.ts b/src/common/utilityScriptSerializers.ts
index e35a1a11d0..43b5db5273 100644
--- a/src/common/utilityScriptSerializers.ts
+++ b/src/common/utilityScriptSerializers.ts
@@ -52,9 +52,6 @@ export function serializeAsCallArgument(value: any, jsHandleSerializer: (value:
 }
 
 function serialize(value: any, jsHandleSerializer: (value: any) => { fallThrough?: any }, visited: Set<any>): any {
-  if (value && typeof value === 'object' && typeof value.then === 'function')
-    return value;
-
   const result = jsHandleSerializer(value);
   if ('fallThrough' in result)
     value = result.fallThrough;
diff --git a/src/injected/utilityScript.ts b/src/injected/utilityScript.ts
index c1cc401130..eae4d33ac0 100644
--- a/src/injected/utilityScript.ts
+++ b/src/injected/utilityScript.ts
@@ -47,8 +47,16 @@ export default class UtilityScript {
       }
     };
 
-    if (value && typeof value === 'object' && typeof value.then === 'function')
-      return value.then(safeJson);
+    if (value && typeof value === 'object' && typeof value.then === 'function') {
+      return (async () => {
+        // By using async function we ensure that return value is a native Promise,
+        // and not some overriden Promise in the page.
+        // This makes Firefox and WebKit debugging protocols recognize it as a Promise,
+        // properly await and return the value.
+        const promiseValue = await value;
+        return safeJson(promiseValue);
+      })();
+    }
     return safeJson(value);
   }
 }
diff --git a/test/evaluation.spec.js b/test/evaluation.spec.js
index 8924585d87..0acc46d384 100644
--- a/test/evaluation.spec.js
+++ b/test/evaluation.spec.js
@@ -200,6 +200,50 @@ describe('Page.evaluate', function() {
     const result = await page.evaluate(() => -Infinity);
     expect(Object.is(result, -Infinity)).toBe(true);
   });
+  it('should work with overwritten Promise', async({page, server}) => {
+    await page.evaluate(() => {
+      const originalPromise = window.Promise;
+      class Promise2 {
+        static all(...arg) {
+          return wrap(originalPromise.all(...arg));
+        }
+        static race(...arg) {
+          return wrap(originalPromise.race(...arg));
+        }
+        static resolve(...arg) {
+          return wrap(originalPromise.resolve(...arg));
+        }
+        constructor(f, r) {
+          this._promise = new originalPromise(f, r);
+        }
+        then(f, r) {
+          return wrap(this._promise.then(f, r));
+        }
+        catch(f) {
+          return wrap(this._promise.catch(f));
+        }
+        finally(f) {
+          return wrap(this._promise.finally(f));
+        }
+      };
+      const wrap = p => {
+        const result = new Promise2(() => {}, () => {});
+        result._promise = p;
+        return result;
+      };
+      window.Promise = Promise2;
+      window.__Promise2 = Promise2;
+    });
+
+    // Sanity check.
+    expect(await page.evaluate(() => {
+      const p = Promise.all([Promise.race([]), new Promise(() => {}).then(() => {})]);
+      return p instanceof window.__Promise2;
+    })).toBe(true);
+
+    // Now, the new promise should be awaitable.
+    expect(await page.evaluate(() => Promise.resolve(42))).toBe(42);
+  });
   it('should throw when passed more than one parameter', async({page, server}) => {
     const expectThrow = async f => {
       let error;
diff --git a/test/frame.spec.js b/test/frame.spec.js
index 911d72f96a..827c18d946 100644
--- a/test/frame.spec.js
+++ b/test/frame.spec.js
@@ -115,7 +115,7 @@ describe('Frame.evaluate', function() {
       iframe.contentDocument.close();
     });
     // Main world should work.
-    expect(await page.frames()[1].evaluate(() => window.top.location.href)).toBe('http://localhost:8907/empty.html');
+    expect(await page.frames()[1].evaluate(() => window.top.location.href)).toBe(server.EMPTY_PAGE);
     // Utility world should work.
     expect(await page.frames()[1].$('div')).toBeTruthy(null);
   });