3 files changed, 172 insertions, 114 deletions
diff --git a/assets/src/network/channel/messages/event_handler.ts b/assets/src/network/channel/messages/event_handler.ts
new file mode 100644
index 0000000..749f289
--- /dev/null
+++ b/assets/src/network/channel/messages/event_handler.ts
@@ -0,0 +1,153 @@
+import type { Channel } from "phoenix";
+import type {
+    AnyMessage,
+    EventName,
+    MessageForEvent,
+    TokenizedMessage,
+} from "./messages";
+
+// The single handler for all events, which is used to dispatch to the correct
+// handler for each event and token.
+// Every event can only have either one single handler, or multiple handlers for different tokens.
+export type EventHandler = {
+    [EN in EventName]?: MessageHandler<MessageForEvent<EN>>;
+};
+
+// A handler for a specific event and message.
+// A message handler can either be a single handler or can have multiple handlers for different tokens
+type MessageHandler<M extends AnyMessage> =
+    // A single handler
+    | { type: "single"; handler: Handler<M> }
+    // A group of handlers for different tokens
+    // Can only be used for messages which have a "token" field.
+    | (M extends TokenizedMessage
+          ? {
+                type: "token";
+                handler: TokenHandler<M>;
+            }
+          : never);
+
+// A map of token to handler for a specific event.
+export type TokenHandler<M extends TokenizedMessage> = Map<string, Handler<M>>;
+
+// A function which handles a message for a specific event.
+export type Handler<M extends AnyMessage> = (message?: M) => void;
+
+// A function that unregisters a single event handler.
+export type Unregister = () => void;
+
+// Adds a single handler for a specific event.
+export function registerHandler<M extends AnyMessage>(
+    eventHandler: EventHandler,
+    channel: Channel,
+    event: M["event_name"],
+    handler: Handler<M>
+) {
+    const messageHandler: MessageHandler<M> = {
+        type: "single",
+        handler,
+    };
+
+    let unregisterChannelEvent: Unregister | null = null;
+    if (typeof eventHandler[event] === "undefined") {
+        // Register a new event handler, since this is the first handler for this event
+        unregisterChannelEvent = registerNewEvent<M>(
+            channel,
+            messageHandler,
+            event
+        );
+    } else {
+        throw new Error("Event already has a handler attached to it.");
+    }
+
+    // TODO: Check if we there is a possibility for TS to accept this.
+    // Technically this should work, because TS is afraid of the generic type parameter
+    // not matching the message type of the handler. But it should match, since
+    // the handler should accept only the exact message type from which the event_name was derived.
+    // @ts-ignore
+    eventHandler[event] = messageHandler;
+
+    // This could cause problems if we would allow to redefine the handler,
+    // as that would cause the Unsubscribe function to no longer apply to this specific handler,
+    // but as we throw an error on redefinition, this should be fine.
+    return () => {
+        delete eventHandler[event];
+
+        // If we registered a new event on the channel, we need to unregister it
+        if (unregisterChannelEvent !== null) {
+            unregisterChannelEvent();
+        }
+    };
+}
+
+export function registerHandlerForSpecificToken<M extends TokenizedMessage>(
+    eventHandler: EventHandler,
+    channel: Channel,
+    event: M["event_name"],
+    token: string,
+    handler: Handler<M>
+): Unregister {
+    if (typeof eventHandler[event] === "undefined") {
+        // TODO: Same as above, this should probably be valid.
+        // @ts-ignore
+        eventHandler[event] = {
+            type: "token",
+            handler: new Map<string, Handler<M>>(),
+        };
+        // @ts-ignore
+        registerNewEvent<M>(channel, eventHandler[event], event);
+    } else {
+        if (eventHandler[event].type === "single") {
+            throw new Error("Event already has a handler attached to it.");
+        }
+    }
+
+    // @ts-ignore This shoudl be valid, as we derive the event name from the message type.
+    const tokenHandler: TokenHandler<M> = eventHandler[event].handler;
+    tokenHandler.set(token, handler);
+
+    return () => {
+        // Unregister the handler for this specific token
+        tokenHandler.delete(token);
+        // If there are no more handlers for this event, we can unregister the event
+        if (tokenHandler.size === 0) {
+            // We should technically use the ref here, but we don't yet have a way to get it, as the event
+            // could have been registered outside of this function call.
+            channel.off(event);
+        }
+    };
+}
+
+// Adds a callback for a new event
+function registerNewEvent<M extends AnyMessage>(
+    channel: Channel,
+    messageHandler: MessageHandler<M>,
+    event: M["event_name"]
+): Unregister {
+    const callback = (data: M) => {
+        onEvent<M>(messageHandler, data);
+    };
+
+    const ref = channel.on(event, callback);
+
+    return () => channel.off(event, ref);
+}
+
+function onEvent<M extends AnyMessage>(
+    messageHandler: MessageHandler<M>,
+    message: M
+) {
+    if (messageHandler.type === "token") {
+        const token: string = message["token"];
+
+        const handler = messageHandler.handler.get(token);
+
+        if (typeof handler !== "undefined") {
+            (handler as Handler<M>)(message);
+        } else {
+            console.warn(`Received message for unknown token: ${token}`);
+        }
+    } else {
+        messageHandler.handler(message);
+    }
+}
diff --git a/assets/src/network/channel/messages/handler.ts b/assets/src/network/channel/messages/handler.ts
deleted file mode 100644
index cc8f005..0000000
--- a/assets/src/network/channel/messages/handler.ts
+++ /dev/null
@@ -1,112 +0,0 @@
-import type { Channel } from "phoenix";
-import type { AnyMessage } from "./messages";
-
-export type Handlers = {
-    [event: string]: EventHandler<unknown>;
-};
-
-type EventHandler<Message extends AnyMessage> = {
-    tokenHandler: TokenHandler<Message>;
-    directHandlers: Handler<Message>[];
-};
-
-export type TokenHandler<Message extends AnyMessage> = {
-    [token: string]: Handler<Message>;
-};
-
-export type Handler<Message extends AnyMessage> = (
-    message?: Message,
-    unregister?: UnregisterHandler
-) => void;
-
-export type UnregisterHandler = () => void;
-
-export function registerTokenHandler<Message extends AnyMessage>(
-    handlers: Handlers,
-    channel: Channel,
-    event: string,
-    token: string | null,
-    handler: Handler<Message>
-): UnregisterHandler {
-    let eventHandler = handlers[event];
-
-    // If this event did not yet have any handlers registered we have to register it
-    if (eventHandler === undefined) {
-        eventHandler = {
-            tokenHandler: {},
-            directHandlers: [],
-        };
-
-        handlers[event] = eventHandler;
-
-        registerNewEvent<Message>(channel, eventHandler, event);
-    }
-
-    let unregister: UnregisterHandler;
-
-    if (token === null) {
-        const directHandlers = eventHandler.directHandlers;
-        directHandlers.push(handler);
-
-        unregister = makeDirectUnregister(directHandlers, handler);
-    } else {
-        const tokenHandler = eventHandler.tokenHandler;
-        tokenHandler[token] = handler;
-
-        unregister = makeTokenUnregister(tokenHandler, token);
-    }
-
-    return unregister;
-}
-
-function registerNewEvent<Message extends AnyMessage>(
-    channel: Channel,
-    eventHandler: EventHandler<Message>,
-    event: string
-) {
-    const callback = (data: Message) => {
-        handleEvent<Message>(eventHandler, data);
-    };
-
-    channel.on(event, callback);
-}
-
-function handleEvent<Message extends AnyMessage>(
-    eventHandler: EventHandler<Message>,
-    message: Message
-) {
-    if (message["token"] !== undefined) {
-        const token = message["token"];
-
-        const tokenHandler = eventHandler.tokenHandler;
-        const handler: Handler<Message> = tokenHandler[token];
-
-        if (handler !== undefined) {
-            handler(message, makeTokenUnregister(tokenHandler, token));
-        }
-    }
-
-    const directHandlers = eventHandler.directHandlers;
-    for (const handler of directHandlers) {
-        handler(message, makeDirectUnregister(directHandlers, handler));
-    }
-}
-
-function makeDirectUnregister<Message extends AnyMessage>(
-    directHandlers: Handler<Message>[],
-    handler: Handler<Message>
-): UnregisterHandler {
-    return () => {
-        const index = directHandlers.findIndex(h => h === handler);
-        directHandlers.splice(index, 1);
-    };
-}
-
-function makeTokenUnregister<Message extends AnyMessage>(
-    tokenHandler: TokenHandler<Message>,
-    token: string
-): UnregisterHandler {
-    return () => {
-        delete tokenHandler[token];
-    };
-}
diff --git a/assets/src/network/channel/messages/messages.ts b/assets/src/network/channel/messages/messages.ts
index 3f2253d..67f70ac 100644
--- a/assets/src/network/channel/messages/messages.ts
+++ b/assets/src/network/channel/messages/messages.ts
@@ -1,5 +1,12 @@
 export type AnyMessage = ShareMessage | RequestMessage;
 
