Reregister unregistered token handlers on re-use

3 files changed, 31 insertions, 16 deletions

diff --git a/assets/src/network/channel/messages/event_handler.ts b/assets/src/network/channel/messages/event_handler.ts
index 749f289..fe950ea 100644
--- a/assets/src/network/channel/messages/event_handler.ts
+++ b/assets/src/network/channel/messages/event_handler.ts
@@ -87,7 +87,9 @@ export function registerHandlerForSpecificToken<M extends TokenizedMessage>(
     token: string,
     handler: Handler<M>
 ): Unregister {
-    if (typeof eventHandler[event] === "undefined") {
+    const messageHandler = eventHandler[event];
+
+    if (typeof messageHandler === "undefined") {
         // TODO: Same as above, this should probably be valid.
         // @ts-ignore
         eventHandler[event] = {
@@ -96,10 +98,15 @@ export function registerHandlerForSpecificToken<M extends TokenizedMessage>(
         };
         // @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.");
-        }
+    } else if (
+        messageHandler.type === "token" &&
+        messageHandler.handler.size === 0
+    ) {
+        // If there is already a token handler with no token, we need to register the event again
+        // @ts-ignore
+        registerNewEvent<M>(channel, messageHandler, event);
+    } else if (messageHandler.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.
@@ -125,7 +132,9 @@ function registerNewEvent<M extends AnyMessage>(
     event: M["event_name"]
 ): Unregister {
     const callback = (data: M) => {
-        onEvent<M>(messageHandler, data);
+        // Add event_name to message, so the type definitions match
+        const message = { event_name: event, ...data };
+        onEvent<M>(messageHandler, message);
     };
 
     const ref = channel.on(event, callback);
diff --git a/assets/src/network/transfer/share_transfer.ts b/assets/src/network/transfer/share_transfer.ts
index 16ad68a..0f88e04 100644
--- a/assets/src/network/transfer/share_transfer.ts
+++ b/assets/src/network/transfer/share_transfer.ts
@@ -6,7 +6,6 @@ import type {
     RequestIceCandidateMessage,
     ShareAcceptedMessage,
 } from "../channel/messages/messages";
-import { connectSocket } from "../channel/socket";
 import {
     createTransfer,
     onIncomingIceCandidate,
diff --git a/assets/src/network/transfer/transfer.ts b/assets/src/network/transfer/transfer.ts
index a0df5bb..27182be 100644
--- a/assets/src/network/transfer/transfer.ts
+++ b/assets/src/network/transfer/transfer.ts
@@ -39,25 +39,27 @@ export function createTransfer(
         negotiated: true,
         id: 0,
     });
-    const state = writable(TransferState.CONNECTING);
 
-    channel.onopen = () => {
-        state.set(TransferState.TRANSFERRING);
-        const completeTransfer = () => state.set(TransferState.DONE);
-        onChannel(channel, completeTransfer);
-    };
+    const state = writable(TransferState.CONNECTING);
 
-    return {
+    const transfer = {
         pc,
         channel,
         state,
     };
+
+    channel.onopen = () => {
+        state.set(TransferState.TRANSFERRING);
+        onChannel(channel, () => onTransferComplete(transfer));
+    };
+
+    return transfer;
 }
 
 export function bindTransfer(
     request: OwnRequest | IncomingRequest,
     transferPromise: Promise<Transfer>,
-    onTransferComplete: () => void
+    completeTransfer: () => void
 ) {
     transferPromise.then(transfer => {
         request.transfer = transfer;
@@ -68,7 +70,7 @@ export function bindTransfer(
                 // Once the data has been transferred we can remove the transfer
                 request.transfer = null;
 
-                onTransferComplete();
+                completeTransfer();
             }
         });
     });
@@ -92,3 +94,8 @@ export function unregisterIceOnComplete(
         }
     };
 }
+
+function onTransferComplete(transfer: Transfer) {
+    transfer.state.set(TransferState.DONE);
+    transfer.pc.close();
+}