From 4962919de7ce53f6e6cbdb36f1ce9c49e3f0d91e Mon Sep 17 00:00:00 2001
From: marlon-erler <viridian035@gmail.com>
Date: Sun, 7 Jul 2024 18:05:31 +0200
Subject: [PATCH] improve audit

---
 index.ts             | 26 ++++++++++++++++++--------
 websocket-handler.ts | 34 ++++++++++++++++++----------------
 2 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/index.ts b/index.ts
index 95b11df..ec0f05d 100755
--- a/index.ts
+++ b/index.ts
@@ -4,7 +4,7 @@ import {
   addConnection,
   connectServers,
   forgetConnection,
-  getStats,
+  getWebSocketStats,
   processMessage,
 } from "./websocket-handler";
 import { writeInfo, writeStat, writeSuccess } from "./utility";
@@ -41,15 +41,15 @@ async function main() {
         case "/":
           return createFileResponse(defaultPage);
         case "/stats":
-          return new Response(JSON.stringify(getStats()));
+          return new Response(JSON.stringify(getAllStats()));
         default:
           return createFileResponse(requestPath);
       }
     },
 
     websocket: {
-      open() {
-        addConnection();
+      open(ws) {
+        addConnection(ws);
       },
       message(ws, message) {
         processMessage(ws, message.toString(), config);
@@ -68,6 +68,18 @@ async function main() {
   // OTHER SERVERS
   connectServers(config);
 
+  // AUDIT
+  function getServerStats(): [string, string][] {
+    return [
+      ["server url", server.url.toString()],
+      ["server ip address", IP.address()],
+    ];
+  }
+
+  function getAllStats(): [string, string | number][] {
+    return [...getServerStats(), ...getWebSocketStats()];
+  }
+
   // CLI
   writeSuccess(`###\nstarted up ${new Date().toLocaleString()}`);
   console.log("relevant output available in the 'logs' file");
@@ -80,11 +92,9 @@ async function main() {
     console.log(new Date().toLocaleString());
     console.log();
 
-    writeInfo("server url", Colors.bold.green(server.url.toString()));
-    writeInfo("server ip address", Colors.bold.green(IP.address()));
+    getServerStats().forEach((entry) => writeInfo(...entry));
     console.log();
-
-    getStats().forEach((entry) => writeStat(...entry));
+    getWebSocketStats().forEach((entry) => writeStat(...entry));
   }
 
   setInterval(updateCLI, 5000);
diff --git a/websocket-handler.ts b/websocket-handler.ts
index 57ea448..fdff8fa 100644
--- a/websocket-handler.ts
+++ b/websocket-handler.ts
@@ -30,13 +30,17 @@ export interface WebSocketMessage {
 // TRACKING
 const knownMessageIds = new Set<string>();
 
-export let connectedClientCount = 0;
-export let disconnectedServerCount = 0;
 export const subscriptions = new SubscriptionMap();
-export const connectedServers = new Set<WebSocket>();
+export const clientsConnected = new Set<WS>();
+export const serversConnectedAsClient = new Set<WS>();
+export const serversConnected = new Set<WebSocket>();
 
-export function addConnection() {
-  connectedClientCount++;
+export function getServerCount() {
+  return serversConnected.size + serversConnectedAsClient.size;
+}
+
+export function addConnection(ws: WS) {
+  clientsConnected.add(ws);
 }
 
 export function trackSubscription(ws: WS, channel: string): void {
@@ -50,7 +54,7 @@ export function removeSubscription(ws: WS, channel: string): void {
 }
 
 export function forgetConnection(ws: WS): void {
-  connectedClientCount--;
+  clientsConnected.delete(ws);
   const channels = subscriptions.getChannelList(ws);
   channels?.forEach((channel) => {
     subscriptions.delete(ws, channel);
@@ -60,7 +64,7 @@ export function forgetConnection(ws: WS): void {
 // OTHER SERVERS
 export function connectServers(config: typeof defaultConfig): void {
   config.connectedServers.forEach((serverAddress) => {
-    const ws = connectServer(serverAddress, config);
+    connectServer(serverAddress, config);
   });
 }
 
@@ -70,7 +74,7 @@ export function connectServer(
 ): void {
   try {
     const ws = new WebSocket(address);
-    connectedServers.add(ws);
+    serversConnected.add(ws);
 
     ws.addEventListener("open", () => {
       writeSuccess(`connected to server at "${address}".`);
@@ -84,12 +88,10 @@ export function connectServer(
     });
 
     ws.addEventListener("close", () => {
-      connectedServers.delete(ws);
-      disconnectedServerCount++;
+      serversConnected.delete(ws);
       writeError(`server at "${address}" disconnected`);
     });
   } catch (error) {
-    disconnectedServerCount++;
     writeError(`failed to connect to server at "${address}": ${error}`);
   }
 }
@@ -109,13 +111,11 @@ export function subscribeChannel(server: WS, channel: string): void {
 }
 
 // AUDIT
-export function getStats(): [string, number][] {
+export function getWebSocketStats(): [string, number][] {
   return [
     ["channels", subscriptions.clientsPerChannel.size],
-    ["clients connected", connectedClientCount],
-    ["clients subscribing", subscriptions.channelsPerClient.size],
-    ["servers connected", connectedServers.size],
-    ["servers disconnected", disconnectedServerCount],
+    ["clients connected", clientsConnected.size],
+    ["servers connected", getServerCount()],
   ];
 }
 
@@ -147,6 +147,8 @@ export function processMessage(
     sendMessage(messageChannel, messageObject.messageBody, messageObject.uuid!);
   });
   if (messageObject.requestingServerConnection == true) {
+    clientsConnected.delete(ws);
+    serversConnectedAsClient.add(ws);
     subscribeChannels(ws, config);
   }
 }