Provide WsRouter to plugins

Author: code-asher

src/node/plugin.ts

@@ -7,6 +7,7 @@ import * as pluginapi from "../../typings/pluginapi"
 import { version } from "./constants"
 import { proxy } from "./proxy"
 import * as util from "./util"
+import { Router as WsRouter, WebsocketRouter } from "./wsRouter"
 const fsp = fs.promises
 
 /**
@@ -21,6 +22,7 @@ require("module")._load = function (request: string, parent: object, isMain: boo
       express,
       field,
       proxy,
+      WsRouter,
     }
   }
   return originalLoad.apply(this, [request, parent, isMain])
@@ -103,14 +105,16 @@ export class PluginAPI {
   }
 
   /**
-   * mount mounts all plugin routers onto r.
+   * mount mounts all plugin routers onto r and websocket routers onto wr.
    */
-  public mount(r: express.Router): void {
+  public mount(r: express.Router, wr: express.Router): void {
     for (const [, p] of this.plugins) {
-      if (!p.router) {
-        continue
+      if (p.router) {
+        r.use(`${p.routerPath}`, p.router())
+      }
+      if (p.wsRouter) {
+        wr.use(`${p.routerPath}`, (p.wsRouter() as WebsocketRouter).router)
       }
-      r.use(`${p.routerPath}`, p.router())
     }
   }
 

src/node/routes/index.ts

@@ -6,15 +6,15 @@ import { promises as fs } from "fs"
 import http from "http"
 import * as path from "path"
 import * as tls from "tls"
+import * as pluginapi from "../../../typings/pluginapi"
 import { HttpCode, HttpError } from "../../common/http"
 import { plural } from "../../common/util"
 import { AuthType, DefaultedArgs } from "../cli"
 import { rootPath } from "../constants"
 import { Heart } from "../heart"
-import { replaceTemplates, redirect } from "../http"
+import { redirect, replaceTemplates } from "../http"
 import { PluginAPI } from "../plugin"
 import { getMediaMime, paths } from "../util"
-import { WebsocketRequest } from "../wsRouter"
 import * as apps from "./apps"
 import * as domainProxy from "./domainProxy"
 import * as health from "./health"
@@ -129,7 +129,7 @@ export const register = async (
 
   const papi = new PluginAPI(logger, process.env.CS_PLUGIN, process.env.CS_PLUGIN_PATH)
   await papi.loadPlugins()
-  papi.mount(app)
+  papi.mount(app, wsApp)
   app.use("/api/applications", apps.router(papi))
 
   app.use(() => {
@@ -170,7 +170,7 @@ export const register = async (
 
   const wsErrorHandler: express.ErrorRequestHandler = async (err, req, res, next) => {
     logger.error(`${err.message} ${err.stack}`)
-    ;(req as WebsocketRequest).ws.end()
+    ;(req as pluginapi.WebsocketRequest).ws.end()
   }
 
   wsApp.use(wsErrorHandler)

src/node/wsRouter.ts

@@ -1,7 +1,7 @@
 import * as express from "express"
 import * as expressCore from "express-serve-static-core"
 import * as http from "http"
-import * as net from "net"
+import * as pluginapi from "../../typings/pluginapi"
 
 export const handleUpgrade = (app: express.Express, server: http.Server): void => {
   server.on("upgrade", (req, socket, head) => {
@@ -20,31 +20,20 @@ export const handleUpgrade = (app: express.Express, server: http.Server): void =
   })
 }
 
-export interface WebsocketRequest extends express.Request {
-  ws: net.Socket
-  head: Buffer
-}
-
-interface InternalWebsocketRequest extends WebsocketRequest {
+interface InternalWebsocketRequest extends pluginapi.WebsocketRequest {
   _ws_handled: boolean
 }
 
-export type WebSocketHandler = (
-  req: WebsocketRequest,
-  res: express.Response,
-  next: express.NextFunction,
-) => void | Promise<void>
-
 export class WebsocketRouter {
   public readonly router = express.Router()
 
-  public ws(route: expressCore.PathParams, ...handlers: WebSocketHandler[]): void {
+  public ws(route: expressCore.PathParams, ...handlers: pluginapi.WebSocketHandler[]): void {
     this.router.get(
       route,
       ...handlers.map((handler) => {
         const wrapped: express.Handler = (req, res, next) => {
           ;(req as InternalWebsocketRequest)._ws_handled = true
-          return handler(req as WebsocketRequest, res, next)
+          return handler(req as pluginapi.WebsocketRequest, res, next)
         }
         return wrapped
       }),

test/httpserver.ts

@@ -1,7 +1,10 @@
+import * as express from "express"
 import * as http from "http"
 import * as nodeFetch from "node-fetch"
+import Websocket from "ws"
 import * as util from "../src/common/util"
 import { ensureAddress } from "../src/node/app"
+import { handleUpgrade } from "../src/node/wsRouter"
 
 // Perhaps an abstraction similar to this should be used in app.ts as well.
 export class HttpServer {
@@ -39,6 +42,13 @@ export class HttpServer {
     })
   }
 
+  /**
+   * Send upgrade requests to an Express app.
+   */
+  public listenUpgrade(app: express.Express): void {
+    handleUpgrade(app, this.hs)
+  }
+
   /**
    * close cleans up the server.
    */
@@ -62,6 +72,13 @@ export class HttpServer {
     return nodeFetch.default(`${ensureAddress(this.hs)}${requestPath}`, opts)
   }
 
+  /**
+   * Open a websocket against the requset path.
+   */
+  public ws(requestPath: string): Websocket {
+    return new Websocket(`${ensureAddress(this.hs).replace("http:", "ws:")}${requestPath}`)
+  }
+
   public port(): number {
     const addr = this.hs.address()
     if (addr && typeof addr === "object") {

test/plugin.test.ts

@@ -23,11 +23,13 @@ describe("plugin", () => {
     await papi.loadPlugins(false)
 
     const app = express.default()
-    papi.mount(app)
+    const wsApp = express.default()
+    papi.mount(app, wsApp)
     app.use("/api/applications", apps.router(papi))
 
     s = new httpserver.HttpServer()
     await s.listen(app)
+    s.listenUpgrade(wsApp)
   })
 
   after(async () => {
@@ -72,4 +74,13 @@ describe("plugin", () => {
     const body = await resp.text()
     assert.equal(body, indexHTML)
   })
+
+  it("/test-plugin/test-app (websocket)", async () => {
+    const ws = s.ws("/test-plugin/test-app")
+    const message = await new Promise((resolve) => {
+      ws.once("message", (message) => resolve(message))
+    })
+    ws.terminate()
+    assert.strictEqual(message, "hello")
+  })
 })

test/test-plugin/src/index.ts

@@ -1,5 +1,8 @@
 import * as cs from "code-server"
 import * as fspath from "path"
+import Websocket from "ws"
+
+const wss = new Websocket.Server({ noServer: true })
 
 export const plugin: cs.Plugin = {
   displayName: "Test Plugin",
@@ -22,6 +25,16 @@ export const plugin: cs.Plugin = {
     return r
   },
 
+  wsRouter() {
+    const wr = cs.WsRouter()
+    wr.ws("/test-app", (req) => {
+      wss.handleUpgrade(req, req.socket, req.head, (ws) => {
+        ws.send("hello")
+      })
+    })
+    return wr
+  },
+
   applications() {
     return [
       {

typings/pluginapi.d.ts

@@ -3,6 +3,9 @@
  */
 import { field, Logger } from "@coder/logger"
 import * as express from "express"
+import * as expressCore from "express-serve-static-core"
+import ProxyServer from "http-proxy"
+import * as net from "net"
 
 /**
  * Overlay
@@ -78,6 +81,27 @@ import * as express from "express"
  * ]
  */
 
+export interface WebsocketRequest extends express.Request {
+  ws: net.Socket
+  head: Buffer
+}
+
+export type WebSocketHandler = (
+  req: WebsocketRequest,
+  res: express.Response,
+  next: express.NextFunction,
+) => void | Promise<void>
+
+export interface WebsocketRouter {
+  readonly router: express.Router
+  ws(route: expressCore.PathParams, ...handlers: WebSocketHandler[]): void
+}
+
+/**
+ * Create a router for websocket routes.
+ */
+export function WsRouter(): WebsocketRouter
+
 /**
  * The Express import used by code-server.
  *
@@ -152,6 +176,15 @@ export interface Plugin {
    */
   router?(): express.Router
 
+  /**
+   * Returns the plugin's websocket router.
+   *
+   * Mounted at <code-sever-root>/<plugin-path>
+   *
+   * If not present, the plugin provides no websockets.
+   */
+  wsRouter?(): WebsocketRouter
+
   /**
    * code-server uses this to collect the list of applications that
    * the plugin can currently provide.