Fix preview proxy REST origin rewriting

packages/runtime-playground/src/preview-server.ts

@@ -1,5 +1,7 @@
 import { createServer as createHttpServer, request as httpRequest, type IncomingHttpHeaders, type IncomingMessage, type ServerResponse } from "node:http"
 import { createServer as createNetServer } from "node:net"
+import { StringDecoder } from "node:string_decoder"
+import { Transform } from "node:stream"
 import type { PreviewLease } from "@automattic/wp-codebox-core"
 
 export interface PlaygroundServerRunResponse {
@@ -81,13 +83,19 @@ export async function withPreviewProxy(server: PlaygroundCliServer, port: number
 async function startPreviewProxy(targetUrl: string, port: number, bind: string): Promise<PlaygroundPreviewProxy> {
   const target = new URL(targetUrl)
   const routes = createPreviewRouteRegistry()
-  const proxy = previewProxyServer(target, routes)
+  let proxyOrigin: string | undefined
+  const proxy = previewProxyServer(target, routes, () => proxyOrigin)
   const servers = [proxy]
 
   await listenPreviewProxy(proxy, port, bind)
 
+  const address = proxy.address()
+  const resolvedPort = address && typeof address === "object" ? address.port : port
+  const reportedHost = bind === "0.0.0.0" ? "127.0.0.1" : bind
+  proxyOrigin = `http://${formatPreviewHost(reportedHost)}:${resolvedPort}`
+
   if (bind === "127.0.0.1") {
-    const ipv6Proxy = previewProxyServer(target, routes)
+    const ipv6Proxy = previewProxyServer(target, routes, () => `http://[::1]:${resolvedPort}`)
     try {
       await listenPreviewProxy(ipv6Proxy, port, "::1")
       servers.push(ipv6Proxy)
@@ -99,12 +107,8 @@ async function startPreviewProxy(targetUrl: string, port: number, bind: string):
     }
   }
 
-  const address = proxy.address()
-  const resolvedPort = address && typeof address === "object" ? address.port : port
-  const reportedHost = bind === "0.0.0.0" ? "127.0.0.1" : bind
-
   return {
-    serverUrl: `http://${formatPreviewHost(reportedHost)}:${resolvedPort}`,
+    serverUrl: proxyOrigin,
     previewRoutes: routes,
     diagnostics: {
       schema: "wp-codebox/preview-proxy-diagnostics/v1",
@@ -120,7 +124,7 @@ async function startPreviewProxy(targetUrl: string, port: number, bind: string):
   }
 }
 
-function previewProxyServer(target: URL, routes: InternalPreviewRouteRegistry): PreviewProxyServer {
+function previewProxyServer(target: URL, routes: InternalPreviewRouteRegistry, proxyOrigin: () => string | undefined): PreviewProxyServer {
   const upstreamQueue = createPreviewProxyQueue()
 
   return createHttpServer(async (incoming, outgoing) => {
@@ -133,7 +137,7 @@ function previewProxyServer(target: URL, routes: InternalPreviewRouteRegistry):
       return
     }
 
-    upstreamQueue(() => proxyPreviewRequest(target, incoming, outgoing)).catch((error: Error) => writeProxyError(outgoing, error))
+    upstreamQueue(() => proxyPreviewRequest(target, proxyOrigin(), incoming, outgoing)).catch((error: Error) => writeProxyError(outgoing, error))
   })
 }
 
@@ -159,7 +163,7 @@ function createPreviewRouteRegistry(): InternalPreviewRouteRegistry {
   }
 }
 
-function proxyPreviewRequest(target: URL, incoming: IncomingMessage, outgoing: ServerResponse): Promise<void> {
+function proxyPreviewRequest(target: URL, proxyOrigin: string | undefined, incoming: IncomingMessage, outgoing: ServerResponse): Promise<void> {
   return new Promise((resolve) => {
     let settled = false
     let targetResponse: IncomingMessage | undefined
@@ -187,7 +191,27 @@ function proxyPreviewRequest(target: URL, incoming: IncomingMessage, outgoing: S
       },
       (response) => {
         targetResponse = response
-        outgoing.writeHead(response.statusCode ?? 502, response.statusMessage, proxyResponseHeaders(response.headers))
+        if (proxyOrigin && shouldRewriteProxyResponse(response.headers)) {
+          const headers = proxyResponseHeaders(response.headers, target.origin, proxyOrigin)
+          delete headers["content-length"]
+          delete headers.etag
+          outgoing.writeHead(response.statusCode ?? 502, response.statusMessage, headers)
+          const rewrite = previewOriginRewriteStream(target.origin, proxyOrigin)
+          response.on("error", (error) => {
+            outgoing.destroy(error)
+            settle()
+          })
+          rewrite.on("error", (error) => {
+            outgoing.destroy(error)
+            settle()
+          })
+          outgoing.on("finish", settle)
+          outgoing.on("close", settle)
+          response.pipe(rewrite).pipe(outgoing)
+          return
+        }
+
+        outgoing.writeHead(response.statusCode ?? 502, response.statusMessage, proxyResponseHeaders(response.headers, target.origin, proxyOrigin))
         response.on("error", (error) => {
           outgoing.destroy(error)
           settle()
@@ -211,6 +235,43 @@ function proxyPreviewRequest(target: URL, incoming: IncomingMessage, outgoing: S
   })
 }
 
+function shouldRewriteProxyResponse(headers: IncomingHttpHeaders): boolean {
+  const contentType = String(headers["content-type"] ?? "").toLowerCase()
+  const contentEncoding = String(headers["content-encoding"] ?? "identity").toLowerCase()
+  return (contentEncoding === "" || contentEncoding === "identity")
+    && (contentType.startsWith("text/") || /(?:^|\/)(?:[^;]+\+)?(?:json|javascript|xml)(?:;|$)/.test(contentType))
+}
+
+function rewritePreviewOrigin(body: string, targetOrigin: string, proxyOrigin: string): string {
+  return body
+    .replaceAll(targetOrigin, proxyOrigin)
+    .replaceAll(targetOrigin.replaceAll("/", "\\/"), proxyOrigin.replaceAll("/", "\\/"))
+}
+
+function previewOriginRewriteStream(targetOrigin: string, proxyOrigin: string): Transform {
+  const decoder = new StringDecoder("utf8")
+  const patterns = [targetOrigin, targetOrigin.replaceAll("/", "\\/")]
+  const retainedCharacters = Math.max(...patterns.map((pattern) => pattern.length)) - 1
+  let pending = ""
+
+  return new Transform({
+    transform(chunk: Buffer, _encoding, callback) {
+      const text = pending + decoder.write(chunk)
+      let splitAt = Math.max(0, text.length - retainedCharacters)
+      for (const pattern of patterns) {
+        for (let index = Math.max(0, splitAt - pattern.length + 1); index < splitAt; index += 1) {
+          if (pattern.startsWith(text.slice(index, splitAt))) splitAt = index
+        }
+      }
+      pending = text.slice(splitAt)
+      callback(null, rewritePreviewOrigin(text.slice(0, splitAt), targetOrigin, proxyOrigin))
+    },
+    flush(callback) {
+      callback(null, rewritePreviewOrigin(pending + decoder.end(), targetOrigin, proxyOrigin))
+    },
+  })
+}
+
 function createPreviewProxyQueue(): (task: () => Promise<void>) => Promise<void> {
   let active = false
   const pending: Array<() => void> = []
@@ -277,11 +338,18 @@ function proxyRequestHeaders(headers: IncomingHttpHeaders): IncomingHttpHeaders
   }
 }
 
-function proxyResponseHeaders(headers: IncomingHttpHeaders): IncomingHttpHeaders {
+function proxyResponseHeaders(headers: IncomingHttpHeaders, targetOrigin?: string, proxyOrigin?: string): IncomingHttpHeaders {
   const forwarded = { ...headers }
   delete forwarded.connection
   delete forwarded["transfer-encoding"]
 
+  if (targetOrigin && proxyOrigin) {
+    for (const [name, value] of Object.entries(forwarded)) {
+      if (typeof value === "string") forwarded[name] = rewritePreviewOrigin(value, targetOrigin, proxyOrigin)
+      else if (Array.isArray(value)) forwarded[name] = value.map((item) => rewritePreviewOrigin(item, targetOrigin, proxyOrigin))
+    }
+  }
+
   return forwarded
 }
 

tests/browser-callback-materialization-contracts.test.ts

@@ -1,5 +1,5 @@
 import assert from "node:assert/strict"
-import { createServer } from "node:http"
+import { createServer, type IncomingHttpHeaders } from "node:http"
 import {
   artifactBundleFileManifest,
   browserRunResultEnvelope,
@@ -138,14 +138,39 @@ assert.deepEqual(browserPreviewAuthCookieUrls("http://localhost:9400", ["PUBLIC.
 
 let activeUpstreamRequests = 0
 let maxActiveUpstreamRequests = 0
-const targetServer = createServer(async (_request, response) => {
+let targetServerUrl = ""
+const restRequests: Array<{ method?: string; headers: IncomingHttpHeaders; body: string }> = []
+const targetServer = createServer(async (request, response) => {
   activeUpstreamRequests += 1
   maxActiveUpstreamRequests = Math.max(maxActiveUpstreamRequests, activeUpstreamRequests)
   await new Promise((resolveDelay) => setTimeout(resolveDelay, 25))
   activeUpstreamRequests -= 1
+
+  if (request.url === "/editor") {
+    response.setHeader("content-type", "text/html; charset=utf-8")
+    const escapedOrigin = targetServerUrl.replaceAll("/", "\\/")
+    const splitAt = Math.floor(escapedOrigin.length / 2)
+    response.write(`<script>window.wpApiSettings={"root":"${escapedOrigin.slice(0, splitAt)}`)
+    response.end(`${escapedOrigin.slice(splitAt)}\\/wp-json\\/"}</script>`)
+    return
+  }
+  if (request.url === "/wp-json/wp/v2/pages/5") {
+    let body = ""
+    for await (const chunk of request) body += chunk.toString()
+    restRequests.push({ method: request.method, headers: request.headers, body })
+    if (request.method === "OPTIONS") {
+      response.writeHead(204, { "access-control-allow-headers": "content-type, x-wp-nonce" })
+      response.end()
+      return
+    }
+    response.setHeader("content-type", "application/json")
+    response.end(JSON.stringify({ updated: true }))
+    return
+  }
+
   response.end("ok")
 })
-const targetServerUrl = await listenLocalHttpServer(targetServer)
+targetServerUrl = await listenLocalHttpServer(targetServer)
 let disposed = false
 const proxied = await withPreviewProxy({
   playground: { async run() { return { text: "" } } },
@@ -165,6 +190,31 @@ try {
   })
   await Promise.all([fetch(proxied.serverUrl), fetch(proxied.serverUrl)])
   assert.equal(maxActiveUpstreamRequests, 1)
+
+  const editorHtml = await fetch(`${proxied.serverUrl}/editor`).then((response) => response.text())
+  const restUrl = editorHtml.match(/root":"([^\"]+)/)?.[1]?.replaceAll("\\/", "/")
+  assert.equal(restUrl, `${proxied.serverUrl}/wp-json/`, "WordPress browser URLs use the preview proxy origin")
+
+  const mutationUrl = new URL("wp/v2/pages/5", restUrl).toString()
+  const preflight = await fetch(mutationUrl, {
+    method: "OPTIONS",
+    headers: {
+      origin: proxied.serverUrl,
+      "access-control-request-method": "POST",
+      "access-control-request-headers": "content-type,x-http-method-override,x-wp-nonce",
+    },
+  })
+  assert.equal(preflight.status, 204)
+  assert.equal(restRequests[0]?.headers["access-control-request-headers"], "content-type,x-http-method-override,x-wp-nonce")
+
+  const mutation = await fetch(mutationUrl, {
+    method: "POST",
+    headers: { "content-type": "application/json", "x-http-method-override": "PUT", "x-wp-nonce": "nonce" },
+    body: JSON.stringify({ content: "updated" }),
+  }).then((response) => response.json())
+  assert.deepEqual(mutation, { updated: true })
+  assert.equal(restRequests[1]?.headers["x-http-method-override"], "PUT")
+  assert.equal(restRequests[1]?.body, JSON.stringify({ content: "updated" }))
 } finally {
   await proxied[Symbol.asyncDispose]()
   await closeHttpServer(targetServer)

tests/editor-actions.test.ts

@@ -281,6 +281,7 @@ assert.equal(nestedValidity.warnings[0]?.clientId, "nested-invalid")
 const frontPageTarget = editorOpenTargetFromArgs(["target=front-page"])
 assert.equal(frontPageTarget.kind, "front-page")
 assert.equal(frontPageTarget.url, "")
+const editorRuntimeSpec = { wp: "latest", environment: {} } as never
 
 // resolveEditorOpenTarget asks the running WordPress for page_on_front and rewrites
 // the target to open that exact post in the editor.
@@ -291,7 +292,7 @@ const resolved = await resolveEditorOpenTarget(frontPageTarget, {
     resolveCalls.push(command)
     return { ok: true, text: "57\n" } as never
   },
-  runtimeSpec: { wp: "latest" } as never,
+  runtimeSpec: editorRuntimeSpec,
   server: { serverUrl: "http://localhost" } as never,
 })
 assert.equal(resolved.kind, "post")
@@ -314,7 +315,7 @@ const resolvedPostSlug = await resolveEditorOpenTarget(postSlugTarget, {
     postSlugCalls.push(command)
     return { ok: true, text: "91\n" } as never
   },
-  runtimeSpec: { wp: "latest" } as never,
+  runtimeSpec: editorRuntimeSpec,
   server: { serverUrl: "http://localhost" } as never,
 })
 assert.equal(resolvedPostSlug.kind, "post")
@@ -326,7 +327,7 @@ await assert.rejects(
   resolveEditorOpenTarget(postSlugTarget, {
     command: "wordpress.editor-open",
     runPlaygroundCommand: async () => ({ ok: true, text: "0" } as never),
-    runtimeSpec: { wp: "latest" } as never,
+    runtimeSpec: editorRuntimeSpec,
     server: { serverUrl: "http://localhost" } as never,
   }),
   /resolved no editable post/,
@@ -338,7 +339,7 @@ await assert.rejects(
   resolveEditorOpenTarget(frontPageTarget, {
     command: "wordpress.editor-validate-blocks",
     runPlaygroundCommand: async () => ({ ok: true, text: "0" } as never),
-    runtimeSpec: { wp: "latest" } as never,
+    runtimeSpec: editorRuntimeSpec,
     server: { serverUrl: "http://localhost" } as never,
   }),
   /no static front page/,
@@ -351,7 +352,7 @@ const passthrough = await resolveEditorOpenTarget(postTarget, {
   runPlaygroundCommand: async () => {
     throw new Error("resolveEditorOpenTarget should not run PHP for a concrete target")
   },
-  runtimeSpec: { wp: "latest" } as never,
+  runtimeSpec: editorRuntimeSpec,
   server: { serverUrl: "http://localhost" } as never,
 })
 assert.equal(passthrough.url, "/wp-admin/post.php?post=12&action=edit")