2 files changed, 11 insertions, 11 deletions
diff --git a/assets/src/network/channel/connection.ts b/assets/src/network/channel/connection.ts
index 6d6f3de..59eb633 100644
--- a/assets/src/network/channel/connection.ts
+++ b/assets/src/network/channel/connection.ts
@@ -69,7 +69,13 @@ export function onWithToken<Message extends AnyMessage>(
     token: string | null,
     handler: Handler<Message>
 ): UnregisterHandler {
-    return registerTokenHandler(connection.handlers, event, token, handler);
+    return registerTokenHandler(
+        connection.handlers,
+        connection.channel,
+        event,
+        token,
+        handler
+    );
 }
 
 export function on<Message extends AnyMessage>(
@@ -87,14 +93,6 @@ export function getOwnToken(): string {
     return connection.token;
 }
 
-export function getChannel(): Channel {
-    if (connection.state <= ConnectionState.FETCHING_TOKEN) {
-        throw new Error("There is no channel yet.");
-    }
-
-    return connection.channel;
-}
-
 function updateState(state: ConnectionState) {
     // TODO: Notify state listeners
     connection.state = state;
diff --git a/assets/src/network/channel/messages/handler.ts b/assets/src/network/channel/messages/handler.ts
index 5cfe327..f7a00d2 100644
--- a/assets/src/network/channel/messages/handler.ts
+++ b/assets/src/network/channel/messages/handler.ts
@@ -20,6 +20,7 @@ export type UnregisterHandler = () => void;
 
 export function registerTokenHandler<Message extends AnyMessage>(
     handlers: Handlers,
+    channel: Channel,
     event: string,
     token: string | null,
     handler: Handler<Message>
@@ -35,7 +36,7 @@ export function registerTokenHandler<Message extends AnyMessage>(
 
         handlers[event] = eventHandler;
 
-        registerNewEvent<Message>(eventHandler, event);
+        registerNewEvent<Message>(channel, eventHandler, event);
     }
 
     let unregister: UnregisterHandler;
@@ -61,6 +62,7 @@ export function registerTokenHandler<Message extends AnyMessage>(
 }
 
 function registerNewEvent<Message extends AnyMessage>(
+    channel: Channel,
     eventHandler: EventHandler<Message>,
     event: string
 ) {
@@ -68,7 +70,7 @@ function registerNewEvent<Message extends AnyMessage>(
         handleEvent<Message>(eventHandler, data);
     };
 
-    getChannel().on(event, callback);
+    channel.on(event, callback);
 }
 
 function handleEvent<Message extends AnyMessage>(