From 59ba78c0281f18fb4b645fb20b41d56072516edb Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Tue, 16 Feb 2021 14:12:58 -0600
Subject: [PATCH 1/3] Force shutdown sockets during tests

---
 test/httpserver.ts | 36 ++++++++++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/test/httpserver.ts b/test/httpserver.ts
index 4fe54f88..399ccda1 100644
--- a/test/httpserver.ts
+++ b/test/httpserver.ts
@@ -1,5 +1,6 @@
 import * as express from "express"
 import * as http from "http"
+import * as net from "net"
 import * as nodeFetch from "node-fetch"
 import Websocket from "ws"
 import * as util from "../src/common/util"
@@ -8,13 +9,21 @@ import { handleUpgrade } from "../src/node/wsRouter"
 
 // Perhaps an abstraction similar to this should be used in app.ts as well.
 export class HttpServer {
-  private hs = http.createServer()
+  private readonly sockets = new Set<net.Socket>()
+  private cleanupTimeout?: NodeJS.Timeout
 
-  public constructor(hs?: http.Server) {
-    // See usage in test/integration.ts
-    if (hs) {
-      this.hs = hs
-    }
+  // See usage in test/integration.ts
+  public constructor(private readonly hs = http.createServer()) {
+    this.hs.on("connection", (socket) => {
+      this.sockets.add(socket)
+      socket.on("close", () => {
+        this.sockets.delete(socket)
+        if (this.cleanupTimeout && this.sockets.size === 0) {
+          clearTimeout(this.cleanupTimeout)
+          this.cleanupTimeout = undefined
+        }
+      })
+    })
   }
 
   /**
@@ -54,6 +63,8 @@ export class HttpServer {
    */
   public close(): Promise<void> {
     return new Promise((res, rej) => {
+      // Close will not actually close anything; it just waits until everything
+      // is closed.
       this.hs.close((err) => {
         if (err) {
           rej(err)
@@ -61,6 +72,19 @@ export class HttpServer {
         }
         res()
       })
+
+      // If there are sockets remaining we might need to force close them or
+      // this promise might never resolve.
+      if (this.sockets.size > 0) {
+        // Give sockets a chance to close up shop.
+        this.cleanupTimeout = setTimeout(() => {
+          this.cleanupTimeout = undefined
+          for (const socket of this.sockets.values()) {
+            console.warn("a socket was left hanging")
+            socket.destroy()
+          }
+        }, 1000)
+      }
     })
   }
 

From 7f80d152d3c298ca0c4d89a47efa985ce81d571c Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Tue, 16 Feb 2021 14:14:21 -0600
Subject: [PATCH 2/3] Add healthz tests

---
 test/health.test.ts | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)
 create mode 100644 test/health.test.ts

diff --git a/test/health.test.ts b/test/health.test.ts
new file mode 100644
index 00000000..4eae9c60
--- /dev/null
+++ b/test/health.test.ts
@@ -0,0 +1,40 @@
+import * as httpserver from "./httpserver"
+import * as integration from "./integration"
+
+describe("health", () => {
+  let codeServer: httpserver.HttpServer | undefined
+
+  afterEach(async () => {
+    if (codeServer) {
+      await codeServer.close()
+      codeServer = undefined
+    }
+  })
+
+  it("/healthz", async () => {
+    ;[, , codeServer] = await integration.setup(["--auth=none"], "")
+    const resp = await codeServer.fetch("/healthz")
+    expect(resp.status).toBe(200)
+    const json = await resp.json()
+    expect(json).toStrictEqual({ lastHeartbeat: 0, status: "expired" })
+  })
+
+  it("/healthz (websocket)", async () => {
+    ;[, , codeServer] = await integration.setup(["--auth=none"], "")
+    const ws = codeServer.ws("/healthz")
+    const message = await new Promise((resolve, reject) => {
+      ws.on("error", console.error)
+      ws.on("message", (message) => {
+        try {
+          const j = JSON.parse(message.toString())
+          resolve(j)
+        } catch (error) {
+          reject(error)
+        }
+      })
+      ws.on("open", () => ws.send(JSON.stringify({ event: "health" })))
+    })
+    ws.terminate()
+    expect(message).toStrictEqual({ event: "health", status: "expired", lastHeartbeat: 0 })
+  })
+})

From 2d8b785fb896724cd0d2400452063845fc7b9267 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Tue, 16 Feb 2021 14:14:52 -0600
Subject: [PATCH 3/3] Fix health socket not getting client messages

Forgot to resume. Went ahead and did the same for the test plugin
although it only sends messages and doesn't receive any.
---
 src/node/routes/health.ts     | 5 +++--
 test/test-plugin/src/index.ts | 3 ++-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/node/routes/health.ts b/src/node/routes/health.ts
index f38bb0ab..faf1f9c7 100644
--- a/src/node/routes/health.ts
+++ b/src/node/routes/health.ts
@@ -13,8 +13,8 @@ router.get("/", (req, res) => {
 export const wsRouter = WsRouter()
 
 wsRouter.ws("/", async (req) => {
-  wss.handleUpgrade(req, req.socket, req.head, (ws) => {
-    ws.on("message", () => {
+  wss.handleUpgrade(req, req.ws, req.head, (ws) => {
+    ws.addEventListener("message", () => {
       ws.send(
         JSON.stringify({
           event: "health",
@@ -23,5 +23,6 @@ wsRouter.ws("/", async (req) => {
         }),
       )
     })
+    req.ws.resume()
   })
 })
diff --git a/test/test-plugin/src/index.ts b/test/test-plugin/src/index.ts
index 592ad372..772b59d8 100644
--- a/test/test-plugin/src/index.ts
+++ b/test/test-plugin/src/index.ts
@@ -28,7 +28,8 @@ export const plugin: cs.Plugin = {
   wsRouter() {
     const wr = cs.WsRouter()
     wr.ws("/test-app", (req) => {
-      cs.wss.handleUpgrade(req, req.socket, req.head, (ws) => {
+      cs.wss.handleUpgrade(req, req.ws, req.head, (ws) => {
+        req.ws.resume()
         ws.send("hello")
       })
     })