Improved some stuff

src/index.ts

@@ -1,16 +1,19 @@
 import { Duplex } from "stream";
+import { WebSocketServer } from "ws";
 
 export interface NextWebsocketConfig {
   app: NextWebsocketApp;
 }
 
 export interface WsRouteHandler {
-  upgrade?: (req: Request, socket: Duplex, head: Buffer<ArrayBuffer>) => void;
+  upgrade?: (req: Request, socket: Duplex, head: NonSharedBuffer) => void;
   post?: (req: Request) => void;
+  get?: (req: Request) => void;
 }
 
 export class NextWebsocketApp {
   handlers: Record<string, WsRouteHandler> = {};
+  websocketHandlers: ((ws: WebSocketServer) => void)[] = [];
 
   on<K extends keyof WsRouteHandler>(
     method: K,
@@ -21,17 +24,22 @@ export class NextWebsocketApp {
     this.handlers[route][method] = handler;
   }
 
-  upgrade(
-    route: string,
-    handler: (req: Request, socket: Duplex, head: Buffer<ArrayBuffer>) => void,
-  ) {
+  wss(handler: (ws: WebSocketServer) => void) {
+    this.websocketHandlers.push(handler);
+  }
+
+  upgrade(route: string, handler: WsRouteHandler["upgrade"]) {
     this.on("upgrade", route, handler);
   }
 
-  post(route: string, handler: () => void) {
+  post(route: string, handler: WsRouteHandler["post"]) {
     this.on("post", route, handler);
   }
 
+  get(route: string, handler: WsRouteHandler["get"]) {
+    this.on("get", route, handler);
+  }
+
   emit<K extends keyof WsRouteHandler>(
     method: K,
     route: string,