First Commit. Duplicate bunext-sqlite preset

src/websocket/(data)/rooms.ts

@@ -0,0 +1,37 @@
+import type { ServerWebSocket } from "bun";
+import type { WebSocketSession } from "../types";
+import type { WebsocketRoomType } from "@/src/types";
+
+export type ClientMetadata = {
+    user_id?: string | number;
+    rooms: Set<WebsocketRoomType>;
+    joinedAt: number;
+};
+
+const clients = new WeakMap<
+    ServerWebSocket<WebSocketSession>,
+    ClientMetadata
+>();
+const rooms = new Map<
+    WebsocketRoomType,
+    Set<ServerWebSocket<WebSocketSession>>
+>();
+
+export function getMetadata(
+    ws: ServerWebSocket<WebSocketSession>,
+): ClientMetadata {
+    let meta = clients.get(ws);
+    if (!meta) {
+        meta = { rooms: new Set(), joinedAt: Date.now() };
+        clients.set(ws, meta);
+    }
+    return meta;
+}
+
+export function getAllClients() {
+    return clients;
+}
+
+export function getAllRooms() {
+    return rooms;
+}

src/websocket/(functions)/broadcast.ts

@@ -0,0 +1,23 @@
+import type { WebSocketDataType, WebsocketRoomType } from "@/src/types";
+import { getAllRooms } from "../(data)/rooms";
+import type { ServerWebSocket } from "bun";
+import type { WebSocketSession } from "../types";
+
+export function broadcast(
+    room: WebsocketRoomType,
+    data: WebSocketDataType,
+    exclude?: ServerWebSocket<WebSocketSession>,
+) {
+    const rooms = getAllRooms();
+    const members = rooms.get(room);
+
+    if (!members) return;
+
+    const message = JSON.stringify(data);
+
+    for (const ws of members) {
+        if (ws !== exclude && ws.readyState === WebSocket.OPEN) {
+            ws.send(message);
+        }
+    }
+}

src/websocket/(functions)/check-session-access.ts

@@ -0,0 +1,33 @@
+import grabUserAuthTypes from "@/src/functions/backend/auth/grab-user-auth-types";
+import type { BUN_SQLITE_BSQLITE_USER_TYPES } from "@/db/types/db";
+import type { WebSocketSession } from "../types";
+
+type Params = {
+    session: WebSocketSession;
+};
+
+type Return = {
+    allowed: boolean;
+    user_types?: BUN_SQLITE_BSQLITE_USER_TYPES[];
+};
+
+/**
+ * # Check WebSocket Session Access
+ * The session is frozen into the socket at upgrade time, so access is
+ * re-checked while the socket is alive. This makes revocation immediate
+ * instead of waiting for the client to disconnect.
+ */
+export default async function checkSessionAccess({
+    session,
+}: Params): Promise<Return> {
+    if (!session.user?.id) {
+        return { allowed: false };
+    }
+
+    const res = await grabUserAuthTypes({ user_id: session.user.id });
+
+    return {
+        allowed: res.success && Boolean(res.user_types?.[0]),
+        user_types: res.user_types ?? undefined,
+    };
+}

src/websocket/(functions)/handle-message.ts

@@ -0,0 +1,65 @@
+import type { WebSocketDataType } from "@/src/types";
+import type { ServerWebSocket } from "bun";
+import type { WebSocketSession } from "../types";
+import { wsSendData } from "./ws-send-data";
+import wsPing from "../client-requests/ping";
+import { logError } from "@/src/functions/backend/logger";
+
+export function handleMessage(
+    ws: ServerWebSocket<WebSocketSession>,
+    data?: WebSocketDataType,
+) {
+    const { user } = ws.data;
+
+    try {
+        if (!data?.event) {
+            return wsSendData({
+                ws,
+                data: {
+                    event: "error",
+                    error: `No Data Provided`,
+                },
+            });
+        }
+
+        const { event } = data;
+
+        if (event === "client:ping") {
+            return wsPing({ data, ws });
+        }
+
+        if (event === "ready" || event === "message") {
+            wsSendData({
+                ws,
+                data: { event: "success", message: "connected" },
+            });
+            return;
+        }
+
+        wsSendData({
+            ws,
+            data: {
+                event: "error",
+                error: `Unknown event: ${event}`,
+            },
+        });
+    } catch (error: any) {
+        logError({
+            title: "WebSocket message error",
+            category: "websocket",
+            user_id: user.id,
+            error,
+            metadata: {
+                event: data?.event,
+            },
+        });
+
+        wsSendData({
+            ws,
+            data: {
+                event: "error",
+                error: `ERROR: ${error.message}`,
+            },
+        });
+    }
+}

src/websocket/(functions)/join-room.ts

@@ -0,0 +1,15 @@
+import { getMetadata, getAllRooms } from "../(data)/rooms";
+import type { ServerWebSocket } from "bun";
+import type { WebSocketSession } from "../types";
+import type { WebsocketRoomType } from "@/src/types";
+
+export function joinRoom(
+    ws: ServerWebSocket<WebSocketSession>,
+    room: WebsocketRoomType,
+) {
+    const meta = getMetadata(ws);
+    const rooms = getAllRooms();
+    meta.rooms.add(room);
+    if (!rooms.has(room)) rooms.set(room, new Set());
+    rooms.get(room)!.add(ws);
+}

src/websocket/(functions)/leave-all-rooms.ts

