chore(ui): test that UI works behind proxy, take 2 (#33771)

tests/config/proxy.ts

@@ -32,6 +32,7 @@ export class TestProxy {
 
   connectHosts: string[] = [];
   requestUrls: string[] = [];
+  wsUrls: string[] = [];
 
   private readonly _server: ProxyServer;
   private readonly _sockets = new Set<net.Socket>();
@@ -58,11 +59,16 @@ export class TestProxy {
     await new Promise(x => this._server.close(x));
   }
 
-  forwardTo(port: number, options?: { allowConnectRequests: boolean }) {
+  forwardTo(port: number, options?: { allowConnectRequests?: boolean, prefix?: string, preserveHostname?: boolean }) {
     this._prependHandler('request', (req: IncomingMessage) => {
       this.requestUrls.push(req.url);
-      const url = new URL(req.url);
-      url.host = `127.0.0.1:${port}`;
+      const url = new URL(req.url, `http://${req.headers.host}`);
+      if (options?.preserveHostname)
+        url.port = '' + port;
+      else
+        url.host = `127.0.0.1:${port}`;
+      if (options?.prefix)
+        url.pathname = url.pathname.replace(options.prefix, '');
       req.url = url.toString();
     });
     this._prependHandler('connect', (req: IncomingMessage) => {
@@ -73,6 +79,17 @@ export class TestProxy {
       this.connectHosts.push(req.url);
       req.url = `127.0.0.1:${port}`;
     });
+    this._prependHandler('upgrade', (req: IncomingMessage) => {
+      this.wsUrls.push(req.url);
+      const url = new URL(req.url, `http://${req.headers.host}`);
+      if (options?.preserveHostname)
+        url.port = '' + port;
+      else
+        url.host = `127.0.0.1:${port}`;
+      if (options?.prefix)
+        url.pathname = url.pathname.replace(options.prefix, '');
+      req.url = url.toString();
+    });
   }
 
   setAuthHandler(handler: (req: IncomingMessage) => boolean) {

tests/playwright-test/ui-mode-trace.spec.ts

@@ -340,6 +340,38 @@ test('should show request source context id', async ({ runUITest, server }) => {
   await expect(page.getByText('api#1')).toBeVisible();
 });
 
+test('should work behind reverse proxy', { annotation: { type: 'issue', description: 'https://github.com/microsoft/playwright/issues/33705' } }, async ({ runUITest, proxyServer: reverseProxy }) => {
+  const { page } = await runUITest({
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test('trace test', async ({ page }) => {
+        await page.setContent('<button>Submit</button>');
+        await page.getByRole('button').click();
+        expect(1).toBe(1);
+      });
+    `,
+  });
+
+  const uiModeUrl = new URL(page.url());
+  reverseProxy.forwardTo(+uiModeUrl.port, { prefix: '/subdir', preserveHostname: true });
+  await page.goto(`${reverseProxy.URL}/subdir${uiModeUrl.pathname}?${uiModeUrl.searchParams}`);
+
+  await page.getByText('trace test').dblclick();
+
+  await expect(page.getByTestId('actions-tree')).toMatchAriaSnapshot(`
+    - tree:
+      - treeitem /Before Hooks \\d+[hmsp]+/
+      - treeitem /page\\.setContent \\d+[hmsp]+/
+      - treeitem /locator\\.clickgetByRole\\('button'\\) \\d+[hmsp]+/
+      - treeitem /expect\\.toBe \\d+[hmsp]+/ [selected]
+      - treeitem /After Hooks \\d+[hmsp]+/
+  `);
+
+  await expect(
+      page.frameLocator('iframe.snapshot-visible[name=snapshot]').locator('button'),
+  ).toHaveText('Submit');
+});
+
 test('should filter actions tab on double-click', async ({ runUITest, server }) => {
   const { page } = await runUITest({
     'a.spec.ts': `

tests/third_party/proxy/index.ts

@@ -3,6 +3,7 @@ import * as net from 'net';
 import * as url from 'url';
 import * as http from 'http';
 import * as os from 'os';
+import { pipeline } from 'stream/promises';
 
 const pkg = { version: '1.0.0' }
 
@@ -33,6 +34,7 @@ export function createProxy(server?: http.Server): ProxyServer {
 	if (!server) server = http.createServer();
 	server.on('request', onrequest);
 	server.on('connect', onconnect);
+	server.on('upgrade', onupgrade);
 	return server;
 }
 
@@ -465,4 +467,29 @@ function requestAuthorization(
 	};
 	res.writeHead(407, headers);
 	res.end('Proxy authorization required');
+}
+
+function onupgrade(req: http.IncomingMessage, socket: net.Socket, head: Buffer) {
+	const proxyReq = http.request(req.url, {
+		method: req.method,
+		headers: req.headers,
+		localAddress: this.localAddress,
+	});
+
+	proxyReq.on('upgrade', async function (proxyRes, proxySocket, proxyHead) {
+		const header = ['HTTP/1.1 101 Switching Protocols'];
+		for (const [key, value] of Object.entries(proxyRes.headersDistinct))
+			header.push(`${key}: ${value}`);
+		socket.write(header.join('\r\n') + '\r\n\r\n');
+		if (proxyHead && proxyHead.length) proxySocket.unshift(proxyHead);
+
+		try {
+			await pipeline(proxySocket, socket, proxySocket);
+		} catch (error) {
+			if (error.code !== "ECONNRESET")
+				throw error;
+		}
+	});
+
+	proxyReq.end(head);
 }