1 files changed, 38 insertions, 10 deletions

diff --git a/lib/sockets.js b/lib/sockets.js
index 242963cf..69fbab07 100644
--- a/lib/sockets.js
+++ b/lib/sockets.js
@@ -6,6 +6,7 @@ wss = new WebSocket.Server({ clientTracking: false, noServer: true });
 
 /* eslint-disable-next-line */
 const pick = (O, ...K) => K.reduce((o, k) => ((o[k] = O[k]), o), {});
+const all = {};
 const subscriptions = [];
 
 emit = (username, type, data) => {
@@ -14,7 +15,11 @@ emit = (username, type, data) => {
   if (type === "payment" && data.amount > 0) {
     for (let i = subscriptions.length - 1; i >= 0; i--) {
       let s = subscriptions[i];
-      if (s.invoice.amount === data.amount || s.invoice.address === data.address || s.invoice.unconfidential === data.address) {
+      if (
+        s.invoice.amount === data.amount ||
+        s.invoice.address === data.address ||
+        s.invoice.unconfidential === data.address
+      ) {
         s.ws.send(JSON.stringify({ type, data }));
         subscriptions.splice(i, 1);
       }
@@ -35,6 +40,10 @@ broadcast = (type, data) => {
       ws.send(JSON.stringify({ type, data }));
     }
   }
+
+  Object.values(all).map((ws) => {
+    ws.send(JSON.stringify({ type, data }));
+  });
 };
 
 server.on("upgrade", function (req, socket, head) {
@@ -45,9 +54,21 @@ server.on("upgrade", function (req, socket, head) {
 
 const names = {};
 
+setInterval(() => {
+  for (const username in sockets) {
+    for (const id in sockets[username]) {
+      let ws = sockets[username][id];
+      ws.beats++;
+      if (ws.beats > 10) ws.close();
+    }
+  }
+}, 2000);
+
 wss.on("connection", function (ws, req) {
   const id = uuidv4();
-  ws.send(JSON.stringify({ type: "connected" }));
+  ws.beats = 0;
+  ws.send(JSON.stringify({ type: "connected", data: id }));
+  all[id] = ws;
 
   ws.on("message", async (message) => {
     let type, data;
@@ -58,20 +79,24 @@ wss.on("connection", function (ws, req) {
     }
 
     switch (type) {
+      case "heartbeat":
+        ws.beats = 0;
+        break;
+
       case "login":
         let token = data;
         try {
-          if (!token || token === "null") throw new Error("Unauthorized socket login");
+          if (!token || token === "null")
+            throw new Error("Unauthorized socket login");
           const username = jwt.decode(token).username;
           if (!username) throw new Error("Invalid JWT token");
 
           if (!sockets[username]) sockets[username] = {};
           const existing = Object.keys(sockets[username]);
-          if (existing.length > 5) {
-            l.warn("too many sockets, logging out", existing[3]);
-            sockets[username][existing[3]].send(
-              JSON.stringify({ type: "login", data: false })
-            );
+          if (existing.length > 4) {
+            const p = existing.find((sid) => sid !== id);
+            l.warn("sending logout to", p);
+            sockets[username][p].send(JSON.stringify({ type: "logout" }));
           }
           sockets[username][id] = ws;
           names[id] = username;
@@ -101,6 +126,7 @@ wss.on("connection", function (ws, req) {
             JSON.stringify({ type: "version", data: config.clientVersion })
           );
         } catch (e) {
+          l.error("caught socket error", e.message);
           ws.close();
         }
         break;
@@ -129,7 +155,6 @@ wss.on("connection", function (ws, req) {
         }
         break;
       case "subscribe":
-        l.info("subscribe", data.amount);
         subscriptions.push({ invoice: data, ws });
         break;
       default:
@@ -140,8 +165,11 @@ wss.on("connection", function (ws, req) {
   ws.on("close", function () {
     try {
       const username = names[id];
-      if (sockets[username] && sockets[username][id])
+      if (sockets[username] && sockets[username][id]) {
         delete sockets[username][id];
+      }
+
+      if (all[id]) delete all[id];
     } catch (e) {
       l.error("problem closing socket", e.message);
     }