From 23d9dbc011711ed3b57a8d9d63baaf69b5856780 Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Mon, 4 Apr 2022 19:56:04 -0800
Subject: [PATCH] chore: rewrite non-web protocol for Electron tracing (#13305)

---
 .../trace/recorder/snapshotterInjected.ts     |  4 ---
 packages/trace-viewer/src/snapshotRenderer.ts | 35 +++++++++++++++++--
 packages/trace-viewer/src/snapshotServer.ts   |  4 +--
 packages/trace-viewer/src/snapshotStorage.ts  |  5 ++-
 4 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/packages/playwright-core/src/server/trace/recorder/snapshotterInjected.ts b/packages/playwright-core/src/server/trace/recorder/snapshotterInjected.ts
index 2dbbb28ee0..697e830b91 100644
--- a/packages/playwright-core/src/server/trace/recorder/snapshotterInjected.ts
+++ b/packages/playwright-core/src/server/trace/recorder/snapshotterInjected.ts
@@ -44,7 +44,6 @@ export function frameSnapshotStreamer(snapshotStreamer: string) {
   const kScrollTopAttribute = '__playwright_scroll_top_';
   const kScrollLeftAttribute = '__playwright_scroll_left_';
   const kStyleSheetAttribute = '__playwright_style_sheet_';
-  const kBlobUrlPrefix = 'http://playwright.bloburl/#';
 
   // Symbols for our own info on Nodes/StyleSheets.
   const kSnapshotFrameId = Symbol('__playwright_snapshot_frameid_');
@@ -220,9 +219,6 @@ export function frameSnapshotStreamer(snapshotStreamer: string) {
     private _sanitizeUrl(url: string): string {
       if (url.startsWith('javascript:'))
         return '';
-      // Rewrite blob urls so that Service Worker can intercept them.
-      if (url.startsWith('blob:'))
-        return kBlobUrlPrefix + url;
       return url;
     }
 
diff --git a/packages/trace-viewer/src/snapshotRenderer.ts b/packages/trace-viewer/src/snapshotRenderer.ts
index bb83d888fc..6ff7f6bd9b 100644
--- a/packages/trace-viewer/src/snapshotRenderer.ts
+++ b/packages/trace-viewer/src/snapshotRenderer.ts
@@ -62,8 +62,9 @@ export class SnapshotRenderer {
           // Never set relative URLs as <iframe src> - they start fetching frames immediately.
           const isFrame = n[0] === 'IFRAME' || n[0] === 'FRAME';
           for (const [attr, value] of Object.entries(n[1] || {})) {
-            const attrToSet = isFrame && attr.toLowerCase() === 'src' ? '__playwright_src__' : attr;
-            builder.push(' ', attrToSet, '="', escapeAttribute(value as string), '"');
+            const attrName = isFrame && attr.toLowerCase() === 'src' ? '__playwright_src__' : attr;
+            const attrValue = attr.toLowerCase() === 'href' || attr.toLowerCase() === 'src' ? rewriteURLForCustomProtocol(value) : value;
+            builder.push(' ', attrName, '="', escapeAttribute(attrValue as string), '"');
           }
           builder.push('>');
           for (let i = 2; i < n.length; i++)
@@ -273,3 +274,33 @@ function snapshotScript() {
 
   return `\n(${applyPlaywrightAttributes.toString()})()`;
 }
+
+const schemas = ['about:', 'blob:', 'data:', 'file:', 'ftp:', 'http:', 'https:', 'mailto:', 'sftp:', 'ws:', 'wss:' ];
+
+const kLegacyBlobPrefix = 'http://playwright.bloburl/#';
+
+export function rewriteURLForCustomProtocol(href: string): string {
+  // Legacy support, we used to prepend this to blobs, strip it away.
+  if (href.startsWith(kLegacyBlobPrefix))
+    href = href.substring(kLegacyBlobPrefix.length);
+
+  try {
+    const url = new URL(href);
+    // Sanitize URL.
+    if (url.protocol === 'javascript:')
+      return 'javascript:void(0)';
+
+    // Pass through if possible.
+    const isBlob = url.protocol === 'blob:';
+    if (!isBlob && schemas.includes(url.protocol))
+      return href;
+
+    // Rewrite blob and custom schemas.
+    const prefix = 'pw-' + url.protocol.slice(0, url.protocol.length - 1);
+    url.protocol = 'https:';
+    url.hostname = url.hostname ? `${prefix}.${url.hostname}` : prefix;
+    return url.toString();
+  } catch {
+    return href;
+  }
+}
diff --git a/packages/trace-viewer/src/snapshotServer.ts b/packages/trace-viewer/src/snapshotServer.ts
index 23d9a701e1..cd4c2f83f6 100644
--- a/packages/trace-viewer/src/snapshotServer.ts
+++ b/packages/trace-viewer/src/snapshotServer.ts
@@ -19,8 +19,6 @@ import type { Point } from '@playwright-core/common/types';
 import { URLSearchParams } from 'url';
 import { SnapshotRenderer } from './snapshotRenderer';
 
-const kBlobUrlPrefix = 'http://playwright.bloburl/#';
-
 export class SnapshotServer {
   private _snapshotStorage: SnapshotStorage;
   private _snapshotIds = new Map<string, SnapshotRenderer>();
@@ -65,7 +63,7 @@ export class SnapshotServer {
 
   async serveResource(requestUrl: string, snapshotUrl: string): Promise<Response> {
     const snapshot = this._snapshotIds.get(snapshotUrl)!;
-    const url = requestUrl.startsWith(kBlobUrlPrefix) ? requestUrl.substring(kBlobUrlPrefix.length) : removeHash(requestUrl);
+    const url = removeHash(requestUrl);
     const resource = snapshot?.resourceByUrl(url);
     if (!resource)
       return new Response(null, { status: 404 });
diff --git a/packages/trace-viewer/src/snapshotStorage.ts b/packages/trace-viewer/src/snapshotStorage.ts
index 120ffdfdc4..ef0e7b83ad 100644
--- a/packages/trace-viewer/src/snapshotStorage.ts
+++ b/packages/trace-viewer/src/snapshotStorage.ts
@@ -16,7 +16,7 @@
 
 import type { FrameSnapshot, ResourceSnapshot } from '@playwright-core/server/trace/common/snapshotTypes';
 import { EventEmitter } from './events';
-import { SnapshotRenderer } from './snapshotRenderer';
+import { rewriteURLForCustomProtocol, SnapshotRenderer } from './snapshotRenderer';
 
 export interface SnapshotStorage {
   resources(): ResourceSnapshot[];
@@ -40,10 +40,13 @@ export abstract class BaseSnapshotStorage  implements SnapshotStorage {
   }
 
   addResource(resource: ResourceSnapshot): void {
+    resource.request.url = rewriteURLForCustomProtocol(resource.request.url);
     this._resources.push(resource);
   }
 
   addFrameSnapshot(snapshot: FrameSnapshot): void {
+    for (const override of snapshot.resourceOverrides)
+      override.url = rewriteURLForCustomProtocol(override.url);
     let frameSnapshots = this._frameSnapshots.get(snapshot.frameId);
     if (!frameSnapshots) {
       frameSnapshots = {