cleanup, forwarding now handled by subscriptionHandler

subscriptionHandler.ts

@@ -0,0 +1,62 @@
+import { WS } from "./index";
+
+export interface Subscriber {
+  send: (message: string) => void;
+}
+
+export class SubscriptionMap {
+  channelsPerSubscriber = new Map<Subscriber, Set<string>>();
+  subscribersPerChannel = new Map<string, Set<Subscriber>>();
+
+  // prepare
+  private prepareChannelSet(subscriber: Subscriber): void {
+    if (this.channelsPerSubscriber.has(subscriber)) return;
+    this.channelsPerSubscriber.set(subscriber, new Set());
+  }
+
+  private prepareSubscriberSet(channel: string): void {
+    if (this.subscribersPerChannel.has(channel)) return;
+    this.subscribersPerChannel.set(channel, new Set());
+  }
+
+  // utility
+  getChannelList(subscriber: Subscriber): Set<string> | undefined {
+    return this.channelsPerSubscriber.get(subscriber);
+  }
+
+  getClientList(channel: string): Set<Subscriber> | undefined {
+    return this.subscribersPerChannel.get(channel);
+  }
+
+  // modifications
+  set(subscriber: Subscriber, channel: string): void {
+    this.prepareChannelSet(subscriber);
+    this.prepareSubscriberSet(channel);
+
+    this.getChannelList(subscriber)?.add(channel);
+    this.getClientList(channel)?.add(subscriber);
+  }
+
+  delete(subscriber: Subscriber, channel: string): void {
+    const channelList = this.getChannelList(subscriber);
+    const clientList = this.getClientList(channel);
+
+    channelList?.delete(channel);
+    if (channelList?.size == 0) {
+      this.channelsPerSubscriber.delete(subscriber);
+    }
+
+    clientList?.delete(subscriber);
+    if (clientList?.size == 0) {
+      this.subscribersPerChannel.delete(channel);
+    }
+  }
+
+  // messages
+  forwardMessage(channels: string[], messageString: string): void {
+    channels.forEach((channel) => {
+      const subscribers = this.getClientList(channel);
+      subscribers?.forEach((subscriber) => subscriber.send(messageString));
+    });
+  }
+}

subscriptionMap.ts

@@ -1,50 +0,0 @@
-import { WS } from "./index";
-
-export class SubscriptionMap {
-  channelsPerClient = new Map<WS, Set<string>>();
-  clientsPerChannel = new Map<string, Set<WS>>();
-
-  // prepare
-  private prepareChannelSet(ws: WS) {
-    if (this.channelsPerClient.has(ws)) return;
-    this.channelsPerClient.set(ws, new Set());
-  }
-
-  private prepareClientSet(channel: string) {
-    if (this.clientsPerChannel.has(channel)) return;
-    this.clientsPerChannel.set(channel, new Set());
-  }
-
-  // utility
-  getChannelList(ws: WS) {
-    return this.channelsPerClient.get(ws);
-  }
-
-  getClientList(channel: string) {
-    return this.clientsPerChannel.get(channel);
-  }
-
-  // modifications
-  set(ws: WS, channel: string) {
-    this.prepareChannelSet(ws);
-    this.prepareClientSet(channel);
-
-    this.getChannelList(ws)?.add(channel);
-    this.getClientList(channel)?.add(ws);
-  }
-
-  delete(ws: WS, channel: string) {
-    const channelList = this.getChannelList(ws);
-    const clientList = this.getClientList(channel);
-
-    channelList?.delete(channel);
-    if (channelList?.size == 0) {
-      this.channelsPerClient.delete(ws);
-    }
-
-    clientList?.delete(ws);
-    if (clientList?.size == 0) {
-      this.clientsPerChannel.delete(channel);
-    }
-  }
-}

websocket-handler.ts

@@ -1,3 +1,7 @@
+import {
+  Subscriber,
+  SubscriptionMap,
+} from "./subscriptionHandler";
 import {
   doIfIsString,
   parseMessage,
@@ -7,7 +11,6 @@ import {
 } from "./utility";
 
 import Crypto from "crypto";
-import { SubscriptionMap } from "./subscriptionMap";
 import { WS } from "./index";
 import { defaultConfig } from "./config-handler";
 
@@ -35,7 +38,9 @@ export interface WebSocketMessage {
 const knownMessageIds = new Set<string>();
 
 export const subscriptionMap = new SubscriptionMap();
+
 export const clientsConnected = new Set<WS>();
+
 export const serversConnectedAsClient = new Set<WS>();
 export const serversConnected = new Set<WebSocket>();
 export const serversDisconnected = new Set<string>();
@@ -58,14 +63,20 @@ export function forgetConnection(ws: WS): void {
   });
 }
 
-export function trackSubscription(ws: WS, channel: string): void {
-  subscriptionMap.set(ws, channel);
-  confirmSubscription(ws, channel);
+export function trackSubscription(
+  subscriber: Subscriber,
+  channel: string
+): void {
+  subscriptionMap.set(subscriber, channel);
+  confirmSubscription(subscriber, channel);
 }
 
-export function forgetSubscription(ws: WS, channel: string): void {
-  subscriptionMap.delete(ws, channel);
-  confirmSubscription(ws, channel);
+export function forgetSubscription(
+  subscriber: Subscriber,
+  channel: string
+): void {
+  subscriptionMap.delete(subscriber, channel);
+  confirmSubscription(subscriber, channel);
 }
 
 // OTHER SERVERS
@@ -131,7 +142,7 @@ export function subscribeChannel(server: WS, channel: string): void {
 // AUDIT
 export function getWebSocketStats(): [string, number][] {
   return [
-    ["channels", subscriptionMap.clientsPerChannel.size],
+    ["channels", subscriptionMap.subscribersPerChannel.size],
     ["clients connected", clientsConnected.size],
     ["servers connected", getServerCount()],
     ["servers disconnected", serversDisconnected.size],
@@ -190,10 +201,7 @@ export function sendMessage(
   const messageObject = { uuid, messageChannel, messageBody };
   const messageString = stringifyMessage(messageObject);
 
-  channels.forEach((channel) => {
-    const subscribers = subscriptionMap.getClientList(channel);
-    subscribers?.forEach((ws) => ws.send(messageString));
-  });
+  subscriptionMap.forwardMessage(channels, messageString);
 }
 
 function requestServerConnection(ws: WS): void {
@@ -204,11 +212,16 @@ function requestServerConnection(ws: WS): void {
   ws.send(messageString);
 }
 
-function confirmSubscription(ws: WS, messageChannel: string): void {
+function confirmSubscription(
+  subscriber: Subscriber,
+  messageChannel: string
+): void {
   const messageObject = {
     messageChannel,
-    subscribed: subscriptionMap.getChannelList(ws) ? true : false,
+    subscribed: subscriptionMap.getChannelList(subscriber)?.has(messageChannel)
+      ? true
+      : false,
   };
   const messageString = stringifyMessage(messageObject);
-  ws.send(messageString);
+  subscriber.send(messageString);
 }