better security for /extension

playwriter/src/cdp-relay.ts

@@ -1,5 +1,6 @@
 import { Hono } from 'hono'
 import { serve } from '@hono/node-server'
+import { getConnInfo } from '@hono/node-server/conninfo'
 import { createNodeWebSocket } from '@hono/node-ws'
 import type { WSContext } from 'hono/ws'
 import type { Protocol } from './cdp-types.js'
@@ -395,23 +396,23 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
     'elnnakgjclnapgflmidlpobefkdmapdm', // Dev extension (loaded unpacked)
   ]
 
-  function isAllowedOrigin(origin: string | undefined): boolean {
-    if (!origin) {
-      return true // Node.js clients don't send Origin
-    }
-    if (origin.startsWith('chrome-extension://')) {
-      const extensionId = origin.replace('chrome-extension://', '')
-      return ALLOWED_EXTENSION_IDS.includes(extensionId)
-    }
-    return false // Reject browser origins (http://, https://, etc.)
-  }
-
   app.get('/cdp/:clientId?', (c, next) => {
     const origin = c.req.header('origin')
-    if (!isAllowedOrigin(origin)) {
-      logger?.log(chalk.red(`Rejecting /cdp WebSocket from origin: ${origin}`))
-      return c.text('Forbidden', 403)
+    
+    // Validate Origin header if present (Node.js clients don't send it)
+    if (origin) {
+      if (origin.startsWith('chrome-extension://')) {
+        const extensionId = origin.replace('chrome-extension://', '')
+        if (!ALLOWED_EXTENSION_IDS.includes(extensionId)) {
+          logger?.log(chalk.red(`Rejecting /cdp WebSocket from unknown extension: ${extensionId}`))
+          return c.text('Forbidden', 403)
+        }
+      } else {
+        logger?.log(chalk.red(`Rejecting /cdp WebSocket from origin: ${origin}`))
+        return c.text('Forbidden', 403)
+      }
     }
+
     if (token) {
       const url = new URL(c.req.url, 'http://localhost')
       const providedToken = url.searchParams.get('token')
@@ -574,11 +575,33 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
   }))
 
   app.get('/extension', (c, next) => {
+    // 1. Host Validation: The extension endpoint must ONLY be accessed from localhost.
+    // This prevents attackers on the network from hijacking the browser session
+    // even if the server is exposed via 0.0.0.0.
+    const info = getConnInfo(c)
+    const remoteAddress = info.remote.address
+    const isLocalhost = remoteAddress === '127.0.0.1' || remoteAddress === '::1'
+
+    if (!isLocalhost) {
+      logger?.log(chalk.red(`Rejecting /extension WebSocket from remote IP: ${remoteAddress}`))
+      return c.text('Forbidden - Extension must be local', 403)
+    }
+
+    // 2. Origin Validation: Prevent browser-based attacks (CSRF).
+    // Browsers cannot spoof the Origin header, so this ensures the connection
+    // is coming from our specific Chrome Extension, not a malicious website.
     const origin = c.req.header('origin')
-    if (!isAllowedOrigin(origin)) {
-      logger?.log(chalk.red(`Rejecting /extension WebSocket from origin: ${origin}`))
+    if (!origin || !origin.startsWith('chrome-extension://')) {
+      logger?.log(chalk.red(`Rejecting /extension WebSocket: origin must be chrome-extension://, got: ${origin || 'none'}`))
       return c.text('Forbidden', 403)
     }
+    
+    const extensionId = origin.replace('chrome-extension://', '')
+    if (!ALLOWED_EXTENSION_IDS.includes(extensionId)) {
+      logger?.log(chalk.red(`Rejecting /extension WebSocket from unknown extension: ${extensionId}`))
+      return c.text('Forbidden', 403)
+    }
+
     return next()
   }, upgradeWebSocket(() => {
     return {

playwriter/test/security.test.ts

@@ -0,0 +1,111 @@
+
+import { describe, it, expect, afterEach } from 'vitest'
+import { startPlayWriterCDPRelayServer } from '../src/cdp-relay.js'
+import { WebSocket } from 'ws'
+import { killPortProcess } from 'kill-port-process'
+import { createFileLogger } from '../src/create-logger.js'
+
+const TEST_PORT = 19999
+
+async function killProcessOnPort(port: number): Promise<void> {
+    try {
+        await killPortProcess(port)
+    } catch (err) {
+        // Ignore if no process is running
+    }
+}
+
+describe('Security Tests', () => {
+    let server: any = null
+
+    afterEach(async () => {
+        if (server) {
+            server.close()
+            server = null
+        }
+        await killProcessOnPort(TEST_PORT)
+    })
+
+    it('should enforce token authentication for /cdp endpoint', async () => {
+        const token = 'secret-token'
+        const logger = createFileLogger()
+        
+        server = await startPlayWriterCDPRelayServer({
+            port: TEST_PORT,
+            token,
+            logger
+        })
+
+        // Helper to try connecting
+        const tryConnect = (tokenParam?: string) => {
+            return new Promise<void>((resolve, reject) => {
+                const url = `ws://127.0.0.1:${TEST_PORT}/cdp${tokenParam ? `?token=${tokenParam}` : ''}`
+                const ws = new WebSocket(url)
+                
+                ws.on('open', () => {
+                    ws.close()
+                    resolve()
+                })
+                
+                ws.on('error', (err) => {
+                    reject(err)
+                })
+
+                 ws.on('unexpected-response', (req, res) => {
+                    reject(new Error(`Unexpected response: ${res.statusCode}`))
+                    ws.close()
+                 })
+            })
+        }
+
+        // 1. No token -> Should fail
+        await expect(tryConnect()).rejects.toThrow(/Unexpected response: (400|401|403)/)
+
+        // 2. Wrong token -> Should fail
+        await expect(tryConnect('wrong-token')).rejects.toThrow(/Unexpected response: (400|401|403)/)
+
+        // 3. Correct token -> Should succeed
+        await expect(tryConnect(token)).resolves.not.toThrow()
+    })
+
+    it('should enforce localhost restrictions for /extension endpoint', async () => {
+        const logger = createFileLogger()
+        server = await startPlayWriterCDPRelayServer({
+            port: TEST_PORT,
+            logger
+        })
+
+        const tryConnectExtension = (origin?: string) => {
+            return new Promise<void>((resolve, reject) => {
+                const url = `ws://127.0.0.1:${TEST_PORT}/extension`
+                const options = origin ? { headers: { Origin: origin } } : {}
+                const ws = new WebSocket(url, options)
+
+                ws.on('open', () => {
+                    ws.close()
+                    resolve()
+                })
+
+                ws.on('error', (err) => {
+                    reject(err)
+                })
+                
+                ws.on('unexpected-response', (req, res) => {
+                    reject(new Error(`Unexpected response: ${res.statusCode}`))
+                    ws.close()
+                })
+            })
+        }
+
+        // 1. Valid chrome-extension origin -> Should succeed
+        // Use a valid extension ID from ALLOWED_EXTENSION_IDS in cdp-relay.ts
+        await expect(tryConnectExtension('chrome-extension://jfeammnjpkecdekppnclgkkffahnhfhe')).resolves.not.toThrow()
+
+        // 2. Invalid origin (e.g., http://evil.com) -> Should fail
+        await expect(tryConnectExtension('http://evil.com')).rejects.toThrow(/Unexpected response: (400|401|403)/)
+        
+        // 3. No origin -> Should likely fail if strict checking is enabled, but typically extension connection requires specific origin handling. 
+        // Based on implementation, usually it checks if it starts with chrome-extension://
+        await expect(tryConnectExtension()).rejects.toThrow(/Unexpected response: (400|401|403)/)
+    })
+})