fix(agent-computer): check every request against the trial, not only the tab

backend/__tests__/unit/services/agentComputerBroker.test.js

@@ -46,11 +46,11 @@ const fakePage = (events = []) => {
   return page;
 };
 
-const setup = (status = 'absent') => {
+const setup = (status = 'absent', resolve = jest.fn(async () => ['151.101.1.1'])) => {
   const pages = [];
   const cdpCommands = [];
   const events = [];
-  const context = { route: jest.fn(), newPage: jest.fn(async () => {
+  const context = { route: jest.fn(), routeWebSocket: jest.fn(), on: jest.fn(), newPage: jest.fn(async () => {
     const page = fakePage(events);
     pages.push(page);
     return page;
@@ -68,7 +68,7 @@ const setup = (status = 'absent') => {
     ? { sessions: status === 'absent' ? [] : [{ id: 'one', status }] }
     : { id: 'one', status: 'live' }));
   const connect = jest.fn(async (endpoint) => (endpoint.includes(OWNER) ? browser : {
-    contexts: () => [{ route: jest.fn(), newCDPSession: context.newCDPSession, newPage: async () => {
+    contexts: () => [{ route: jest.fn(), routeWebSocket: jest.fn(), on: jest.fn(), newCDPSession: context.newCDPSession, newPage: async () => {
       const page = fakePage(events);
       pages.push(page);
       return page;
@@ -97,7 +97,7 @@ const setup = (status = 'absent') => {
     allowedHosts: new Set(['www.linkedin.com']),
   };
   return {
-    broker: new AgentComputerBroker(config, request, connect, wake, takeoverState),
+    broker: new AgentComputerBroker(config, request, connect, wake, takeoverState, resolve),
     pages, context, request, connect, wake, takeoverState, config, cdpCommands, events,
   };
 };
@@ -256,6 +256,86 @@ describe('agent computer broker (fake Steel and CDP)', () => {
     expect(pages[0].mouse.click).toHaveBeenCalledTimes(1);
   });
 
+  test('every frame and every subresource is held to the trial, not only the tab', async () => {
+    const resolve = jest.fn(async (host) => (host === 'intranet.example' ? ['10.1.2.3'] : ['151.101.1.1']));
+    const { broker, pages, context } = setup('live', resolve);
+    await broker.snapshot(OWNER, FIRST);
+    const intercept = context.route.mock.calls[0][1];
+    const verdict = async (url, { navigation = false, subframe = false } = {}) => {
+      const route = { request: () => ({ method: () => 'GET', url: () => url, isNavigationRequest: () => navigation,
+        frame: () => ({ page: () => pages[0], parentFrame: () => (subframe ? {} : null) }) }),
+      abort: jest.fn(), fallback: jest.fn() };
+      await intercept(route);
+      if (route.fallback.mock.calls.length) return 'allowed';
+      return route.abort.mock.calls.length ? 'refused' : 'held';
+    };
+    expect(await verdict('https://www.linkedin.com/feed/', { navigation: true })).toBe('allowed');
+    expect(await verdict('https://www.linkedin.com/embed', { navigation: true, subframe: true })).toBe('allowed');
+    expect(await verdict('https://elsewhere.example/', { navigation: true, subframe: true })).toBe('refused');
+    expect(await verdict('https://elsewhere.example/', { navigation: true })).toBe('refused');
+    expect(await verdict('https://10.0.0.1/pixel.png')).toBe('refused');
+    expect(await verdict('http://169.254.169.254/computeMetadata/v1/')).toBe('refused');
+    expect(await verdict('http://[::ffff:10.0.0.1]/pixel.png')).toBe('refused');
+    expect(await verdict('https://intranet.example/logo.png')).toBe('refused');
+    expect(await verdict('https://cdn.example/app.js')).toBe('allowed');
+    expect(resolve).toHaveBeenCalledWith('cdn.example');
+  });
+
+  test('a WebSocket is held to the subresource rule before any tab opens', async () => {
+    const resolve = jest.fn(async (host) => (host === 'intranet.example' ? ['10.1.2.3'] : ['151.101.1.1']));
+    const { broker, context } = setup('live', resolve);
+    await broker.snapshot(OWNER, FIRST);
+    expect(context.routeWebSocket).toHaveBeenCalledTimes(1);
+    // Playwright routes only sockets created after the route is installed.
+    expect(context.routeWebSocket.mock.invocationCallOrder[0])
+      .toBeLessThan(context.newPage.mock.invocationCallOrder[0]);
+    const [matches, guard] = context.routeWebSocket.mock.calls[0];
+    expect(matches(new URL('wss://anything.example/socket'))).toBe(true);
+    const verdict = async (url) => {
+      const ws = { url: () => url, connectToServer: jest.fn(), close: jest.fn(async () => {}) };
+      await guard(ws);
+      if (ws.connectToServer.mock.calls.length) return ws.close.mock.calls.length ? 'both' : 'connected';
+      return ws.close.mock.calls.length ? 'closed' : 'held';
+    };
+    expect(await verdict('wss://www.linkedin.com/realtime')).toBe('connected');
+    expect(await verdict('wss://cdn.example/socket')).toBe('connected');
+    expect(resolve).toHaveBeenCalledWith('cdn.example');
+    expect(await verdict('ws://10.0.0.1:8080/')).toBe('closed');
+    expect(await verdict('ws://169.254.169.254/')).toBe('closed');
+    expect(await verdict('ws://[::ffff:10.0.0.1]/')).toBe('closed');
+    expect(await verdict('wss://intranet.example/socket')).toBe('closed');
+    expect(await verdict('not a url')).toBe('closed');
+  });
+
+  test('a redirect hop that leaves the trial closes the tab', async () => {
+    const resolve = jest.fn(async (host) => (host === 'intranet.example' ? ['10.1.2.3'] : ['151.101.1.1']));
+    const { broker, pages, context } = setup('live', resolve);
+    await broker.snapshot(OWNER, FIRST);
+    const onRequest = context.on.mock.calls.find(([event]) => event === 'request')[1];
+    const hop = (url, navigation, page = pages[0]) => ({ redirectedFrom: () => ({}), url: () => url,
+      isNavigationRequest: () => navigation, frame: () => ({ page: () => page }) });
+    await onRequest(hop('https://www.linkedin.com/login', true));
+    await onRequest(hop('https://cdn.example/app.js', false));
+    expect(pages[0].close).not.toHaveBeenCalled();
+    await onRequest(hop('https://elsewhere.example/', true));
+    expect(pages[0].close).toHaveBeenCalledTimes(1);
+    // The closed tab is gone from the broker; the next action gets a fresh one.
+    await broker.snapshot(OWNER, FIRST);
+    expect(context.newPage).toHaveBeenCalledTimes(2);
+    await onRequest(hop('https://intranet.example/logo.png', false, pages[1]));
+    expect(pages[1].close).toHaveBeenCalledTimes(1);
+  });
+
+  test('a navigation that lands off the trial is refused and its tab closed', async () => {
+    const { broker, pages } = setup('live');
+    await broker.snapshot(OWNER, FIRST);
+    pages[0].goto.mockImplementationOnce(async () => {});
+    pages[0].url = () => 'https://elsewhere.example/landing';
+    await expect(broker.navigate(OWNER, FIRST, 'https://www.linkedin.com/feed/'))
+      .rejects.toThrow('outside this computer trial');
+    expect(pages[0].close).toHaveBeenCalledTimes(1);
+  });
+
   test('the first request from an unowned popup is denied at the context boundary', async () => {
     const { broker, context } = setup('live');
     await broker.snapshot(OWNER, FIRST);

backend/services/agentComputerBroker.ts

@@ -3,9 +3,15 @@
  * broker owns its single live session and gives each native installation a
  * separate tab in that profile. Agents never receive a Steel URL or CDP port.
  */
-import { chromium, type Browser, type BrowserContext, type Page, type Route } from 'playwright-core';
+import {
+  chromium, type Browser, type BrowserContext, type Page, type Request as BrowserRequest, type Route,
+  type WebSocketRoute,
+} from 'playwright-core';
 import { randomUUID } from 'crypto';
+import dns from 'node:dns';
+import net from 'node:net';
 import { AgentComputerTakeover } from '../models/AgentComputerTakeover';
+import { classifyAddress } from './outboundUrlGuard';
 
 type SteelSession = { id: string; status: string };
 type SteelSessions = { sessions: SteelSession[] };
@@ -15,6 +21,7 @@ type ComputerConfig = {
   allowedHosts: ReadonlySet<string>;
 };
 type Connect = (endpoint: string) => Promise<Browser>;
+type Resolve = (host: string) => Promise<string[]>;
 type TakeoverStore = {
   get: (userId: string, installationId: string) => Promise<boolean>;
   set: (userId: string, installationId: string, active: boolean) => Promise<void>;
@@ -91,6 +98,10 @@ const wakeSteel = async (userId: string): Promise<void> => {
   if (!response.ok) throw new Error(`Computer waker failed (${response.status}).`);
 };
 
+const resolveAddresses: Resolve = async (host) => (
+  await dns.promises.lookup(host, { all: true, verbatim: true })
+).map(({ address }) => address);
+
 const takeoverStore: TakeoverStore = {
   get: async (userId, installationId) => Boolean(await AgentComputerTakeover.exists({
     _id: installationId, ownerId: userId, active: true,
@@ -148,8 +159,70 @@ export class AgentComputerBroker {
     private readonly connect: Connect = (endpoint) => chromium.connectOverCDP(endpoint, { timeout: SESSION_TIMEOUT_MS }),
     private readonly wake: (userId: string) => Promise<void> = wakeSteel,
     private readonly takeoverState: TakeoverStore = takeoverStore,
+    private readonly resolve: Resolve = resolveAddresses,
   ) {} // eslint-disable-line no-empty-function
 
+  // Pages closed because a redirect left the trial, so navigate can name why
+  // its goto failed.
+  private leftTrial = new WeakSet<Page>();
+
+  // Any frame's document is a page the agent reads, so every frame's
+  // navigation is held to the trial's hosts, not only the tab's.
+  private onList(url: URL): boolean {
+    return url.protocol === 'https:' && !url.username && !url.password && !url.port
+      && this.config.allowedHosts.has(url.hostname.toLowerCase());
+  }
+
+  // The browser runs inside the cluster, so a subresource may reach only a
+  // public address. The browser does its own DNS: an answer that changes
+  // between this lookup and the browser's is the egress policy's to stop.
+  private async publicDestination(url: URL): Promise<boolean> {
+    if (url.protocol !== 'https:' && url.protocol !== 'http:') return false;
+    const host = url.hostname.replace(/^\[|\]$/g, '').toLowerCase();
+    if (this.config.allowedHosts.has(host)) return true;
+    const addresses = net.isIP(host) ? [host] : await this.resolve(host).catch(() => [] as string[]);
+    return addresses.length > 0 && addresses.every((address) => classifyAddress(address) === 'public');
+  }
+
+  // context.route never sees a WebSocket handshake, so a socket is held to the
+  // subresource rule: an allowlisted host or a public address. Playwright routes
+  // sockets inside the page, which a worker or service worker never runs, so a
+  // socket opened from one is the egress policy's to stop.
+  private async guardWebSocket(ws: WebSocketRoute): Promise<void> {
+    let target: URL | null = null;
+    try { target = new URL(ws.url()); } catch { /* unparseable: closed below */ }
+    if (target && (target.protocol === 'ws:' || target.protocol === 'wss:')) {
+      target.protocol = target.protocol === 'wss:' ? 'https:' : 'http:';
+      if (await this.publicDestination(target)) { ws.connectToServer(); return; }
+    }
+    await ws.close().catch(() => {});
+  }
+
+  private async closeTab(state: UserBrowser, installationId: string, page: Page): Promise<void> {
+    this.leftTrial.add(page);
+    if (state.tabs.get(installationId) === page) state.tabs.delete(installationId);
+    state.frames.delete(installationId);
+    await this.abortPendingWrites(state, installationId);
+    await page.close().catch(() => {});
+  }
+
+  // Playwright continues a redirect hop itself and never calls the route
+  // handler for it, so a hop cannot be held. The request event still names
+  // the hop: a tab whose hop leaves the trial is closed before the agent reads
+  // what came back.
+  private async checkRedirect(request: BrowserRequest, state: UserBrowser): Promise<void> {
+    if (!request.redirectedFrom()) return;
+    let page: Page;
+    try { page = request.frame().page(); } catch { return; }
+    const installationId = [...state.tabs].find(([, tab]) => tab === page)?.[0];
+    if (!installationId) return;
+    let target: URL | null = null;
+    try { target = new URL(request.url()); } catch { /* unparseable: closed below */ }
+    const allowed = target !== null && (request.isNavigationRequest()
+      ? this.onList(target) : await this.publicDestination(target));
+    if (!allowed) await this.closeTab(state, installationId, page);
+  }
+
   private async attachPage(page: Page): Promise<void> {
     page.setDefaultTimeout(ACTION_TIMEOUT_MS);
     // Request interception is installed on the shared context before any tab
@@ -175,12 +248,10 @@ export class AgentComputerBroker {
     const installationId = [...state.tabs].find(([, tab]) => tab === page)?.[0];
     // A popup is not an installation tab. Its first request must be denied.
     if (!installationId) { await route.abort(); return; }
-    if (request.isNavigationRequest() && !request.frame().parentFrame()) {
-      let target: URL;
-      try { target = new URL(request.url()); } catch { await route.abort(); return; }
-      if (target.protocol !== 'https:' || !this.config.allowedHosts.has(target.hostname.toLowerCase())) {
-        await route.abort(); return;
-      }
+    let target: URL;
+    try { target = new URL(request.url()); } catch { await route.abort(); return; }
+    if (request.isNavigationRequest() ? !this.onList(target) : !(await this.publicDestination(target))) {
+      await route.abort(); return;
     }
     if (['GET', 'HEAD', 'OPTIONS'].includes(request.method())) { await route.fallback(); return; }
     // A takeover request alone grants no write permission. A request after a
@@ -233,6 +304,8 @@ export class AgentComputerBroker {
       epoch: randomUUID(),
     };
     await context.route('**/*', (route) => this.guardRequest(route, state));
+    await context.routeWebSocket(() => true, (ws) => this.guardWebSocket(ws));
+    context.on('request', (request) => this.checkRedirect(request, state).catch(() => {}));
     const bindings = await this.takeoverState.bindings(userId);
     const byTarget = new Map(bindings.filter(({ installationId }) => INSTALLATION_ID.test(installationId))
       .map(({ installationId, targetId }) => [targetId, installationId]));
@@ -306,7 +379,18 @@ export class AgentComputerBroker {
       throw new Error('That address is outside this computer trial.');
     }
     const page = await this.tab(userId, installationId);
-    await page.goto(url.toString(), { waitUntil: 'domcontentloaded', timeout: ACTION_TIMEOUT_MS });
+    try {
+      await page.goto(url.toString(), { waitUntil: 'domcontentloaded', timeout: ACTION_TIMEOUT_MS });
+    } catch (error) {
+      if (this.leftTrial.has(page)) throw new Error('That address is outside this computer trial.');
+      throw error;
+    }
+    let landed: URL | null = null;
+    try { landed = new URL(page.url()); } catch { /* closed below */ }
+    if (!landed || !this.onList(landed)) {
+      await this.closeTab(await this.user(userId), installationId, page);
+      throw new Error('That address is outside this computer trial.');
+    }
     return { url: page.url(), title: await page.title() };
   }