itpdp/api/src/routes/party-socket.ts

import { Elysia } from "elysia";

import { betterAuthElysia } from "../auth";

import { db } from "../db";
import { getMemberRecord, getPartyStatus } from "../party-data";
import type { PartySocketEvent } from "../party-types";

export function userTopic(userId: string) {
	return `user:${userId}`;
}

export function partyTopic(partyId: string) {
	return `party:${partyId}`;
}

export const socketPartyId = new WeakMap<object, string>();

export const pubsub = {
	_server: null as ReturnType<typeof Bun.serve> | null,
	setServer(server: ReturnType<typeof Bun.serve> | null) {
		this._server = server;
	},
	publish(topic: string, data: string) {
		this._server?.publish(topic, data);
	},
	publishPartyData(partyId: string, data: PartySocketEvent) {
		pubsub.publish(`party:${partyId}`, JSON.stringify(data));
	},
};

async function subscribeWsToParty(
	ws: {
		subscribe: (topic: string) => void;
		unsubscribe: (topic: string) => void;
		send: (message: string) => void;
	},
	userId: string,
) {
	const membership = await getMemberRecord(db, userId);
	if (!membership) return null;

	const nextPartyId = membership.partyId;
	const currentPartyId = socketPartyId.get(ws as object);
	if (currentPartyId && currentPartyId !== nextPartyId) {
		ws.unsubscribe(partyTopic(currentPartyId));
	}

	socketPartyId.set(ws as object, nextPartyId);
	ws.subscribe(partyTopic(nextPartyId));

	const snapshot = await getPartyStatus(nextPartyId);
	if (snapshot) {
		ws.send(
			JSON.stringify({
				type: "party_status",
				party: snapshot.party,
				members: snapshot.members,
			} satisfies PartySocketEvent),
		);
	}

	return nextPartyId;
}

export const topic = {
	user: userTopic,
	party: partyTopic,
};

export const partySocketApp = new Elysia()
	.use(betterAuthElysia)
	.group("/party-socket", (app) =>
		app
			.get("/test", () => ({ ok: 1 }))
			.ws("/ws", {
				auth: true,
				publishToSelf: true,
				open: async (ws) => {
					const user = ws.data.user;
					if (!user) return;

					ws.subscribe(userTopic(user.id));

					const subscribedPartyId = await subscribeWsToParty(ws, user.id);
					if (!subscribedPartyId) {
						ws.send(
							JSON.stringify({
								type: "party_status",
								party: null,
								members: [],
							} as PartySocketEvent),
						);
						return;
					}
				},
				message: async (ws, message) => {
					const data = ws.data;
					const user = data.user;
					if (!user) return;

					if (typeof message !== "string") return;

					let parsed: { type: string; payload?: unknown };
					try {
						parsed = JSON.parse(message);
					} catch {
						ws.send(JSON.stringify({ type: "error", message: "Invalid JSON" }));
						return;
					}

					if (parsed.type === "ping") {
						ws.send(JSON.stringify({ type: "pong" }));
						return;
					}

					if (parsed.type === "subscribe_party") {
						const payload = parsed as {
							type: "subscribe_party";
							partyId: string;
						};
						if (typeof payload.partyId !== "string") return;
						await subscribeWsToParty(ws, user.id);
						return;
					}

					if (parsed.type !== "member_payload") return;

					const MAX_MEMBER_PAYLOAD_SIZE = 8_000;
					const payloadString = JSON.stringify(parsed.payload);
					if (payloadString.length > MAX_MEMBER_PAYLOAD_SIZE) {
						ws.send(
							JSON.stringify({ type: "error", message: "Payload too large." }),
						);
						return;
					}

					const membership = await getMemberRecord(db, user.id);
					if (!membership) return;

					const currentParty = await db.query.party.findFirst({
						where: { id: membership.partyId },
					});
					if (!currentParty) return;

					ws.publish(
						partyTopic(membership.partyId),
						JSON.stringify({
							type: "member_payload",
							fromUserId: user.id,
							payload: parsed.payload,
						} as PartySocketEvent),
					);
				},
				close: async (ws) => {
					const user = ws.data.user;
					if (!user) return;

					ws.unsubscribe(userTopic(user.id));

					const partyId = socketPartyId.get(ws);
					if (!partyId) return;

					ws.unsubscribe(partyTopic(partyId));
				},
			}),
	);