Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 17 additions & 7 deletions apps/realtime/src/handlers/connection.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import { createLogger } from '@sim/logger'
import { parseRoomName, type RoomRef, roomName } from '@sim/realtime-protocol/rooms'
import { parseRoomName, ROOM_TYPES, type RoomRef, roomName } from '@sim/realtime-protocol/rooms'
import { cleanupFileDocForSocket } from '@/handlers/file-doc'
import { cleanupPendingSubblocksForSocket } from '@/handlers/subblocks'
import { cleanupPendingVariablesForSocket } from '@/handlers/variables'
Expand All @@ -8,6 +8,14 @@ import type { IRoomManager } from '@/rooms'

const logger = createLogger('ConnectionHandlers')

/**
* Room types whose presence lives in the room manager (Redis-backed), so a disconnect must
* remove the socket + broadcast a correction. The workspace-files and file-doc rooms are
* NOT here: workspace-files carries no presence (native Socket.IO membership only), and
* file-doc broadcasts its own server-authenticated roster via `cleanupFileDocForSocket`.
*/
const PRESENCE_BEARING_TYPES = new Set<RoomRef['type']>([ROOM_TYPES.WORKFLOW, ROOM_TYPES.TABLE])

export function setupConnectionHandlers(socket: AuthenticatedSocket, roomManager: IRoomManager) {
socket.on('error', (error) => {
logger.error(`Socket ${socket.id} error:`, error)
Expand All @@ -33,8 +41,9 @@ export function setupConnectionHandlers(socket: AuthenticatedSocket, roomManager
cleanupPendingSubblocksForSocket(socket.id)
cleanupPendingVariablesForSocket(socket.id)
// Clear the socket's collaborative-document awareness (removes its caret for
// everyone else) and drop the room if it was the last editor.
cleanupFileDocForSocket(socket.id, roomManager.io)
// everyone else) and drop the room if it was the last editor. `endOfLife` drops the
// socket's join-generation entry — safe only here, on true disconnect (see cleanup).
cleanupFileDocForSocket(socket.id, roomManager.io, true)

// A socket may occupy multiple rooms (one per type). Remove it from every
// room the manager knows about.
Expand All @@ -47,12 +56,13 @@ export function setupConnectionHandlers(socket: AuthenticatedSocket, roomManager
const wasInRooms = new Map<string, RoomRef>()
for (const room of removedRooms) wasInRooms.set(roomName(room), room)
for (const name of liveRoomNames) {
// `wasInRooms.has(name)` already excludes every room the manager removed
// (same room-name key via the roomName/parseRoomName bijection), so any
// room reaching here was NOT in `removedRooms` and needs a removal attempt.
// `wasInRooms.has(name)` already excludes every room the manager removed (same
// room-name key via the roomName/parseRoomName bijection). Skip room types with no
// manager-tracked presence (workspace-files, file-doc): removing there is a no-op and
// broadcasting a correction would emit a dead presence-update no client listens to.
if (name === socket.id || wasInRooms.has(name)) continue
const ref = parseRoomName(name)
if (!ref) continue
if (!ref || !PRESENCE_BEARING_TYPES.has(ref.type)) continue
wasInRooms.set(name, ref)
await roomManager.removeUserFromRoom(ref, socket.id)
}
Expand Down
87 changes: 85 additions & 2 deletions apps/realtime/src/handlers/file-doc.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,8 @@ function createSocket(id: string, overrides?: Record<string, unknown>) {
id,
userId: 'user-1',
userName: 'Test User',
// Set so the server's roster resolves the avatar from the socket (never the DB).
userImage: 'avatar.png',
disconnected: false,
on: vi.fn((event: string, handler: Handler) => {
handlers[event] = handler
Expand Down Expand Up @@ -133,7 +135,9 @@ describe('setupWorkspaceFileDocHandlers', () => {
afterEach(() => {
// The room store is module-global; drop every room the test's sockets opened.
const { io } = createIo()
for (const id of createdSocketIds) cleanupFileDocForSocket(id, io)
// Simulate a full disconnect between tests (`endOfLife`) so the module-global join-generation
// map is cleared and never bleeds a counter into the next test.
for (const id of createdSocketIds) cleanupFileDocForSocket(id, io, true)
createdSocketIds.clear()
vi.clearAllTimers()
vi.useRealTimers()
Expand Down Expand Up @@ -435,7 +439,7 @@ describe('setupWorkspaceFileDocHandlers', () => {

const pending = s.handlers[FILE_DOC_EVENTS.JOIN]({ fileId: 'file-1', clientId: 1 })
s.socket.disconnected = true
cleanupFileDocForSocket('socket-a', io) // disconnect cleanup — no-op, nothing registered yet
cleanupFileDocForSocket('socket-a', io, true) // disconnect cleanup — no-op, nothing registered yet
resolveAuth({ allowed: true, status: 200, workspacePermission: 'write' })
await pending

Expand All @@ -462,6 +466,49 @@ describe('setupWorkspaceFileDocHandlers', () => {
expect(s.socket.join).toHaveBeenCalledWith('workspace-file-doc:file-2')
})

it('does not reset the join generation on a leave, so an in-flight join still binds', async () => {
const { io } = createIo()
const s = setup('socket-a', io)

// file-1 join completes; the socket is registered in file-1.
await s.handlers[FILE_DOC_EVENTS.JOIN]({ fileId: 'file-1', clientId: 1 })

// file-2 join goes in-flight (authorize deferred).
let resolveAuth: (v: unknown) => void = () => {}
mockAuthorizeRoom.mockReturnValueOnce(new Promise((resolve) => (resolveAuth = resolve)))
const pending = s.handlers[FILE_DOC_EVENTS.JOIN]({ fileId: 'file-2', clientId: 1 })

// A deferred leave for the prior file-1 lands while file-2's join awaits authorization. Its
// cleanup must NOT reset the monotonic join generation, or file-2's guard would see an emptied
// map (`undefined !== generation`) and abort the join the client actually wants.
s.handlers[FILE_DOC_EVENTS.LEAVE]({ fileId: 'file-1' })

resolveAuth({ allowed: true, status: 200, workspacePermission: 'write' })
await pending

expect(joinSuccessFileId(s.socket)).toBe('file-2')
expect(s.socket.join).toHaveBeenCalledWith('workspace-file-doc:file-2')
})

it('cancels an in-flight join when the client leaves that same file (no ghost owner)', async () => {
const { io, sent } = createIo()
let resolveAuth: (v: unknown) => void = () => {}
mockAuthorizeRoom.mockReturnValueOnce(new Promise((resolve) => (resolveAuth = resolve)))
const s = setup('socket-a', io)

// Join file-1 is awaiting authorization when the client leaves file-1 (fast open→close).
const pending = s.handlers[FILE_DOC_EVENTS.JOIN]({ fileId: 'file-1', clientId: 1 })
s.handlers[FILE_DOC_EVENTS.LEAVE]({ fileId: 'file-1' })
resolveAuth({ allowed: true, status: 200, workspacePermission: 'write' })
await pending

// The stale join must not register: no success, no room join, and no presence broadcast that
// would leave a ghost collaborator until disconnect.
expect(s.socket.join).not.toHaveBeenCalled()
expect(joinSuccessFileId(s.socket)).toBeUndefined()
expect(sent.some((m) => m.event === FILE_DOC_EVENTS.PRESENCE)).toBe(false)
})

it('scopes LEAVE to the named file (a leave for a different file is a no-op)', async () => {
const { io } = createIo()
const a = setup('socket-a', io)
Expand Down Expand Up @@ -574,4 +621,40 @@ describe('setupWorkspaceFileDocHandlers', () => {
// No re-election: the successful seed cancelled the deadline.
expect(sent.find((m) => m.event === FILE_DOC_EVENTS.SEED_REQUEST)).toBeUndefined()
})

it('broadcasts a server-authenticated presence roster on join, one entry per session', async () => {
const { io, sent } = createIo()
const a = setup('socket-a', io, { userId: 'user-a', userName: 'Ada', userImage: 'ada.png' })
const b = setup('socket-b', io, { userId: 'user-b', userName: 'Bob', userImage: 'bob.png' })

await a.handlers[FILE_DOC_EVENTS.JOIN]({ fileId: 'file-1', clientId: 1 })
await b.handlers[FILE_DOC_EVENTS.JOIN]({ fileId: 'file-1', clientId: 2 })

const roster = sent.filter((m) => m.event === FILE_DOC_EVENTS.PRESENCE).at(-1)?.payload as {
fileId: string
users: Array<{ socketId: string; userId: string; userName: string; avatarUrl: string | null }>
}
expect(roster.fileId).toBe('file-1')
// Identity is each socket's authenticated session — not any client-supplied value.
expect([...roster.users].sort((x, y) => x.userId.localeCompare(y.userId))).toEqual([
{ socketId: 'socket-a', userId: 'user-a', userName: 'Ada', avatarUrl: 'ada.png' },
{ socketId: 'socket-b', userId: 'user-b', userName: 'Bob', avatarUrl: 'bob.png' },
])
})

it('keeps a per-session entry for two sockets of the SAME user (no server-side user dedup)', async () => {
const { io, sent } = createIo()
// Two tabs of one account: the client self-excludes its own socket, so the roster must carry
// BOTH sessions or a client could never see the other tab as present.
const a = setup('socket-a', io, { userId: 'user-a', userName: 'Ada', userImage: 'ada.png' })
const b = setup('socket-b', io, { userId: 'user-a', userName: 'Ada', userImage: 'ada.png' })

await a.handlers[FILE_DOC_EVENTS.JOIN]({ fileId: 'file-1', clientId: 1 })
await b.handlers[FILE_DOC_EVENTS.JOIN]({ fileId: 'file-1', clientId: 2 })

const roster = sent.filter((m) => m.event === FILE_DOC_EVENTS.PRESENCE).at(-1)?.payload as {
users: Array<{ socketId: string; userId: string }>
}
expect([...roster.users].map((u) => u.socketId).sort()).toEqual(['socket-a', 'socket-b'])
})
})
86 changes: 73 additions & 13 deletions apps/realtime/src/handlers/file-doc.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import {
FILE_DOC_EVENTS,
FILE_DOC_MESSAGE_TYPE,
FILE_DOC_SEED,
type FileDocPresenceUser,
type JoinFileDocPayload,
type LeaveFileDocPayload,
toFileDocBytes,
Expand All @@ -15,6 +16,7 @@ import type { Server } from 'socket.io'
import * as awarenessProtocol from 'y-protocols/awareness'
import * as syncProtocol from 'y-protocols/sync'
import * as Y from 'yjs'
import { resolveAvatarUrl } from '@/handlers/avatar'
import type { AuthenticatedSocket } from '@/middleware/auth'
import type { IRoomManager } from '@/rooms'

Expand Down Expand Up @@ -57,6 +59,10 @@ interface FileDocOwner {
/** The owning user — used to tell a reconnect (same user reusing its Yjs client
* id) from a spoof (a different user binding a peer's id). */
userId: string
/** Server-authenticated display identity for the presence roster (from the socket's
* session, never the client-set awareness — so a peer cannot spoof it). */
userName: string
avatarUrl: string | null
}

interface FileDocRoom {
Expand Down Expand Up @@ -116,6 +122,28 @@ function broadcast(io: Server, name: string, payload: Uint8Array, exceptSocketId
channel.emit(FILE_DOC_EVENTS.MESSAGE, payload)
}

/**
* Broadcast the room's collaborator roster to everyone in it, for the avatar stack. One entry
* PER SESSION (socket) — the client excludes its own socket and dedupes the remainder per user
* for display, so a second tab of the same account still registers as present (mirroring the
* canvas presence model). Deduping here instead would drop the current user's other sessions
* asymmetrically (only one socket survives), so each client could never reliably self-exclude.
* Identity comes from each owner's server-authenticated session — never the client-set awareness
* — so a peer cannot spoof or suppress an entry.
*/
function broadcastFileDocPresence(io: Server, name: string, room: FileDocRoom) {
const users: FileDocPresenceUser[] = []
for (const [socketId, owner] of room.owners) {
users.push({
socketId,
userId: owner.userId,
userName: owner.userName,
avatarUrl: owner.avatarUrl,
})
}
io.to(name).emit(FILE_DOC_EVENTS.PRESENCE, { fileId: room.fileId, users })
}

/** Whether the client has recorded that it seeded the document's initial content. */
function isDocSeeded(doc: Y.Doc): boolean {
return doc.getMap(FILE_DOC_SEED.configMap).get(FILE_DOC_SEED.flag) === true
Expand Down Expand Up @@ -316,10 +344,15 @@ function handleMessage(socket: AuthenticatedSocket, data: unknown) {
* seeding, and drop the room's document when the last collaborator leaves.
* Exported for the disconnect handler; safe to call for a socket in no room.
*/
export function cleanupFileDocForSocket(socketId: string, io: Server): void {
// Drop the join generation so an in-flight JOIN for this socket aborts after
// its authorize resolves, and the map never leaks across the socket's life.
joinGeneration.delete(socketId)
export function cleanupFileDocForSocket(socketId: string, io: Server, endOfLife = false): void {
// The join-generation counter is monotonic for the socket's WHOLE life and must survive a room
// switch/leave: resetting it here would let the next join reuse a low number that a still
// in-flight earlier join also holds, so that stale join passes the generation guard and rebinds
// the socket to the wrong document. Drop it ONLY when the socket is truly gone (disconnect),
// which is also the only place the map would otherwise leak. An in-flight join is already
// aborted on disconnect by the `socket.disconnected` check, and on a switch by a newer join
// bumping the generation — neither needs this delete.
if (endOfLife) joinGeneration.delete(socketId)

const name = socketToRoomName.get(socketId)
if (!name) return
Expand All @@ -334,6 +367,8 @@ export function cleanupFileDocForSocket(socketId: string, io: Server): void {
// Fires the awareness `update` handler with a non-socket origin → the removal
// is broadcast to every remaining client, so the departed caret vanishes.
awarenessProtocol.removeAwarenessStates(room.awareness, [owner.clientId], null)
// Refresh the roster for whoever remains (server-authenticated identity).
broadcastFileDocPresence(io, name, room)
}

// Hand off the seeder role: if the elected seeder left before it seeded, elect
Expand All @@ -351,12 +386,22 @@ export function cleanupFileDocForSocket(socketId: string, io: Server): void {
* file id; joining requires workspace `write` (editing a document). Mirrors the
* workspace-files join shape (auth → readiness → validate → authorize → join),
* then runs the Yjs sync/awareness handshake.
*
* The avatar roster is derived from this room's own `owners` map and broadcast as
* `FILE_DOC_EVENTS.PRESENCE` — NOT the Redis-backed room-manager presence the workflow /
* table rooms use — because the file-doc room already owns an authoritative in-memory Y.Doc
* pinned to a single replica, so the session identity is right here with no extra store.
*/
export function setupWorkspaceFileDocHandlers(
socket: AuthenticatedSocket,
roomManager: IRoomManager
) {
const io = roomManager.io
// The file this socket currently intends to edit (set when a join starts). A leave targeting it
// — or an unscoped leave — advances the join generation to cancel an in-flight join, so a join
// awaiting authorization can't complete after the client left and register a ghost owner. A
// leave for a DIFFERENT file must NOT cancel it (a document switch), mirroring workspace-files.
let currentFileId: string | null = null

socket.on(FILE_DOC_EVENTS.JOIN, async ({ fileId, clientId }: JoinFileDocPayload) => {
try {
Expand All @@ -376,9 +421,11 @@ export function setupWorkspaceFileDocHandlers(
return
}

// Claim this JOIN's generation before the async authorize below.
// Claim this JOIN's generation before the async authorize below, and record the file the
// socket now intends to edit so a leave for it can cancel this join if it's still in-flight.
const generation = (joinGeneration.get(socket.id) ?? 0) + 1
joinGeneration.set(socket.id, generation)
currentFileId = fileId

const room = fileDocRoom(fileId)
const name = roomName(room)
Expand Down Expand Up @@ -408,9 +455,13 @@ export function setupWorkspaceFileDocHandlers(
return
}

// Abort a JOIN superseded during authorization: the socket disconnected, or
// a newer JOIN (a document switch) bumped the generation. Registering here
// would leak a dead socket's room or bind the socket to the wrong document.
// Server-authenticated identity for the presence roster (never trusts the client-set
// awareness). Resolved here so the generation guard below also covers this await.
const avatarUrl = await resolveAvatarUrl(socket, userId)

// Abort a JOIN superseded during authorization/identity resolution: the socket
// disconnected, or a newer JOIN (a document switch) bumped the generation. Registering
// here would leak a dead socket's room or bind the socket to the wrong document.
Comment thread
waleedlatif1 marked this conversation as resolved.
if (socket.disconnected || joinGeneration.get(socket.id) !== generation) return

// Switched documents on the same socket — leave the previous one first (a
Expand Down Expand Up @@ -449,11 +500,13 @@ export function setupWorkspaceFileDocHandlers(
awarenessProtocol.removeAwarenessStates(entry.awareness, [previous.clientId], null)
}

entry.owners.set(socket.id, { clientId, userId })
entry.owners.set(socket.id, { clientId, userId, userName, avatarUrl })
socketToRoomName.set(socket.id, name)
socket.join(name)

socket.emit(FILE_DOC_EVENTS.JOIN_SUCCESS, { fileId })
// Server-authenticated roster → everyone in the room, including this joiner.
broadcastFileDocPresence(io, name, entry)
Comment thread
waleedlatif1 marked this conversation as resolved.

// Begin the sync handshake: send the server's state (sync step 1). The
// client replies with its updates and requests the server's in return.
Expand Down Expand Up @@ -496,10 +549,17 @@ export function setupWorkspaceFileDocHandlers(

socket.on(FILE_DOC_EVENTS.LEAVE, (payload?: LeaveFileDocPayload) => {
try {
// Only affect a REGISTERED room; never touch the join generation here. A
// leave that raced ahead of an in-flight join (no room registered yet) is a
// no-op — bumping the generation would silently abort an unrelated join for
// a different file (a document switch), leaving the socket bound to nothing.
// Cancel an in-flight join whose file the client is now leaving (or an unscoped leave): a
// join still awaiting authorization would otherwise complete after the client left, register
// as an owner, and broadcast a ghost collaborator until disconnect. Guard on the current
// file intent so a stale/deferred leave for a DIFFERENT file can't abort the join the client
// has since switched to (bumping the generation blindly caused that regression in #5941).
if (!payload?.fileId || payload.fileId === currentFileId) {
joinGeneration.set(socket.id, (joinGeneration.get(socket.id) ?? 0) + 1)
currentFileId = null
}
// Tear down membership only for a REGISTERED room; a leave that raced ahead of an in-flight
// join (nothing registered yet) has already cancelled it above.
const name = socketToRoomName.get(socket.id)
if (!name) return
// Scope the leave to the named file when provided: a deferred leave from a
Expand Down
Loading