From 48d679234ae23faab65a51b7a031f1baadd96149 Mon Sep 17 00:00:00 2001
From: Louis Lam <louislam@users.noreply.github.com>
Date: Tue, 5 Apr 2022 19:41:29 +0800
Subject: [PATCH] Stop bree and cloudflared while the server shutting down

---
 server/jobs.js                                    | 15 +++++++++++----
 server/server.js                                  |  9 ++++++---
 .../socket-handlers/cloudflared-socket-handler.js |  5 +++++
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/server/jobs.js b/server/jobs.js
index 0469d5cab..d33adb98c 100644
--- a/server/jobs.js
+++ b/server/jobs.js
@@ -1,7 +1,7 @@
 const path = require("path");
 const Bree = require("bree");
 const { SHARE_ENV } = require("worker_threads");
-
+let bree;
 const jobs = [
     {
         name: "clear-old-data",
@@ -10,7 +10,7 @@ const jobs = [
 ];
 
 const initBackgroundJobs = function (args) {
-    const bree = new Bree({
+    bree = new Bree({
         root: path.resolve("server", "jobs"),
         jobs,
         worker: {
@@ -26,6 +26,13 @@ const initBackgroundJobs = function (args) {
     return bree;
 };
 
-module.exports = {
-    initBackgroundJobs
+const stopBackgroundJobs = function () {
+    if (bree) {
+        bree.stop();
+    }
+};
+
+module.exports = {
+    initBackgroundJobs,
+    stopBackgroundJobs
 };
diff --git a/server/server.js b/server/server.js
index 48be4e2b7..dfff97e30 100644
--- a/server/server.js
+++ b/server/server.js
@@ -65,7 +65,7 @@ debug("Importing Database");
 const Database = require("./database");
 
 debug("Importing Background Jobs");
-const { initBackgroundJobs } = require("./jobs");
+const { initBackgroundJobs, stopBackgroundJobs } = require("./jobs");
 const { loginRateLimiter, twoFaRateLimiter } = require("./rate-limiter");
 
 const { basicAuth } = require("./auth");
@@ -137,7 +137,7 @@ const { statusPageSocketHandler } = require("./socket-handlers/status-page-socke
 const databaseSocketHandler = require("./socket-handlers/database-socket-handler");
 const TwoFA = require("./2fa");
 const StatusPage = require("./model/status_page");
-const { cloudflaredSocketHandler, autoStart: cloudflaredAutoStart } = require("./socket-handlers/cloudflared-socket-handler");
+const { cloudflaredSocketHandler, autoStart: cloudflaredAutoStart, stop: cloudflaredStop } = require("./socket-handlers/cloudflared-socket-handler");
 
 app.use(express.json());
 
@@ -1459,7 +1459,7 @@ exports.entryPage = "dashboard";
 
     server.once("error", async (err) => {
         console.error("Cannot listen: " + err.message);
-        await Database.close();
+        await shutdownFunction();
     });
 
     server.listen(port, hostname, () => {
@@ -1670,6 +1670,9 @@ async function shutdownFunction(signal) {
     }
     await sleep(2000);
     await Database.close();
+
+    stopBackgroundJobs();
+    await cloudflaredStop();
 }
 
 function finalFunction() {
diff --git a/server/socket-handlers/cloudflared-socket-handler.js b/server/socket-handlers/cloudflared-socket-handler.js
index 3d65cda5b..37c12256d 100644
--- a/server/socket-handlers/cloudflared-socket-handler.js
+++ b/server/socket-handlers/cloudflared-socket-handler.js
@@ -83,3 +83,8 @@ module.exports.autoStart = async (token) => {
         cloudflared.start();
     }
 };
+
+module.exports.stop = async () => {
+    console.log("Stop cloudflared");
+    cloudflared.stop();
+};