server/src/socket/hapticLinkServer.ts

import { ZodSchema } from "zod";
import { generateSessionToken } from "../helpers";
import { Room } from "./room";
import { WebSocketInterface } from "./WebSocketAdapter";

type RouteHandler<T> = (context: Context<T>) => void;

export interface Route<T> {
    name: string;
    schema: ZodSchema<T>;
    handler: RouteHandler<T>;
}

export class User {
    static userIdLength: number = 32;
    id: string;
    username?: string;
    socket: WebSocketInterface;
    currentRoom?: Room;

    constructor(socket: WebSocketInterface) {
        this.id = generateSessionToken(32);
        this.socket = socket;
        this.username = "Unknown";
    }
}

export interface Context<T> {
    ws: WebSocketInterface;
    user: User;
    payload: T;
    server: HapticLinkServer;
}

export class HapticLinkServer {
    routes: { [key: string]: Route<any> };
    rooms: { [key: string]: Room };
    users: Map<WebSocketInterface, User>;

    constructor() {
        this.routes = {};
        this.users = new Map();
        this.rooms = {};
    }

    removeUser(ws: WebSocketInterface): boolean {
        const user = this.users.get(ws);
        if (!user) return false;
        if (user.currentRoom) {
            user.currentRoom.removeUserById(user.id);
        }
        return true;
    }

    addRoute<T>(name: string, schema: ZodSchema<T>, handler: RouteHandler<T>): boolean {
        if (name in this.routes) {
            return false;
        }

        this.routes[name] = {
            name,
            schema,
            handler,
        };

        return true;
    }

    handleRoute(ws: WebSocketInterface, message: string) {
        // Parse JSON
        let payload: any;
        try {
            payload = JSON.parse(message);
        } catch (e) {
            return ws.send(JSON.stringify({ error: "message not in JSON format" }));
        }

        if (typeof payload != "object" || !Object.keys(payload).includes("route")) {
            return ws.send(JSON.stringify({ error: "missing route" }));
        }

        if (!(payload.route in this.routes)) {
            return ws.send(JSON.stringify({ error: "route not found" }));
        }

        const route = this.routes[payload.route];
        delete payload.route;

        let user: User;

        if (this.users.has(ws)) {
            user = this.users.get(ws)!;
        } else {
            const newUser = new User(ws);
            this.users.set(ws, newUser);
            user = newUser;
        }

        let context: Context<any> = {
            ws,
            payload,
            server: this,
            user,
        };

        if (!route) {
            return ws.send(JSON.stringify({ error: "invalid route" }));
        }

        if (route.schema.safeParse(payload).success) {
            route.handler(context);
        } else {
            return ws.send(JSON.stringify({ error: "invalid payload format" }));
        }
    }
}