plugin.ts: Hijack plugin HTML responses and inject app-overlay.js

nhooyr · nhooyr · commit 3e82a479529d · 2020-11-06T14:38:44.000-05:00
diff --git a/package.json b/package.json
@@ -76,6 +76,7 @@
     "env-paths": "^2.2.0",
     "express": "^5.0.0-alpha.8",
     "fs-extra": "^9.0.1",
+    "hijackresponse": "^5.0.0",
     "http-proxy": "^1.18.0",
     "httpolyglot": "^0.1.2",
     "js-yaml": "^3.13.1",
diff --git a/src/node/plugin.ts b/src/node/plugin.ts
@@ -1,6 +1,7 @@
 import { Logger, field } from "@coder/logger"
 import * as express from "express"
 import * as fs from "fs"
+import * as hijackresponse from "hijackresponse"
 import * as path from "path"
 import * as semver from "semver"
 import * as pluginapi from "../../typings/pluginapi"
@@ -86,7 +87,12 @@ export class PluginAPI {
    */
   public mount(r: express.Router): void {
     for (const [, p] of this.plugins) {
-      r.use(`${p.routerPath}`, p.router())
+      const pr = p.router()
+      r.use(`${p.routerPath}`, (req, res, next) => {
+        // All HTML responses need the overlay javascript injected.
+        tryInjectOverlay(res)
+        pr(req, res, next)
+      })
     }
   }
 
@@ -241,3 +247,58 @@ function q(s: string | undefined): string {
   }
   return JSON.stringify(s)
 }
+
+function tryInjectOverlay(res: express.Response) {
+  hijackresponse.default(res).then((hj) => {
+    if (!res.get("Content-Type").includes("text/html")) {
+      hj.readable.pipe(hj.writable)
+      return
+    }
+    injectOverlay(res, hj)
+  })
+}
+
+// injectOverlay injects the script tag for the overlay into the HTML response
+// in res.
+// A point of improvement is to make it stream instead of buffer the entire response.
+// See for example https://www.npmjs.com/package/stream-buffer-replace
+async function injectOverlay(res: express.Response, hj: hijackresponse.HJ): Promise<void> {
+  res.removeHeader("Content-Length")
+
+  try {
+    const bodyPromise = new Promise<string>((res, rej) => {
+      hj.readable.on("close", rej)
+      hj.writable.on("close", rej)
+      hj.readable.on("error", rej)
+      hj.writable.on("error", rej)
+
+      const chunks: Buffer[] = []
+      hj.readable.on("data", (chunk: Buffer) => {
+        chunks.push(chunk)
+      })
+      hj.readable.on("end", () => {
+        res(String(Buffer.concat(chunks)))
+      })
+    })
+    let body = await bodyPromise
+    body = injectOverlayHTML(body)
+    hj.writable.write(body)
+    hj.writable.end()
+  } catch (err) {
+    hj.destroyAndRestore()
+    res.status(500)
+    res.json({ error: "overlay script injection failed" })
+  }
+}
+
+/**
+ * injectOverlayString injects the app-overlay.js script tag
+ * into the html.
+ */
+export function injectOverlayHTML(html: string): string {
+  return html.replace(
+    "</head>",
+    `  <script defer data-cfasync="false" src="./dist/app-overlay.js"></script>
+  </head>`,
+  )
+}
diff --git a/src/node/routes/index.ts b/src/node/routes/index.ts
@@ -15,8 +15,8 @@ import { replaceTemplates } from "../http"
 import { PluginAPI } from "../plugin"
 import { getMediaMime, paths } from "../util"
 import { WebsocketRequest } from "../wsRouter"
-import * as domainProxy from "./domainProxy"
 import * as apps from "./apps"
+import * as domainProxy from "./domainProxy"
 import * as health from "./health"
 import * as login from "./login"
 import * as proxy from "./pathProxy"
