From d969256eeaca66c1a1a077fda527b609a127cd3b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=8E?= <sliphua@outlook.com>
Date: Sat, 23 Jul 2022 04:26:44 +0800
Subject: [PATCH] feat(route): add request .frame and .serviceWorker (#7)

---
 src/WS/WSClient.ts           |  2 +-
 src/WS/WSServer.ts           | 15 +++++++++++++++
 src/WS/handle/NodeHandle.ts  | 10 ++++++++++
 src/WS/message.ts            |  3 +++
 src/WS/route/RouteRequest.ts | 21 ++++++++++++++++++++-
 5 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/src/WS/WSClient.ts b/src/WS/WSClient.ts
index cb6a5e2..ad0d9a4 100644
--- a/src/WS/WSClient.ts
+++ b/src/WS/WSClient.ts
@@ -130,7 +130,7 @@ export default class WSClient {
       });
       body = await bodyBlob.arrayBuffer();
     }
-    const request = new RouteRequest(meta, body ?? null);
+    const request = new RouteRequest(meta, body ?? null, this.ws);
     const route = new Route(this.ws, meta.id, request);
     const requestURL = new URL(meta.url, window.location.origin);
     const handlers = this.routes.filter((r) => r.matches(requestURL));
diff --git a/src/WS/WSServer.ts b/src/WS/WSServer.ts
index b699112..37f27aa 100644
--- a/src/WS/WSServer.ts
+++ b/src/WS/WSServer.ts
@@ -186,13 +186,22 @@ export default class WSServer {
       }
       const bodyBuffer = request.postDataBuffer();
       const hasBody = bodyBuffer !== null;
+      let frame: number | null;
+      try {
+        frame = this.createHandleID(request.frame());
+      } catch {
+        frame = null;
+      }
+      const sw = request.serviceWorker();
       client.send(createRouteMeta({
         id: this.routeList.length,
         hasBody,
+        frame,
         headersArray,
         isNavigationRequest: request.isNavigationRequest(),
         method: request.method(),
         resourceType: request.resourceType(),
+        serviceWorker: sw ? this.createHandleID(sw) : null,
         url: request.url(),
       }));
       if (hasBody) {
@@ -262,6 +271,12 @@ export default class WSServer {
         result = e instanceof Error ? e : String(e);
       }
     }
+    if (resolveID === null) {
+      if (error) {
+        throw result;
+      }
+      return;
+    }
     source.send(createHandleMeta({
       action: 'resolve',
       id,
diff --git a/src/WS/handle/NodeHandle.ts b/src/WS/handle/NodeHandle.ts
index 1a725f1..d7e4f9a 100644
--- a/src/WS/handle/NodeHandle.ts
+++ b/src/WS/handle/NodeHandle.ts
@@ -53,6 +53,16 @@ const finalizationRegistry = new FinalizationRegistry(({ id, ws }: {
 });
 
 export default class NodeHandle<T = unknown> extends Handle {
+  /**
+   * Share the handle ID with an object so that the matching handles will keep referencing the node
+   * target until all ID users are disposed or garbage-collected.
+   * @internal
+   */
+  static share(id: number, ws: WebSocket, withObj: object) {
+    idReferenceCounts[id] = (idReferenceCounts[id] ?? 0) + 1;
+    finalizationRegistry.register(withObj, { id, ws });
+  }
+
   constructor(
     id: number,
     private readonly ws: WebSocket,
diff --git a/src/WS/message.ts b/src/WS/message.ts
index 8374abc..b075e1c 100644
--- a/src/WS/message.ts
+++ b/src/WS/message.ts
@@ -5,10 +5,13 @@ export interface RouteMetaBase {
 export interface RouteRequestMeta extends RouteMetaBase {
   id: number;
   hasBody: boolean;
+
+  frame: number | null;
   headersArray: { name: string, value: string }[];
   isNavigationRequest: boolean;
   method: string;
   resourceType: string;
+  serviceWorker: number | null;
   url: string;
 }
 
diff --git a/src/WS/route/RouteRequest.ts b/src/WS/route/RouteRequest.ts
index 9104b9e..2107bb5 100644
--- a/src/WS/route/RouteRequest.ts
+++ b/src/WS/route/RouteRequest.ts
@@ -1,3 +1,5 @@
+import type playwright from 'playwright-core';
+import NodeHandle from '../handle/NodeHandle.js';
 import { RouteRequestMeta } from '../message.js';
 import { FallbackOverrides } from './Route.js';
 
@@ -5,7 +7,12 @@ export default class RouteRequest {
   constructor(
     private readonly requestMeta: RouteRequestMeta,
     private readonly requestBody: ArrayBuffer | null,
-  ) {}
+    private readonly ws: WebSocket,
+  ) {
+    const { frame, serviceWorker } = requestMeta;
+    if (frame !== null) NodeHandle.share(frame, ws, this);
+    if (serviceWorker !== null) NodeHandle.share(serviceWorker, ws, this);
+  }
 
   private fallbackOverrides: FallbackOverrides = {};
 
@@ -39,6 +46,13 @@ export default class RouteRequest {
     return this.cachedAllHeaders;
   }
 
+  frame() {
+    if (this.requestMeta.frame === null) {
+      throw new Error('Service Worker requests do not have an associated frame');
+    }
+    return new NodeHandle<playwright.Frame>(this.requestMeta.frame, this.ws);
+  }
+
   headerValue(name: string) {
     return this.allHeaders()[name.toLowerCase()];
   }
@@ -111,6 +125,11 @@ export default class RouteRequest {
     return this.requestMeta.resourceType;
   }
 
+  serviceWorker() {
+    if (this.requestMeta.serviceWorker === null) return null;
+    return new NodeHandle<playwright.Worker>(this.requestMeta.serviceWorker, this.ws);
+  }
+
   url() {
     return this.fallbackOverrides.url ?? this.requestMeta.url;
   }