frm.so/server/ws/ws.js

import { WebSocket, WebSocketServer } from 'ws';
import { z } from 'zod';
import jwt from 'jsonwebtoken';

import * as serverFunctions from "../../ui/_/code/bridge/serverFunctions.js"
import ForumHandler from "./handlers/ForumHandler.js"
import MessagesHandler from "./handlers/MessagesHandler.js"

export default class Socket {
    wss;
    messageSchema = z.object({
        id: z.string(),
        app: z.string().optional(),
        operation: z.string().optional(),
        msg: z.union([
            z.object({}).passthrough(), // allows any object
            z.array(z.any())            // allows any array
        ]).optional()
    })
    .superRefine((data, ctx) => {
        if (data.operation !== "GET" && data.msg === undefined) {
            ctx.addIssue({
                code: z.ZodIssueCode.custom,
                path: ["msg"],
                message: "msg is required when operation is not GET"
            })
        }
    })
    .strict()

    constructor(server) {
        this.wss = new WebSocketServer({ server });

        this.wss.on('connection', (ws, req) => {
            console.log('✅ New WebSocket client connected');

            function parseCookies(cookieHeader = "") {
                return Object.fromEntries(
                    cookieHeader.split(";").map(c => {
                        const [key, ...v] = c.trim().split("=");
                        return [key, decodeURIComponent(v.join("="))];
                    })
                );
            }

            const cookies = parseCookies(req.headers.cookie);
            const token = cookies.auth_token;
            if (!token) throw new Error("No auth token");
            const payload = jwt.verify(token, process.env.JWT_SECRET);
            ws.userEmail = payload.email;
            
            ws.on('message', (msg) => {
                this.handleMessage(msg, ws);
            });

            ws.on('close', () => {
                console.log('Client disconnected');
            });
        });

        console.log('WebSocket server initialized');
    }

    // Build a system where the ws obj is updated every time on navigate, so it already has context
    // this way, we can only send broadcast messages to clients that actually have that app / subapp open
    handleMessage = async (msg, ws) => {
        try {
            const text = msg.toString();
            const req = JSON.parse(text);
            if(!this.messageSchema.safeParse(req).success) throw new Error("Socket.handleMessage: Incoming ws message has incorrect format!")
            
            let responseData;
            switch (req.app) {
                case "FORUM":
                    responseData = await ForumHandler.handle(req.operation, req.msg, ws)
                    break;

                case "MESSAGES":
                    responseData = MessagesHandler.handle(req.operation, req.msg, ws)
                    break;

                default:
                    if(!req.app) {
                        let func = req.msg
                        responseData = serverFunctions[func.name](...args)
                    } else {
                        console.error("unknown ws message")
                    }
            }

            let response = {
                ...req
            }
            response.msg = responseData

            if(!this.messageSchema.safeParse(response).success) throw new Error("Socket.handleMessage: Outgoing ws message has incorrect format!")
            ws.send(JSON.stringify(response))

        } catch (e) {
            console.error("Invalid WS message:", e);
        }
    }

    broadcast(event) {
        if (!this.wss) return;

        let message = JSON.stringify(event)

        this.wss.clients.forEach(ws => {
                if (ws.readyState === WebSocket.OPEN) {
                ws.send(message);
            }
        });
    }
}