@@ -0,0 +1,11 @@
+import { getMetadata } from "../(data)/rooms";
+import { leaveRoom } from "./leave-room";
+import type { ServerWebSocket } from "bun";
+import type { WebSocketSession } from "../types";
+
+export function leaveAllRooms(ws: ServerWebSocket<WebSocketSession>) {
+    const meta = getMetadata(ws);
+    for (const room of meta.rooms) {
+        leaveRoom(ws, room);
+    }
+}

src/websocket/(functions)/leave-room.ts

@@ -0,0 +1,15 @@
+import { getMetadata, getAllRooms } from "../(data)/rooms";
+import type { ServerWebSocket } from "bun";
+import type { WebSocketSession } from "../types";
+import type { WebsocketRoomType } from "@/src/types";
+
+export function leaveRoom(
+    ws: ServerWebSocket<WebSocketSession>,
+    room: WebsocketRoomType,
+) {
+    const meta = getMetadata(ws);
+    const rooms = getAllRooms();
+    meta.rooms.delete(room);
+    rooms.get(room)?.delete(ws);
+    if (rooms.get(room)?.size === 0) rooms.delete(room);
+}

src/websocket/(functions)/ws-send-data.ts

@@ -0,0 +1,14 @@
+import type { WebSocketDataType } from "@/src/types";
+import type { ServerWebSocket } from "bun";
+import type { WebSocketSession } from "../types";
+
+type Params = {
+    ws: ServerWebSocket<WebSocketSession>;
+    data: WebSocketDataType;
+};
+
+export function wsSendData({ data, ws }: Params) {
+    if (ws.readyState === WebSocket.OPEN) {
+        ws.send(JSON.stringify(data));
+    }
+}

src/websocket/client-requests/ping/index.ts

@@ -0,0 +1,20 @@
+import type { WebSocketDataType } from "@/src/types";
+import type { ServerWebSocket } from "bun";
+import type { WebSocketSession } from "../../types";
+import { wsSendData } from "../../(functions)/ws-send-data";
+
+type Params = {
+    data: WebSocketDataType;
+    ws: ServerWebSocket<WebSocketSession>;
+};
+
+export default function wsPing({ ws }: Params) {
+    if (ws.readyState === WebSocket.OPEN) {
+        wsSendData({
+            ws,
+            data: {
+                event: "server:pong",
+            },
+        });
+    }
+}

src/websocket/index.ts

@@ -0,0 +1,85 @@
+import type { ServerWebSocket } from "bun";
+import type { WebSocketSession } from "./types";
+import { joinRoom } from "./(functions)/join-room";
+import { leaveRoom } from "./(functions)/leave-room";
+import { leaveAllRooms } from "./(functions)/leave-all-rooms";
+import { broadcast } from "./(functions)/broadcast";
+import { handleMessage } from "./(functions)/handle-message";
+import { getAllClients, getAllRooms } from "./(data)/rooms";
+import type { WebSocketDataType } from "@/src/types";
+import { wsSendData } from "./(functions)/ws-send-data";
+import checkSessionAccess from "./(functions)/check-session-access";
+import logEvent from "@/src/functions/backend/logger";
+
+/**
+ * The twui `useWebSocket` hook sends a raw keepalive string to keep the
+ * connection alive. Non-JSON messages are safe to ignore.
+ */
+function parseMessage(message: string | Buffer): WebSocketDataType | undefined {
+    const text = message.toString();
+
+    try {
+        return JSON.parse(text) as WebSocketDataType | undefined;
+    } catch (error) {
+        return undefined;
+    }
+}
+
+export const websocket = {
+    open(ws: ServerWebSocket<WebSocketSession>) {
+        const { user } = ws.data;
+
+        joinRoom(ws, "public");
+
+        wsSendData({
+            ws,
+            data: {
+                event: "ready",
+                message: "WebSocket connected",
+            },
+        });
+
+        logEvent({
+            title: "WebSocket connected",
+            category: "websocket",
+            user_id: user.id,
+            plain_text: `User #${user.id} [${user.first_name} ${user.last_name}] connected via WebSocket.`,
+            metadata: {
+                user_id: user.id,
+            },
+        });
+    },
+
+    async message(
+        ws: ServerWebSocket<WebSocketSession>,
+        message: string | Buffer,
+    ) {
+        const data = parseMessage(message);
+
+        if (!data) return;
+
+        const access = await checkSessionAccess({ session: ws.data });
+
+        if (!access.allowed) {
+            ws.close(4401, "Unauthorized");
+            return;
+        }
+
+        if (access.user_types) {
+            ws.data.user_types = access.user_types;
+        }
+
+        handleMessage(ws, data);
+    },
+
+    close(ws: ServerWebSocket<WebSocketSession>) {
+        leaveAllRooms(ws);
+    },
+
+    drain(ws: ServerWebSocket<WebSocketSession>) {},
+    broadcast,
+    joinRoom,
+    leaveRoom,
+    getClients: getAllClients,
+    getRooms: getAllRooms,
+};

src/websocket/types.ts

@@ -0,0 +1,7 @@
+import type { User } from "@/src/types";
+import type { BUN_SQLITE_BSQLITE_USER_TYPES } from "@/db/types/db";
+
+export type WebSocketSession = {
+    user: User;
+    user_types: BUN_SQLITE_BSQLITE_USER_TYPES[];
+};