From f1401395342488566820cbf189b8eeeb8c53da7f Mon Sep 17 00:00:00 2001
From: Timothy Carambat <rambat1010@gmail.com>
Date: Fri, 17 May 2024 14:03:25 -0700
Subject: [PATCH] Patch WSS upgrade for manual HTTPS certs (#1429)

* Patch WSS upgrade for manual HTTPS certs

* update comment

* refactor
---
 server/.gitignore          |  5 ++++-
 server/index.js            | 15 +++++++++------
 server/utils/boot/index.js | 11 +++++++----
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/server/.gitignore b/server/.gitignore
index b22a054f..adcf7aa4 100644
--- a/server/.gitignore
+++ b/server/.gitignore
@@ -18,4 +18,7 @@ public/
 # For legacy copies of repo
 documents
 vector-cache
-yarn-error.log
\ No newline at end of file
+yarn-error.log
+
+# Local SSL Certs for HTTPS
+sslcert
\ No newline at end of file
diff --git a/server/index.js b/server/index.js
index 7874045b..59d8fec6 100644
--- a/server/index.js
+++ b/server/index.js
@@ -36,7 +36,12 @@ app.use(
   })
 );
 
-require("express-ws")(app);
+if (!!process.env.ENABLE_HTTPS) {
+  bootSSL(app, process.env.SERVER_PORT || 3001);
+} else {
+  require("express-ws")(app); // load WebSockets in non-SSL mode.
+}
+
 app.use("/api", apiRouter);
 systemEndpoints(apiRouter);
 extensionEndpoints(apiRouter);
@@ -109,8 +114,6 @@ app.all("*", function (_, response) {
   response.sendStatus(404);
 });
 
-if (!!process.env.ENABLE_HTTPS) {
-  bootSSL(app, process.env.SERVER_PORT || 3001);
-} else {
-  bootHTTP(app, process.env.SERVER_PORT || 3001);
-}
+// In non-https mode we need to boot at the end since the server has not yet
+// started and is `.listen`ing.
+if (!process.env.ENABLE_HTTPS) bootHTTP(app, process.env.SERVER_PORT || 3001);
diff --git a/server/utils/boot/index.js b/server/utils/boot/index.js
index ea95e1f5..2022f66e 100644
--- a/server/utils/boot/index.js
+++ b/server/utils/boot/index.js
@@ -12,16 +12,18 @@ function bootSSL(app, port = 3001) {
     const privateKey = fs.readFileSync(process.env.HTTPS_KEY_PATH);
     const certificate = fs.readFileSync(process.env.HTTPS_CERT_PATH);
     const credentials = { key: privateKey, cert: certificate };
+    const server = https.createServer(credentials, app);
 
-    https
-      .createServer(credentials, app)
+    server
       .listen(port, async () => {
         await setupTelemetry();
         new CommunicationKey(true);
         console.log(`Primary server in HTTPS mode listening on port ${port}`);
       })
       .on("error", catchSigTerms);
-    return app;
+
+    require("express-ws")(app, server); // Apply same certificate + server for WSS connections
+    return { app, server };
   } catch (e) {
     console.error(
       `\x1b[31m[SSL BOOT FAILED]\x1b[0m ${e.message} - falling back to HTTP boot.`,
@@ -46,7 +48,8 @@ function bootHTTP(app, port = 3001) {
       console.log(`Primary server in HTTP mode listening on port ${port}`);
     })
     .on("error", catchSigTerms);
-  return app;
+
+  return { app, server: null };
 }
 
 function catchSigTerms() {