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 of route
    name: string;
    
    // ZodSchema used to validate request body
    schema: ZodSchema<T>;

    // Handler Function. To be called after request body is validated
    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";
    }
}

// Context holds all available information for any given request
export interface Context<T> {
    ws: WebSocketInterface;
    user: User;
    payload: T;
    server: HapticLinkServer;
}

export class HapticLinkServer {
    // Router map
    routes: { [key: string]: Route<any> };

    // All rooms are stored in this map
    rooms: { [key: string]: Room };

    // All users are stored in this map and should be dropped when their connection does.
    users: Map<WebSocketInterface, User>;

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

    /**
     * Removes user from registered users
     * @param {WebSocketInterface} ws 
     * @returns {boolean} true/false whether the user was successfully removed
     */
    removeUser(ws: WebSocketInterface): boolean {
        const user = this.users.get(ws);
        if (!user) return false;
        if (user.currentRoom) {
            user.currentRoom.removeUserById(user.id);
        }
        return true;
    }

    /**
     * Registers a route with router
     * @param {string} name The name of the route which be used by the client to identify the route.
     * @param {ZodSchema<T>} schema A ZodSchema to match the body of the WS message
     * @param {RouteHandler<T>} handler A handler function that should have a matching schema
     * @returns 
     */
    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;
    }

    /**
     * Compares message param with names of registered routes
     * If route exists, validates the schema and then calls
     * the handler function with the formatted payload
     * @param {WebSocketInterface} ws
     * @param {string} message WebSocket message
     * @returns 
     */
    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" }));
        }

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

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

        const route = this.routes[payload.route];
        // Removes route from body
        delete payload.route;

        let user: User;

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

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

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

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