From c21d6b791c537dc341b89257726ea651e9234f2b Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Thu, 21 Oct 2021 10:23:49 -0800
Subject: [PATCH] chore: fix the route(times) to disable interception (#9671)

---
 .../src/client/browserContext.ts              | 26 ++++++++++---------
 .../playwright-core/src/client/network.ts     | 18 ++++++-------
 packages/playwright-core/src/client/page.ts   | 23 +++++++++-------
 3 files changed, 35 insertions(+), 32 deletions(-)

diff --git a/packages/playwright-core/src/client/browserContext.ts b/packages/playwright-core/src/client/browserContext.ts
index 53018929c6..d9a982327a 100644
--- a/packages/playwright-core/src/client/browserContext.ts
+++ b/packages/playwright-core/src/client/browserContext.ts
@@ -141,20 +141,18 @@ export class BrowserContext extends ChannelOwner<channels.BrowserContextChannel,
   }
 
   _onRoute(route: network.Route, request: network.Request) {
-    let handled = false;
     for (const routeHandler of this._routes) {
       if (routeHandler.matches(request.url())) {
-        routeHandler.handle(route, request);
-        handled = true;
-        break;
+        if (routeHandler.handle(route, request)) {
+          this._routes.splice(this._routes.indexOf(routeHandler), 1);
+          if (!this._routes.length)
+            this._wrapApiCall(channel => this._disableInterception(channel), undefined, true).catch(() => {});
+        }
+        return;
       }
     }
-    if (!handled) {
-      // it can race with BrowserContext.close() which then throws since its closed
-      route._internalContinue();
-    } else {
-      this._routes = this._routes.filter(route => !route.expired());
-    }
+    // it can race with BrowserContext.close() which then throws since its closed
+    route._internalContinue();
   }
 
   async _onBinding(bindingCall: BindingCall) {
@@ -284,11 +282,15 @@ export class BrowserContext extends ChannelOwner<channels.BrowserContextChannel,
   async unroute(url: URLMatch, handler?: network.RouteHandlerCallback): Promise<void> {
     return this._wrapApiCall(async (channel: channels.BrowserContextChannel) => {
       this._routes = this._routes.filter(route => route.url !== url || (handler && route.handler !== handler));
-      if (this._routes.length === 0)
-        await channel.setNetworkInterceptionEnabled({ enabled: false });
+      if (!this._routes.length)
+        await this._disableInterception(channel);
     });
   }
 
+  private async _disableInterception(channel: channels.BrowserContextChannel) {
+    await channel.setNetworkInterceptionEnabled({ enabled: false });
+  }
+
   async waitForEvent(event: string, optionsOrPredicate: WaitForEventOptions = {}): Promise<any> {
     return this._wrapApiCall(async (channel: channels.BrowserContextChannel) => {
       const timeout = this._timeoutSettings.timeout(typeof optionsOrPredicate === 'function'  ? {} : optionsOrPredicate);
diff --git a/packages/playwright-core/src/client/network.ts b/packages/playwright-core/src/client/network.ts
index 0a0690915e..b8463cbd16 100644
--- a/packages/playwright-core/src/client/network.ts
+++ b/packages/playwright-core/src/client/network.ts
@@ -620,29 +620,27 @@ export function validateHeaders(headers: Headers) {
 export class RouteHandler {
   private handledCount = 0;
   private readonly _baseURL: string | undefined;
-  private readonly _times: number | undefined;
+  private readonly _times: number;
   readonly url: URLMatch;
   readonly handler: RouteHandlerCallback;
 
-  constructor(baseURL: string | undefined, url: URLMatch, handler: RouteHandlerCallback, times?: number) {
+  constructor(baseURL: string | undefined, url: URLMatch, handler: RouteHandlerCallback, times: number = Number.MAX_SAFE_INTEGER) {
     this._baseURL = baseURL;
     this._times = times;
     this.url = url;
     this.handler = handler;
   }
 
-  public expired(): boolean {
-    return !!this._times && this.handledCount >= this._times;
-  }
-
   public matches(requestURL: string): boolean {
     return urlMatches(this._baseURL, requestURL, this.url);
   }
 
-  public handle(route: Route, request: Request) {
-    this.handler(route, request);
-    if (this._times)
-      this.handledCount++;
+  public handle(route: Route, request: Request): boolean {
+    try {
+      this.handler(route, request);
+    } finally {
+      return ++this.handledCount >= this._times;
+    }
   }
 }
 
diff --git a/packages/playwright-core/src/client/page.ts b/packages/playwright-core/src/client/page.ts
index 7c2094fda0..f4d78e07b6 100644
--- a/packages/playwright-core/src/client/page.ts
+++ b/packages/playwright-core/src/client/page.ts
@@ -168,18 +168,17 @@ export class Page extends ChannelOwner<channels.PageChannel, channels.PageInitia
   }
 
   private _onRoute(route: Route, request: Request) {
-    let handled = false;
     for (const routeHandler of this._routes) {
       if (routeHandler.matches(request.url())) {
-        routeHandler.handle(route, request);
-        handled = true;
-        break;
+        if (routeHandler.handle(route, request)) {
+          this._routes.splice(this._routes.indexOf(routeHandler), 1);
+          if (!this._routes.length)
+            this._wrapApiCall(channel => this._disableInterception(channel), undefined, true).catch(() => {});
+        }
+        return;
       }
     }
-    if (!handled)
-      this._browserContext._onRoute(route, request);
-    else
-      this._routes = this._routes.filter(route => !route.expired());
+    this._browserContext._onRoute(route, request);
   }
 
   async _onBinding(bindingCall: BindingCall) {
@@ -467,11 +466,15 @@ export class Page extends ChannelOwner<channels.PageChannel, channels.PageInitia
   async unroute(url: URLMatch, handler?: RouteHandlerCallback): Promise<void> {
     return this._wrapApiCall(async (channel: channels.PageChannel) => {
       this._routes = this._routes.filter(route => route.url !== url || (handler && route.handler !== handler));
-      if (this._routes.length === 0)
-        await channel.setNetworkInterceptionEnabled({ enabled: false });
+      if (!this._routes.length)
+        await this._disableInterception(channel);
     });
   }
 
+  private async _disableInterception(channel: channels.PageChannel) {
+    await channel.setNetworkInterceptionEnabled({ enabled: false });
+  }
+
   async screenshot(options: channels.PageScreenshotOptions & { path?: string } = {}): Promise<Buffer> {
     return this._wrapApiCall(async (channel: channels.PageChannel) => {
       const copy = { ...options };