From ef57d5d71e3be2ee3a34a713d7546ccef2947b11 Mon Sep 17 00:00:00 2001
From: Alex Renoki <alex@renoki.org>
Date: Tue, 27 Jun 2023 13:02:51 +0300
Subject: [PATCH] Implemented global router Added extra args for handlers in WS

---
 src/router.ts    | 15 +++++++++++++++
 src/ws/router.ts | 37 ++++++++++++++-----------------------
 2 files changed, 29 insertions(+), 23 deletions(-)
 create mode 100644 src/router.ts

diff --git a/src/router.ts b/src/router.ts
new file mode 100644
index 0000000..920e253
--- /dev/null
+++ b/src/router.ts
@@ -0,0 +1,15 @@
+export class Router {
+    static protocols = new Map<string, { handler: Function; }>();
+
+    static registerHandler(protocol: string, handler: Function): void {
+        this.protocols.set(protocol, { handler });
+    }
+
+    static getProtocol(protocol: string): { handler: Function; } {
+        return this.protocols.get(protocol);
+    }
+
+    static async handle(protocol: string, ...args): Promise<void> {
+        await this.getProtocol(protocol).handler(...args);
+    }
+}
diff --git a/src/ws/router.ts b/src/ws/router.ts
index ea80527..9244d0e 100644
--- a/src/ws/router.ts
+++ b/src/ws/router.ts
@@ -1,50 +1,41 @@
 import type * as FN from '@soketi/impl/types';
+import { Router as BaseRouter } from '../router';
 
-export class Router {
-    static protocols = new Map<string, { handler: Function; }>();
-
+export class Router extends BaseRouter {
     static ON_NEW_CONNECTION = 'onNewConnection';
     static ON_CONNECTION_CLOSED = 'onConnectionClosed';
     static ON_MESSAGE = 'onMessage';
     static ON_ERROR = 'onError';
 
-    static registerHandler(protocol: string, handler: Function): void {
-        this.protocols.set(protocol, { handler });
-    }
-
-    static onNewConnection(cb: (connection: FN.WS.Connection) => any): void {
+    static onNewConnection(cb: (connection: FN.WS.Connection, ...args) => any): void {
         this.registerHandler(this.ON_NEW_CONNECTION, cb);
     }
 
-    static onConnectionClosed(cb: (connection: FN.WS.Connection, code?: number, message?: any) => any): void {
+    static onConnectionClosed(cb: (connection: FN.WS.Connection, code?: number, message?: any, ...args) => any): void {
         this.registerHandler(this.ON_CONNECTION_CLOSED, cb);
     }
 
-    static onMessage(cb: (connection: FN.WS.Connection, message: any) => any): void {
+    static onMessage(cb: (connection: FN.WS.Connection, message: any, ...args) => any): void {
         this.registerHandler(this.ON_MESSAGE, cb);
     }
 
-    static onError(cb: (connection: FN.WS.Connection, error: any) => any): void {
+    static onError(cb: (connection: FN.WS.Connection, error: any, ...args) => any): void {
         this.registerHandler(this.ON_ERROR, cb);
     }
 
-    static async handleNewConnection(connection: FN.WS.Connection): Promise<void> {
-        await this.getProtocol(this.ON_NEW_CONNECTION).handler(connection);
-    }
-
-    static async handleConnectionClosed(connection: FN.WS.Connection, code?: number, message?: any): Promise<void> {
-        await this.getProtocol(this.ON_CONNECTION_CLOSED).handler(connection, code, message);
+    static async handleNewConnection(connection: FN.WS.Connection, ...args): Promise<void> {
+        await this.handle(this.ON_NEW_CONNECTION, connection, ...args);
     }
 
-    static async handleMessage(connection: FN.WS.Connection, message: any): Promise<void> {
-        await this.getProtocol(this.ON_MESSAGE).handler(connection, message);
+    static async handleConnectionClosed(connection: FN.WS.Connection, code?: number, message?: any, ...args): Promise<void> {
+        await this.handle(this.ON_CONNECTION_CLOSED, connection, code, message, ...args);
     }
 
-    static async handleError(connection: FN.WS.Connection, error: any): Promise<void> {
-        await this.getProtocol(this.ON_ERROR).handler(connection, error);
+    static async handleMessage(connection: FN.WS.Connection, message: any, ...args): Promise<void> {
+        await this.handle(this.ON_MESSAGE, connection, message, ...args);
     }
 
-    static getProtocol(protocol: string): { handler: Function; } {
-        return this.protocols.get(protocol);
+    static async handleError(connection: FN.WS.Connection, error: any, ...args): Promise<void> {
+        await this.handle(this.ON_ERROR, connection, error, ...args);
     }
 }