+export type EventName = AnyMessage["event_name"];
+export type MessageForEvent<EN> = Extract<AnyMessage, { event_name: EN }>;
+
+export type TokenizedMessage = {
+    token: string;
+} & AnyMessage;
+
 // Messages for the sharer
 
 export type ShareMessage =
@@ -9,20 +16,24 @@ export type ShareMessage =
     | RequestIceCandidateMessage;
 
 export type NewRequestMessage = {
+    event_name: "new_request";
     token: string;
 };
 
 export type RequestCancelledMessage = {
+    event_name: "request_cancelled";
     token: string;
 };
 
 export type ShareAcceptedMessage = {
+    event_name: "share_accepted";
     token: string;
     sdp: string;
     type: RTCSdpType;
 };
 
 export type RequestIceCandidateMessage = {
+    event_name: "request_ice_candidate";
     token: string;
     candidate: RTCIceCandidateInit;
 };
@@ -35,15 +46,21 @@ export type RequestMessage =
     | ShareCancelledMessage
     | ShareIceCandidateMessage;
 
-export type RequestAcknowledgedMessage = {};
+export type RequestAcknowledgedMessage = {
+    event_name: "request_acknowledged";
+};
 
 export type RequestAcceptedMessage = {
+    event_name: "request_accepted";
     sdp: string;
     type: RTCSdpType;
 };
 
-export type ShareCancelledMessage = {};
+export type ShareCancelledMessage = {
+    event_name: "share_cancelled";
+};
 
 export type ShareIceCandidateMessage = {
+    event_name: "share_ice_candidate";
     candidate: RTCIceCandidateInit;
 };