diff --git a/test/plugin.test.ts b/test/plugin.test.ts
@@ -1,23 +1,23 @@
 import { logger } from "@coder/logger"
+import * as express from "express"
+import * as fs from "fs"
 import { describe } from "mocha"
 import * as path from "path"
-import { PluginAPI } from "../src/node/plugin"
 import * as supertest from "supertest"
-import * as express from "express"
+import * as plugin from "../src/node/plugin"
 import * as apps from "../src/node/routes/apps"
-import * as fs from "fs"
 const fsp = fs.promises
 
 /**
  * Use $LOG_LEVEL=debug to see debug logs.
  */
 describe("plugin", () => {
-  let papi: PluginAPI
+  let papi: plugin.PluginAPI
   let app: express.Application
   let agent: supertest.SuperAgentTest
 
   before(async () => {
-    papi = new PluginAPI(logger, path.resolve(__dirname, "test-plugin") + ":meow")
+    papi = new plugin.PluginAPI(logger, path.resolve(__dirname, "test-plugin") + ":meow")
     await papi.loadPlugins()
 
     app = express.default()
@@ -54,9 +54,10 @@ describe("plugin", () => {
   })
 
   it("/test-plugin/test-app", async () => {
-    const indexHTML = await fsp.readFile(path.join(__dirname, "test-plugin/public/index.html"), {
+    let indexHTML = await fsp.readFile(path.join(__dirname, "test-plugin/public/index.html"), {
       encoding: "utf8",
     })
+    indexHTML = plugin.injectOverlayHTML(indexHTML)
     await agent.get("/test-plugin/test-app").expect(200, indexHTML)
   })
 })
diff --git a/test/test-plugin/public/index.html b/test/test-plugin/public/index.html
@@ -1,10 +1,10 @@
 <!DOCTYPE html>
 <html lang="en">
-<head>
-  <meta charset="UTF-8">
-  <title>Test Plugin</title>
-</head>
-<body>
-  <p>Welcome to the test plugin!</p>
-</body>
+  <head>
+    <meta charset="UTF-8" />
+    <title>Test Plugin</title>
+  </head>
+  <body>
+    <p>Welcome to the test plugin!</p>
+  </body>
 </html>
diff --git a/typings/hijackresponse/index.d.ts b/typings/hijackresponse/index.d.ts
@@ -0,0 +1,13 @@
+declare module "hijackresponse" {
+  import * as express from "express"
+  import * as stream from "stream"
+
+  function e(res: express.Response): Promise<HJ>
+  export default e
+
+  export interface HJ {
+    readable: stream.Readable
+    writable: stream.Writable
+    destroyAndRestore: () => void
+  }
+}
diff --git a/yarn.lock b/yarn.lock
@@ -3867,6 +3867,11 @@ hex-color-regex@^1.1.0:
   resolved "https://registry.yarnpkg.com/hex-color-regex/-/hex-color-regex-1.1.0.tgz#4c06fccb4602fe2602b3c93df82d7e7dbf1a8a8e"
   integrity sha512-l9sfDFsuqtOqKDsQdqrMRk0U85RZc0RtOR9yPI7mRVOa4FsR/BVnZ0shmQRM96Ji99kYZP/7hn1cedc1+ApsTQ==
 
+hijackresponse@^5.0.0:
+  version "5.0.0"
+  resolved "https://registry.yarnpkg.com/hijackresponse/-/hijackresponse-5.0.0.tgz#03174e70e582a244034af9254c71e7d0a7a763cd"
+  integrity sha512-9ErmDCl4BdTEgLTQh36ZFP62EXPeKyDd68K/FmQoQFaV3i2QY9l/kacf9Ygw+iuSvF/qsNIkwUz6NOCIZAaqMw==
+
 hmac-drbg@^1.0.0:
   version "1.0.1"
   resolved "https://registry.yarnpkg.com/hmac-drbg/-/hmac-drbg-1.0.1.tgz#d2745701025a6c775a6c545793ed502fc0c649a1"
