625 lines
18 KiB
TypeScript
625 lines
18 KiB
TypeScript
import type { FastifyInstance } from 'fastify';
|
|
import type { WebSocket } from 'ws';
|
|
import { verifyJwt } from '../utils/auth.js';
|
|
import { getDb, schema } from '../db/index.js';
|
|
import { eq, inArray, desc } from 'drizzle-orm';
|
|
import { handleClientEvent } from './events.js';
|
|
import type {
|
|
User,
|
|
ServerWithChannelsAndMembers,
|
|
MemberWithUser,
|
|
Channel,
|
|
DmChannel,
|
|
ServerEvent,
|
|
ServerFolder,
|
|
ReadState,
|
|
} from '@opencord/shared';
|
|
|
|
function sanitizeUser(row: typeof schema.users.$inferSelect): User {
|
|
return {
|
|
id: row.id,
|
|
username: row.username,
|
|
displayName: row.displayName,
|
|
avatar: row.avatar,
|
|
status: (row.status ?? 'offline') as User['status'],
|
|
customStatus: row.customStatus,
|
|
createdAt: row.createdAt,
|
|
};
|
|
}
|
|
|
|
export interface AuthenticatedSocket {
|
|
ws: WebSocket;
|
|
userId: string;
|
|
username: string;
|
|
}
|
|
|
|
class ConnectionManager {
|
|
// userId → Set of WebSocket connections (multiple tabs)
|
|
private connections: Map<string, Set<WebSocket>> = new Map();
|
|
// userId → Set of server IDs the user belongs to
|
|
private userServers: Map<string, Set<string>> = new Map();
|
|
// channelId → Set of userIds in voice channel
|
|
private voiceStates: Map<string, Set<string>> = new Map();
|
|
// ws → userId (reverse lookup)
|
|
private wsToUser: Map<WebSocket, string> = new Map();
|
|
// dmChannelId → { callerId, startedAt } — active DM calls
|
|
private activeCalls: Map<string, { callerId: string; startedAt: number }> = new Map();
|
|
// userId → { isMuted, isDeafened } — voice user status (mute/deafen state)
|
|
private voiceUserStates: Map<string, { isMuted: boolean; isDeafened: boolean }> = new Map();
|
|
// userId → Timeout
|
|
private pendingOfflineTimeouts: Map<string, NodeJS.Timeout> = new Map();
|
|
|
|
addConnection(userId: string, ws: WebSocket): void {
|
|
if (!this.connections.has(userId)) {
|
|
this.connections.set(userId, new Set());
|
|
}
|
|
this.connections.get(userId)!.add(ws);
|
|
this.wsToUser.set(ws, userId);
|
|
|
|
// If they were pending offline, cancel it!
|
|
this.cancelDisconnect(userId);
|
|
}
|
|
|
|
removeConnection(ws: WebSocket): string | undefined {
|
|
const userId = this.wsToUser.get(ws);
|
|
if (!userId) return undefined;
|
|
|
|
this.wsToUser.delete(ws);
|
|
const userConnections = this.connections.get(userId);
|
|
if (userConnections) {
|
|
userConnections.delete(ws);
|
|
if (userConnections.size === 0) {
|
|
this.connections.delete(userId);
|
|
// Schedule disconnect cleanup
|
|
this.scheduleDisconnect(userId);
|
|
}
|
|
}
|
|
return userId;
|
|
}
|
|
|
|
private scheduleDisconnect(userId: string) {
|
|
if (this.pendingOfflineTimeouts.has(userId)) return;
|
|
|
|
const timeout = setTimeout(() => {
|
|
this.finalizeDisconnect(userId);
|
|
this.pendingOfflineTimeouts.delete(userId);
|
|
}, 5000); // 5 second grace period
|
|
|
|
this.pendingOfflineTimeouts.set(userId, timeout);
|
|
}
|
|
|
|
private cancelDisconnect(userId: string) {
|
|
const timeout = this.pendingOfflineTimeouts.get(userId);
|
|
if (timeout) {
|
|
clearTimeout(timeout);
|
|
this.pendingOfflineTimeouts.delete(userId);
|
|
console.log(`[ConnectionManager] Rescued session for user ${userId}`);
|
|
}
|
|
}
|
|
|
|
private finalizeDisconnect(userId: string) {
|
|
// Double check they are still offline
|
|
if (this.isUserOnline(userId)) return;
|
|
|
|
console.log(`[ConnectionManager] Finalizing disconnect for user ${userId}`);
|
|
const db = getDb();
|
|
db.update(schema.users).set({ status: 'offline' }).where(eq(schema.users.id, userId)).run();
|
|
|
|
// Leave voice if in one
|
|
const leftChannel = this.leaveAllVoice(userId);
|
|
this.clearVoiceUserStatus(userId);
|
|
if (leftChannel) {
|
|
// Get channel's server to broadcast
|
|
const channel = db.select().from(schema.channels).where(eq(schema.channels.id, leftChannel)).get();
|
|
if (channel) {
|
|
this.sendToServer(channel.serverId, {
|
|
type: 'voice_state_update',
|
|
channelId: leftChannel,
|
|
userId: userId,
|
|
action: 'leave',
|
|
});
|
|
}
|
|
}
|
|
|
|
// Broadcast offline to all servers
|
|
const userServers = this.getUserServers(userId);
|
|
for (const serverId of userServers) {
|
|
this.sendToServer(serverId, {
|
|
type: 'presence_update',
|
|
userId: userId,
|
|
status: 'offline',
|
|
});
|
|
}
|
|
}
|
|
|
|
getUserConnections(userId: string): Set<WebSocket> {
|
|
return this.connections.get(userId) ?? new Set();
|
|
}
|
|
|
|
isUserOnline(userId: string): boolean {
|
|
const conns = this.connections.get(userId);
|
|
return conns !== undefined && conns.size > 0;
|
|
}
|
|
|
|
setUserServers(userId: string, serverIds: string[]): void {
|
|
this.userServers.set(userId, new Set(serverIds));
|
|
}
|
|
|
|
addUserServer(userId: string, serverId: string): void {
|
|
if (!this.userServers.has(userId)) {
|
|
this.userServers.set(userId, new Set());
|
|
}
|
|
this.userServers.get(userId)!.add(serverId);
|
|
}
|
|
|
|
getUserServers(userId: string): Set<string> {
|
|
return this.userServers.get(userId) ?? new Set();
|
|
}
|
|
|
|
// Voice state management
|
|
joinVoice(channelId: string, userId: string): void {
|
|
// Leave any existing voice channel first
|
|
this.leaveAllVoice(userId);
|
|
if (!this.voiceStates.has(channelId)) {
|
|
this.voiceStates.set(channelId, new Set());
|
|
}
|
|
this.voiceStates.get(channelId)!.add(userId);
|
|
}
|
|
|
|
leaveVoice(channelId: string, userId: string): void {
|
|
const users = this.voiceStates.get(channelId);
|
|
if (users) {
|
|
users.delete(userId);
|
|
if (users.size === 0) {
|
|
this.voiceStates.delete(channelId);
|
|
}
|
|
}
|
|
}
|
|
|
|
leaveAllVoice(userId: string): string | null {
|
|
let leftChannelId: string | null = null;
|
|
for (const [channelId, users] of this.voiceStates) {
|
|
if (users.has(userId)) {
|
|
users.delete(userId);
|
|
if (users.size === 0) {
|
|
this.voiceStates.delete(channelId);
|
|
}
|
|
leftChannelId = channelId;
|
|
break;
|
|
}
|
|
}
|
|
return leftChannelId;
|
|
}
|
|
|
|
getVoiceUsers(channelId: string): Set<string> {
|
|
return this.voiceStates.get(channelId) ?? new Set();
|
|
}
|
|
|
|
getUserVoiceChannel(userId: string): string | null {
|
|
for (const [channelId, users] of this.voiceStates) {
|
|
if (users.has(userId)) {
|
|
return channelId;
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
// Voice user status management
|
|
setVoiceUserStatus(userId: string, isMuted: boolean, isDeafened: boolean): void {
|
|
this.voiceUserStates.set(userId, { isMuted, isDeafened });
|
|
}
|
|
|
|
getVoiceUserStatus(userId: string): { isMuted: boolean; isDeafened: boolean } | undefined {
|
|
return this.voiceUserStates.get(userId);
|
|
}
|
|
|
|
clearVoiceUserStatus(userId: string): void {
|
|
this.voiceUserStates.delete(userId);
|
|
}
|
|
|
|
getAllVoiceUserStates(): Map<string, { isMuted: boolean; isDeafened: boolean }> {
|
|
return this.voiceUserStates;
|
|
}
|
|
|
|
// DM call management
|
|
startCall(dmChannelId: string, callerId: string): boolean {
|
|
if (this.activeCalls.has(dmChannelId)) return false; // Already in a call
|
|
this.activeCalls.set(dmChannelId, { callerId, startedAt: Date.now() });
|
|
return true;
|
|
}
|
|
|
|
endCall(dmChannelId: string): void {
|
|
this.activeCalls.delete(dmChannelId);
|
|
}
|
|
|
|
getActiveCall(dmChannelId: string): { callerId: string; startedAt: number } | undefined {
|
|
return this.activeCalls.get(dmChannelId);
|
|
}
|
|
|
|
// Send to a specific user (all their connections)
|
|
sendToUser(userId: string, event: ServerEvent): void {
|
|
const connections = this.getUserConnections(userId);
|
|
const message = JSON.stringify(event);
|
|
for (const ws of connections) {
|
|
if (ws.readyState === 1) { // WebSocket.OPEN
|
|
ws.send(message);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Send to all members of a server
|
|
sendToServer(serverId: string, event: ServerEvent, excludeUserId?: string): void {
|
|
const message = JSON.stringify(event);
|
|
for (const [userId, serverIds] of this.userServers) {
|
|
if (serverIds.has(serverId) && userId !== excludeUserId) {
|
|
const connections = this.getUserConnections(userId);
|
|
for (const ws of connections) {
|
|
if (ws.readyState === 1) {
|
|
ws.send(message);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Send to all connections of all online users (for global events)
|
|
sendToAll(event: ServerEvent, excludeUserId?: string): void {
|
|
const message = JSON.stringify(event);
|
|
for (const [userId, connections] of this.connections) {
|
|
if (userId !== excludeUserId) {
|
|
for (const ws of connections) {
|
|
if (ws.readyState === 1) {
|
|
ws.send(message);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
getAllOnlineUserIds(): string[] {
|
|
return Array.from(this.connections.keys());
|
|
}
|
|
}
|
|
|
|
export const connectionManager = new ConnectionManager();
|
|
|
|
function buildReadyPayload(userId: string): {
|
|
user: User;
|
|
servers: ServerWithChannelsAndMembers[];
|
|
dmChannels: DmChannel[];
|
|
folders: ServerFolder[];
|
|
voiceStates: Record<string, string[]>;
|
|
voiceUserStates: Record<string, { isMuted: boolean; isDeafened: boolean }>;
|
|
readStates: ReadState[];
|
|
} {
|
|
const db = getDb();
|
|
|
|
// Get user
|
|
const userRow = db.select().from(schema.users).where(eq(schema.users.id, userId)).get();
|
|
if (!userRow) {
|
|
throw new Error('User not found');
|
|
}
|
|
const user = sanitizeUser(userRow);
|
|
|
|
// Get user's server memberships
|
|
const memberships = db.select()
|
|
.from(schema.serverMembers)
|
|
.where(eq(schema.serverMembers.userId, userId))
|
|
.all();
|
|
|
|
const serverIds = memberships.map(m => m.serverId);
|
|
|
|
const servers: ServerWithChannelsAndMembers[] = [];
|
|
|
|
if (serverIds.length > 0) {
|
|
const serverRows = db.select()
|
|
.from(schema.servers)
|
|
.where(inArray(schema.servers.id, serverIds))
|
|
.all();
|
|
|
|
for (const serverRow of serverRows) {
|
|
const channels = db.select()
|
|
.from(schema.channels)
|
|
.where(eq(schema.channels.serverId, serverRow.id))
|
|
.all();
|
|
|
|
const roles = db.select()
|
|
.from(schema.roles)
|
|
.where(eq(schema.roles.serverId, serverRow.id))
|
|
.orderBy(schema.roles.position)
|
|
.all();
|
|
|
|
const memberRows = db.select()
|
|
.from(schema.serverMembers)
|
|
.where(eq(schema.serverMembers.serverId, serverRow.id))
|
|
.all();
|
|
|
|
const memberUserIds = memberRows.map(m => m.userId);
|
|
const users = memberUserIds.length > 0
|
|
? db.select().from(schema.users).where(inArray(schema.users.id, memberUserIds)).all()
|
|
: [];
|
|
const userMap = new Map(users.map(u => [u.id, u]));
|
|
|
|
const memberRoleRows = db.select()
|
|
.from(schema.memberRoles)
|
|
.where(eq(schema.memberRoles.serverId, serverRow.id))
|
|
.all();
|
|
|
|
const members: MemberWithUser[] = memberRows
|
|
.map(m => {
|
|
const u = userMap.get(m.userId);
|
|
if (!u) return null;
|
|
|
|
const assignedRoleIds = memberRoleRows
|
|
.filter(mr => mr.userId === m.userId)
|
|
.map(mr => mr.roleId);
|
|
|
|
const memberRoles = roles
|
|
.filter(r => assignedRoleIds.includes(r.id))
|
|
.map(r => ({
|
|
id: r.id,
|
|
serverId: r.serverId,
|
|
name: r.name,
|
|
color: r.color ?? '#b9bbbe',
|
|
position: r.position ?? 0,
|
|
createdAt: r.createdAt,
|
|
}));
|
|
|
|
return {
|
|
serverId: m.serverId,
|
|
userId: m.userId,
|
|
role: (m.role ?? 'member') as MemberWithUser['role'],
|
|
nickname: m.nickname,
|
|
joinedAt: m.joinedAt,
|
|
user: sanitizeUser(u),
|
|
roles: memberRoles,
|
|
};
|
|
})
|
|
.filter((m): m is MemberWithUser => m !== null);
|
|
|
|
servers.push({
|
|
id: serverRow.id,
|
|
name: serverRow.name,
|
|
icon: serverRow.icon,
|
|
ownerId: serverRow.ownerId,
|
|
inviteCode: serverRow.inviteCode,
|
|
createdAt: serverRow.createdAt,
|
|
channels: channels.map(ch => {
|
|
const lastMsg = db.select({ id: schema.messages.id })
|
|
.from(schema.messages)
|
|
.where(eq(schema.messages.channelId, ch.id))
|
|
.orderBy(desc(schema.messages.createdAt))
|
|
.limit(1)
|
|
.get();
|
|
return {
|
|
id: ch.id,
|
|
serverId: ch.serverId,
|
|
name: ch.name,
|
|
type: ch.type as Channel['type'],
|
|
topic: ch.topic,
|
|
position: ch.position ?? 0,
|
|
createdAt: ch.createdAt,
|
|
lastMessageId: lastMsg?.id ?? null,
|
|
};
|
|
}),
|
|
members,
|
|
roles: roles.map(r => ({
|
|
id: r.id,
|
|
serverId: r.serverId,
|
|
name: r.name,
|
|
color: r.color ?? '#b9bbbe',
|
|
position: r.position ?? 0,
|
|
createdAt: r.createdAt,
|
|
})),
|
|
});
|
|
}
|
|
}
|
|
|
|
// Store user's server IDs for broadcasting
|
|
connectionManager.setUserServers(userId, serverIds);
|
|
|
|
// Get DM channels
|
|
const dmMemberships = db.select()
|
|
.from(schema.dmMembers)
|
|
.where(eq(schema.dmMembers.userId, userId))
|
|
.all();
|
|
|
|
const dmChannels: DmChannel[] = [];
|
|
|
|
for (const dm of dmMemberships) {
|
|
const dmChannel = db.select()
|
|
.from(schema.dmChannels)
|
|
.where(eq(schema.dmChannels.id, dm.dmChannelId))
|
|
.get();
|
|
|
|
if (!dmChannel) continue;
|
|
|
|
const dmMemberRows = db.select()
|
|
.from(schema.dmMembers)
|
|
.where(eq(schema.dmMembers.dmChannelId, dm.dmChannelId))
|
|
.all();
|
|
|
|
const dmMemberUserIds = dmMemberRows.map(m => m.userId);
|
|
const dmUsers = dmMemberUserIds.length > 0
|
|
? db.select().from(schema.users).where(inArray(schema.users.id, dmMemberUserIds)).all()
|
|
: [];
|
|
|
|
// Get last message
|
|
const lastMessage = db.select()
|
|
.from(schema.dmMessages)
|
|
.where(eq(schema.dmMessages.dmChannelId, dm.dmChannelId))
|
|
.orderBy(schema.dmMessages.createdAt)
|
|
.all();
|
|
|
|
const last = lastMessage.length > 0 ? lastMessage[lastMessage.length - 1] : null;
|
|
|
|
dmChannels.push({
|
|
id: dmChannel.id,
|
|
createdAt: dmChannel.createdAt,
|
|
members: dmUsers.map(sanitizeUser),
|
|
lastMessage: last ? {
|
|
id: last.id,
|
|
dmChannelId: last.dmChannelId,
|
|
userId: last.userId,
|
|
content: last.content,
|
|
createdAt: last.createdAt,
|
|
} : null,
|
|
});
|
|
}
|
|
|
|
// Get Server Folders
|
|
const folderRows = db.select()
|
|
.from(schema.serverFolders)
|
|
.where(eq(schema.serverFolders.userId, userId))
|
|
.orderBy(schema.serverFolders.position)
|
|
.all();
|
|
|
|
const folders: any[] = [];
|
|
for (const folder of folderRows) {
|
|
const serverIds = db.select()
|
|
.from(schema.serverFolderMembers)
|
|
.where(eq(schema.serverFolderMembers.folderId, folder.id))
|
|
.all()
|
|
.map(m => m.serverId);
|
|
|
|
folders.push({
|
|
id: folder.id,
|
|
userId: folder.userId,
|
|
name: folder.name,
|
|
color: folder.color,
|
|
position: folder.position,
|
|
serverIds,
|
|
});
|
|
}
|
|
|
|
// Build voice states — tell the client who is currently in voice channels
|
|
// across all their servers
|
|
const voiceStates: Record<string, string[]> = {};
|
|
for (const srv of servers) {
|
|
for (const ch of srv.channels) {
|
|
if (ch.type === 'voice' || ch.type === 'video') {
|
|
const users = connectionManager.getVoiceUsers(ch.id);
|
|
if (users.size > 0) {
|
|
voiceStates[ch.id] = Array.from(users);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Build voice user states — tell the client mute/deafen status of voice users
|
|
const voiceUserStates: Record<string, { isMuted: boolean; isDeafened: boolean }> = {};
|
|
for (const chId of Object.keys(voiceStates)) {
|
|
const usersInChannel = voiceStates[chId];
|
|
if (usersInChannel) {
|
|
for (const uid of usersInChannel) {
|
|
const status = connectionManager.getVoiceUserStatus(uid);
|
|
if (status) {
|
|
voiceUserStates[uid] = status;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Fetch read states for unread tracking
|
|
const readStateRows = db.select()
|
|
.from(schema.readStates)
|
|
.where(eq(schema.readStates.userId, userId))
|
|
.all();
|
|
|
|
const readStates: ReadState[] = readStateRows.map(rs => ({
|
|
channelId: rs.channelId,
|
|
lastReadMessageId: rs.lastReadMessageId,
|
|
}));
|
|
|
|
return { user, servers, dmChannels, folders, voiceStates, voiceUserStates, readStates };
|
|
}
|
|
|
|
export async function registerWebSocket(app: FastifyInstance): Promise<void> {
|
|
app.get('/ws', { websocket: true }, (socket, request) => {
|
|
const ws = socket as unknown as WebSocket;
|
|
let authenticated = false;
|
|
let userId: string | undefined;
|
|
let username: string | undefined;
|
|
|
|
// Set auth timeout - must authenticate within 10 seconds
|
|
const authTimeout = setTimeout(() => {
|
|
if (!authenticated) {
|
|
ws.send(JSON.stringify({ type: 'error', message: 'Authentication timeout' }));
|
|
ws.close();
|
|
}
|
|
}, 10000);
|
|
|
|
ws.on('message', (data: Buffer | string) => {
|
|
let parsed: Record<string, unknown>;
|
|
try {
|
|
const raw = typeof data === 'string' ? data : data.toString('utf-8');
|
|
parsed = JSON.parse(raw) as Record<string, unknown>;
|
|
} catch {
|
|
ws.send(JSON.stringify({ type: 'error', message: 'Invalid JSON' }));
|
|
return;
|
|
}
|
|
|
|
if (!authenticated) {
|
|
// First message must be auth
|
|
if (parsed.type !== 'auth' || typeof parsed.token !== 'string') {
|
|
ws.send(JSON.stringify({ type: 'error', message: 'First message must be auth' }));
|
|
ws.close();
|
|
return;
|
|
}
|
|
|
|
try {
|
|
const payload = verifyJwt(parsed.token);
|
|
userId = payload.userId;
|
|
username = payload.username;
|
|
authenticated = true;
|
|
clearTimeout(authTimeout);
|
|
|
|
// Update user status to online
|
|
const db = getDb();
|
|
db.update(schema.users).set({ status: 'online' }).where(eq(schema.users.id, userId)).run();
|
|
|
|
// Add connection
|
|
connectionManager.addConnection(userId, ws);
|
|
|
|
// Build and send ready payload
|
|
const readyData = buildReadyPayload(userId);
|
|
ws.send(JSON.stringify({
|
|
type: 'ready',
|
|
...readyData,
|
|
}));
|
|
|
|
// Broadcast presence update to all servers
|
|
const userServers = connectionManager.getUserServers(userId);
|
|
for (const serverId of userServers) {
|
|
connectionManager.sendToServer(serverId, {
|
|
type: 'presence_update',
|
|
userId,
|
|
status: 'online',
|
|
}, userId);
|
|
}
|
|
} catch {
|
|
ws.send(JSON.stringify({ type: 'error', message: 'Invalid token' }));
|
|
ws.close();
|
|
}
|
|
return;
|
|
}
|
|
|
|
// Handle authenticated events
|
|
if (userId && username) {
|
|
handleClientEvent(parsed, userId, username);
|
|
}
|
|
});
|
|
|
|
ws.on('close', () => {
|
|
clearTimeout(authTimeout);
|
|
if (userId) {
|
|
connectionManager.removeConnection(ws);
|
|
}
|
|
});
|
|
|
|
ws.on('error', () => {
|
|
clearTimeout(authTimeout);
|
|
});
|
|
});
|
|
}
|