Started better typing for message handler

8 files changed, 214 insertions, 143 deletions

diff --git a/assets/src/network/channel/connection.ts b/assets/src/network/channel/connection.ts
index 59e07e8..7e9b13c 100644
--- a/assets/src/network/channel/connection.ts
+++ b/assets/src/network/channel/connection.ts
@@ -2,11 +2,12 @@ import { Channel, Push, Socket } from "phoenix";
 import { get, Readable, writable, Writable } from "svelte/store";
 import {
     Handler,
-    Handlers,
-    registerTokenHandler,
-    UnregisterHandler,
-} from "./messages/handler";
-import type { AnyMessage } from "./messages/messages";
+    EventHandler,
+    registerHandlerForSpecificToken,
+    Unregister,
+    registerHandler,
+} from "./messages/event_handler";
+import type { AnyMessage, TokenizedMessage } from "./messages/messages";
 import { connectSocket, fetchToken } from "./socket";
 
 export enum ConnectionState {
@@ -22,7 +23,7 @@ export type Connection = {
     channel: Channel | null;
     token: string | null;
     state: Writable<ConnectionState>;
-    handlers: Handlers;
+    handlers: EventHandler;
 };
 
 const connection: Connection = {
@@ -50,12 +51,12 @@ export function send(event: string, data: any): Push {
     return connection.channel.push(event, data);
 }
 
-export function onWithToken<Message extends AnyMessage>(
-    event: string,
+export function onWithToken<M extends TokenizedMessage>(
+    event: M["event_name"],
     token: string | null,
-    handler: Handler<Message>
-): UnregisterHandler {
-    return registerTokenHandler(
+    handler: Handler<M>
+): Unregister {
+    return registerHandlerForSpecificToken(
         connection.handlers,
         connection.channel,
         event,
@@ -64,11 +65,16 @@ export function onWithToken<Message extends AnyMessage>(
     );
 }
 
-export function on<Message extends AnyMessage>(
-    event: string,
-    handler: Handler<Message>
-): UnregisterHandler {
-    return onWithToken(event, null, handler);
+export function on<M extends AnyMessage>(
+    event: M["event_name"],
+    handler: Handler<M>
+): Unregister {
+    return registerHandler(
+        connection.handlers,
+        connection.channel,
+        event,
+        handler
+    );
 }
 
 export function getOwnToken(): string {
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;
 };
diff --git a/assets/src/network/channel/share_connection.ts b/assets/src/network/channel/share_connection.ts
index 415cb7f..1bb4f14 100644
--- a/assets/src/network/channel/share_connection.ts
+++ b/assets/src/network/channel/share_connection.ts
@@ -7,7 +7,7 @@ import {
     start,
     updateState,
 } from "./connection";
-import type { UnregisterHandler } from "./messages/handler";
+import type { Unregister } from "./messages/event_handler";
 import type {
     NewRequestMessage,
     RequestCancelledMessage,
@@ -38,12 +38,18 @@ function onNewRequest(message: NewRequestMessage) {
     const request = newIncomingRequest(token);
     requests.addRequest(request);
 
-    onWithToken("request_cancelled", token, onRequestCancelled);
+    const unregister = onWithToken(
+        "request_cancelled",
+        token,
+        (message: RequestCancelledMessage) => {
+            onRequestCancelled(message, unregister);
+        }
+    );
 }
 
 function onRequestCancelled(
     message: RequestCancelledMessage,
-    unregister: UnregisterHandler
+    unregister: Unregister
 ) {
     const token = message.token;
     requests.removeRequest(token);
diff --git a/assets/src/network/transfer/request_transfer.ts b/assets/src/network/transfer/request_transfer.ts
index 317f5e5..e69a9a8 100644
--- a/assets/src/network/transfer/request_transfer.ts
+++ b/assets/src/network/transfer/request_transfer.ts
@@ -1,6 +1,6 @@
 import data from "../../stores/data";
 import { on, send } from "../channel/connection";
-import type { RequestIceCandidateMessage } from "../channel/messages/messages";
+import type { ShareIceCandidateMessage } from "../channel/messages/messages";
 import {
     createTransfer,
     onIncomingIceCandidate,
@@ -27,8 +27,8 @@ export async function createAnswerTransfer(
     };
 
     const unregisterIce = on(
-        "ice_candidate",
-        (message: RequestIceCandidateMessage) =>
+        "share_ice_candidate",
+        (message: ShareIceCandidateMessage) =>
             onIncomingIceCandidate(transfer, message)
     );
 
diff --git a/assets/src/network/transfer/share_transfer.ts b/assets/src/network/transfer/share_transfer.ts
index fc6df9c..16ad68a 100644
--- a/assets/src/network/transfer/share_transfer.ts
+++ b/assets/src/network/transfer/share_transfer.ts
@@ -1,11 +1,12 @@
 import { get } from "svelte/store";
 import dataStore from "../../stores/data";
 import { onWithToken, send } from "../channel/connection";
-import type { UnregisterHandler } from "../channel/messages/handler";
+import type { Unregister } from "../channel/messages/event_handler";
 import type {
     RequestIceCandidateMessage,
     ShareAcceptedMessage,
 } from "../channel/messages/messages";
+import { connectSocket } from "../channel/socket";
 import {
     createTransfer,
     onIncomingIceCandidate,
@@ -34,10 +35,10 @@ export async function createOfferTransfer(
         type: offer.type,
     });
 
-    onWithToken(
+    const unregister: Unregister = onWithToken(
         "share_accepted",
         request_token,
-        (message: ShareAcceptedMessage, unregister) =>
+        (message: ShareAcceptedMessage) =>
             onShareAccepted(transfer, message, unregister)
     );
 
@@ -47,7 +48,7 @@ export async function createOfferTransfer(
 function onShareAccepted(
     transfer: Transfer,
     message: ShareAcceptedMessage,
-    unregister: UnregisterHandler
+    unregister: Unregister
 ) {
     const token = message.token;
 
@@ -55,7 +56,7 @@ function onShareAccepted(
     transfer.pc.setRemoteDescription(answerDescription);
 
     const unregisterIce = onWithToken(
-        "ice_candidate",
+        "request_ice_candidate",
         token,
         (message: RequestIceCandidateMessage) =>
             onIncomingIceCandidate(transfer, message)
diff --git a/assets/src/network/transfer/transfer.ts b/assets/src/network/transfer/transfer.ts
index 7641c17..a0df5bb 100644
--- a/assets/src/network/transfer/transfer.ts
+++ b/assets/src/network/transfer/transfer.ts
@@ -1,7 +1,7 @@
 import { Writable, writable } from "svelte/store";
 import type { IncomingRequest } from "../../models/incoming_request";
 import type { OwnRequest } from "../../models/own_request";
-import type { UnregisterHandler } from "../channel/messages/handler";
+import type { Unregister } from "../channel/messages/event_handler";
 import type {
     RequestIceCandidateMessage,
     ShareIceCandidateMessage,
@@ -83,7 +83,7 @@ export function onIncomingIceCandidate(
 
 export function unregisterIceOnComplete(
     transfer: Transfer,
-    unregister: UnregisterHandler
+    unregister: Unregister
 ) {
     transfer.pc.onicegatheringstatechange = event => {
         const connection = event.target as any;