@tldraw/sync-core 5.5.0-next.9665a55b2322 → 5.5.0-next.b6a9f7cd7be2

package/src/lib/ClientWebSocketAdapter.test.ts

@@ -12,6 +12,7 @@ vi.mock('@tldraw/utils', async (importOriginal) => {
 // NOTE: setupVitest.js replaces the global WebSocket with the 'ws' package's WebSocket,
 // matching the WebSocketServer the tests connect to.
 import { WebSocketServer, WebSocket as WsWebSocket } from 'ws'
+import { MAX_ASSEMBLED_MESSAGE_CHARS } from './chunk'
 import {
 	ACTIVE_MAX_DELAY,
 	ACTIVE_MIN_DELAY,
@@ -73,24 +74,38 @@ describe('ClientWebSocketAdapter', () => {
 		connectedServerSocket = socket
 	})
 
+	// Bound per test on an OS-assigned port. A fixed port raced against the previous test's
+	// server, which closes asynchronously, and intermittently failed with EADDRINUSE.
+	let port: number
+
 	let consoleWarnSpy: ReturnType<typeof vi.spyOn>
-	beforeEach(() => {
+	beforeEach(async () => {
 		consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {})
-		adapter = new ClientWebSocketAdapter(() => 'ws://localhost:2233')
-		wsServer = new WebSocketServer({ port: 2233 })
+		wsServer = new WebSocketServer({ port: 0 })
+		// Attach before the adapter exists: the server is already listening by the time the
+		// adapter is constructed, so a handler added afterwards can miss the first connection.
 		wsServer.on('connection', connectMock as any)
+		await new Promise<void>((resolve, reject) => {
+			wsServer.once('listening', resolve)
+			wsServer.once('error', reject)
+		})
+		port = (wsServer.address() as { port: number }).port
+		adapter = new ClientWebSocketAdapter(() => `ws://localhost:${port}`)
 	})
 
 	afterEach(() => {
 		consoleWarnSpy.mockRestore()
 		adapter.close()
+		// Not awaited: the next test binds its own OS-assigned port, so a server still closing
+		// here can't collide with it. Awaiting would deliver in-flight messages to the socket
+		// adapter.close() just orphaned.
 		wsServer.close()
 		connectMock.mockClear()
 	})
 
 	describe('connection and initial state (CW1, CW2)', () => {
 		it('[CW2] starts with connectionStatus offline while the internal status is initial', () => {
-			const newAdapter = new ClientWebSocketAdapter(() => 'ws://localhost:2233')
+			const newAdapter = new ClientWebSocketAdapter(() => `ws://localhost:${port}`)
 			try {
 				expect(newAdapter._connectionStatus.get()).toBe('initial')
 				expect(newAdapter.connectionStatus).toBe('offline')
@@ -113,7 +128,7 @@ describe('ClientWebSocketAdapter', () => {
 			let uriCallCount = 0
 			const dynamicAdapter = new ClientWebSocketAdapter(() => {
 				uriCallCount++
-				return `ws://localhost:2233?attempt=${uriCallCount}`
+				return `ws://localhost:${port}?attempt=${uriCallCount}`
 			})
 			try {
 				await waitFor(() => dynamicAdapter._ws?.readyState === WebSocket.OPEN)
@@ -143,7 +158,7 @@ describe('ClientWebSocketAdapter', () => {
 				// no socket can be created until the promise resolves
 				expect(asyncAdapter._ws).toBeNull()
 
-				resolveUri!('ws://localhost:2233')
+				resolveUri!(`ws://localhost:${port}`)
 
 				await waitFor(() => asyncAdapter._ws?.readyState === WebSocket.OPEN)
 				expect(asyncAdapter.connectionStatus).toBe('online')
@@ -301,7 +316,7 @@ describe('ClientWebSocketAdapter', () => {
 
 	describe('URI conversion (CW1)', () => {
 		it('[CW1] converts http URIs to ws and connects', async () => {
-			const httpAdapter = new ClientWebSocketAdapter(() => 'http://localhost:2233')
+			const httpAdapter = new ClientWebSocketAdapter(() => `http://localhost:${port}`)
 			try {
 				await waitFor(() => httpAdapter._ws?.readyState === WebSocket.OPEN)
 				expect(httpAdapter._ws!.url).toMatch(/^ws:\/\//)
@@ -312,7 +327,7 @@ describe('ClientWebSocketAdapter', () => {
 		})
 
 		it('[CW1] converts https URIs to wss', async () => {
-			const httpsAdapter = new ClientWebSocketAdapter(() => 'https://localhost:2233')
+			const httpsAdapter = new ClientWebSocketAdapter(() => `https://localhost:${port}`)
 			try {
 				await waitFor(() => httpsAdapter._ws !== null)
 				expect(httpsAdapter._ws!.url).toMatch(/^wss:\/\//)
@@ -386,6 +401,32 @@ describe('ClientWebSocketAdapter', () => {
 			)
 		})
 
+		it('[CW6] refuses to send a message the server would reject as too large', async () => {
+			const onMessage = vi.fn()
+			connectMock.mockImplementationOnce((ws: any) => {
+				ws.on('message', onMessage)
+			})
+			await waitFor(() => adapter._ws?.readyState === WebSocket.OPEN)
+
+			const onStatusChange = vi.fn()
+			adapter.onStatusChange(onStatusChange)
+
+			const message = {
+				...connectMessage(),
+				largeData: 'x'.repeat(MAX_ASSEMBLED_MESSAGE_CHARS),
+			} as any
+			adapter.sendMessage(message)
+
+			// nothing goes on the wire: sending would only earn a close and a reconnect that
+			// re-sends the same message
+			expect(onMessage).not.toHaveBeenCalled()
+			expect(adapter.connectionStatus).toBe('error')
+			expect(onStatusChange).toHaveBeenCalledWith({
+				status: 'error',
+				reason: TLSyncErrorCloseEventReason.MESSAGE_TOO_LARGE,
+			})
+		})
+
 		it('[CW6] silently drops the message when there is no socket', async () => {
 			await waitFor(() => adapter._ws?.readyState === WebSocket.OPEN)
 
@@ -529,7 +570,7 @@ describe('ClientWebSocketAdapter', () => {
 		})
 
 		it('[CW9][RM5] closing immediately does not start another getUri call', async () => {
-			const getUri = vi.fn(() => 'ws://localhost:2233')
+			const getUri = vi.fn(() => `ws://localhost:${port}`)
 			const testAdapter = new ClientWebSocketAdapter(getUri)
 			testAdapter.close()
 			const callsAtClose = getUri.mock.calls.length
@@ -544,7 +585,7 @@ describe('ClientWebSocketAdapter', () => {
 			let uriCallCount = 0
 			const testAdapter = new ClientWebSocketAdapter(() => {
 				uriCallCount++
-				return 'ws://localhost:2233'
+				return `ws://localhost:${port}`
 			})
 			const onStatusChange = vi.fn()
 			testAdapter.onStatusChange(onStatusChange)
@@ -606,17 +647,26 @@ describe('ReconnectManager', () => {
 		connectedServerSocket = socket
 	})
 
+	// See the note on the other suite's port: bound per test on an OS-assigned port.
+	let port: number
+
 	let consoleWarnSpy: ReturnType<typeof vi.spyOn>
-	beforeEach(() => {
+	beforeEach(async () => {
 		consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {})
-		adapter = new ClientWebSocketAdapter(() => 'ws://localhost:2234')
-		wsServer = new WebSocketServer({ port: 2234 })
+		wsServer = new WebSocketServer({ port: 0 })
 		wsServer.on('connection', connectMock as any)
+		await new Promise<void>((resolve, reject) => {
+			wsServer.once('listening', resolve)
+			wsServer.once('error', reject)
+		})
+		port = (wsServer.address() as { port: number }).port
+		adapter = new ClientWebSocketAdapter(() => `ws://localhost:${port}`)
 	})
 
 	afterEach(() => {
 		consoleWarnSpy.mockRestore()
 		adapter.close()
+		// Not awaited - see the note on the other suite's teardown.
 		wsServer.close()
 		connectMock.mockClear()
 	})
@@ -631,7 +681,7 @@ describe('ReconnectManager', () => {
 	})
 
 	it('[RM1] reconnection delays back off exponentially, bounded by the active delays', () => {
-		const testAdapter = new ClientWebSocketAdapter(() => 'ws://localhost:2234')
+		const testAdapter = new ClientWebSocketAdapter(() => `ws://localhost:${port}`)
 		try {
 			const manager = testAdapter._reconnectManager
 			// make every disconnected() call take the "attempt now" branch
@@ -656,7 +706,7 @@ describe('ReconnectManager', () => {
 
 	it('[RM1] uses the inactive delay bounds when the tab is hidden', () => {
 		const hiddenMock = vi.spyOn(document, 'hidden', 'get').mockReturnValue(true)
-		const testAdapter = new ClientWebSocketAdapter(() => 'ws://localhost:2234')
+		const testAdapter = new ClientWebSocketAdapter(() => `ws://localhost:${port}`)
 		try {
 			const manager = testAdapter._reconnectManager
 			;(manager as any).lastAttemptStart = Date.now() - 1_000_000
@@ -828,7 +878,7 @@ describe('ReconnectManager', () => {
 					if (failure === 'throws') throw error
 					return Promise.reject(error)
 				}
-				return 'ws://localhost:2234'
+				return `ws://localhost:${port}`
 			})
 			try {
 				await waitFor(() => testAdapter._ws?.readyState === WebSocket.OPEN)
@@ -843,7 +893,7 @@ describe('ReconnectManager', () => {
 	)
 
 	it('[RM5] close cancels timers and removes the reconnect event listeners', () => {
-		const testAdapter = new ClientWebSocketAdapter(() => 'ws://localhost:2234')
+		const testAdapter = new ClientWebSocketAdapter(() => `ws://localhost:${port}`)
 		const manager = testAdapter._reconnectManager
 
 		testAdapter.close()
@@ -862,6 +912,10 @@ describe('ReconnectManager', () => {
 })
 
 describe('URI failure boundaries', () => {
+	// These tests never open a real socket - the adapter is closed first, or the socket is a
+	// mock - so the URI only has to parse. No server, hence no port to bind.
+	const UNUSED_URI = 'ws://localhost:1'
+
 	it.each(['resolves', 'rejects'])('ignores getUri that %s after close', async (outcome) => {
 		const uri = promiseWithResolve<string>()
 		const getUri = vi.fn(() => uri)
@@ -872,7 +926,7 @@ describe('URI failure boundaries', () => {
 		try {
 			await Promise.resolve()
 			testAdapter.close()
-			if (outcome === 'resolves') uri.resolve('ws://localhost:2234')
+			if (outcome === 'resolves') uri.resolve(UNUSED_URI)
 			else uri.reject(new Error('token endpoint unavailable'))
 			await vi.advanceTimersByTimeAsync(INACTIVE_MAX_DELAY)
 
@@ -890,7 +944,7 @@ describe('URI failure boundaries', () => {
 	it.each([
 		{ uri: 'not a URL', openSocket: false, error: 'Invalid URL' },
 		{
-			uri: 'ws://localhost:2234',
+			uri: UNUSED_URI,
 			openSocket: true,
 			error: 'There should be no connection attempts while already connected',
 		},

package/src/lib/ClientWebSocketAdapter.ts

@@ -1,7 +1,7 @@
 import { atom, Atom } from '@tldraw/state'
 import { TLRecord } from '@tldraw/tlschema'
 import { assert, warnOnce } from '@tldraw/utils'
-import { chunk } from './chunk'
+import { chunk, MAX_ASSEMBLED_MESSAGE_CHARS } from './chunk'
 import { TLSocketClientSentEvent, TLSocketServerSentEvent } from './protocol'
 import {
 	TLPersistentClientSocket,
@@ -280,7 +280,20 @@ export class ClientWebSocketAdapter implements TLPersistentClientSocket<
 
 		if (!this._ws) return
 		if (this.connectionStatus === 'online') {
-			const chunks = chunk(JSON.stringify(msg))
+			const stringified = JSON.stringify(msg)
+			// The server rejects a session whose assembled message exceeds this, so sending would
+			// only earn a close and a reconnect that re-sends the same message. Fail here the way
+			// the server would, so onSyncError runs once instead of the client looping forever.
+			if (stringified.length > MAX_ASSEMBLED_MESSAGE_CHARS) {
+				this._handleDisconnect(
+					'closed',
+					TLSyncErrorCloseEventCode,
+					true,
+					TLSyncErrorCloseEventReason.MESSAGE_TOO_LARGE
+				)
+				return
+			}
+			const chunks = chunk(stringified)
 			for (const part of chunks) {
 				this._ws.send(part)
 			}

package/src/lib/InMemorySyncStorage.test.ts

@@ -51,6 +51,27 @@ describe('InMemorySyncStorage', () => {
 			expect(storage.getClock()).toBe(25)
 		})
 
+		it('[IM1] handles snapshots with far more records and tombstones than fit in a spread call', () => {
+			const documents = []
+			for (let i = 0; i < 150_000; i++) {
+				documents.push({
+					state: { ...contractRecords[0], id: `shape:${i}` } as TLRecord,
+					lastChangedClock: i,
+				})
+			}
+			const tombstones: Record<string, number> = {}
+			for (let i = 0; i < 150_000; i++) tombstones[`shape:gone-${i}`] = 150_000 + i
+			const storage = new InMemorySyncStorage<TLRecord>({
+				snapshot: makeContractSnapshot(contractRecords, {
+					documents,
+					tombstones,
+					documentClock: 0,
+					tombstoneHistoryStartsAtClock: 0,
+				}),
+			})
+			expect(storage.getClock()).toBe(299_999)
+		})
+
 		it('[IM2] clamps tombstoneHistoryStartsAtClock down to the document clock', () => {
 			const storage = new InMemorySyncStorage<TLRecord>({
 				snapshot: makeContractSnapshot(contractRecords, {

package/src/lib/InMemorySyncStorage.ts

@@ -153,27 +153,27 @@ export class InMemorySyncStorage<R extends UnknownRecord> implements TLSyncStora
 		onChange?(arg: TLSyncStorageOnChangeCallbackProps): unknown
 	} = {}) {
 		this.objectTypes = new Set(objectTypes ?? [])
-		const maxClockValue = Math.max(
-			0,
-			...Object.values(snapshot.tombstones ?? {}),
-			...Object.values(snapshot.documents.map((d) => d.lastChangedClock))
-		)
+		// a loop rather than `Math.max(0, ...clocks)`: spreading a large room's clocks as call
+		// arguments overflows the stack at roughly 100k+ records/tombstones
+		let maxClockValue = 0
 		// route snapshot entries into their partitions (a seed snapshot may carry object-lane
 		// records merged in, e.g. loaded from a separate persistence lane)
-		const toEntry = (
-			d: RoomSnapshot['documents'][number]
-		): [string, { state: R; lastChangedClock: number }] => [
-			d.state.id,
-			{ state: devFreeze(d.state) as R, lastChangedClock: d.lastChangedClock },
-		]
-		this.documents = new AtomMap(
-			'room documents',
-			snapshot.documents.filter((d) => !this.objectTypes.has(d.state.typeName)).map(toEntry)
-		)
-		this.objects = new AtomMap(
-			'room objects',
-			snapshot.documents.filter((d) => this.objectTypes.has(d.state.typeName)).map(toEntry)
-		)
+		const documentEntries: [string, { state: R; lastChangedClock: number }][] = []
+		const objectEntries: [string, { state: R; lastChangedClock: number }][] = []
+		for (const d of snapshot.documents) {
+			if (d.lastChangedClock > maxClockValue) maxClockValue = d.lastChangedClock
+			const entries = this.objectTypes.has(d.state.typeName) ? objectEntries : documentEntries
+			entries.push([
+				d.state.id,
+				{ state: devFreeze(d.state) as R, lastChangedClock: d.lastChangedClock },
+			])
+		}
+		const tombstoneEntries = objectMapEntries(snapshot.tombstones ?? {})
+		for (const [, clock] of tombstoneEntries) {
+			if (clock > maxClockValue) maxClockValue = clock
+		}
+		this.documents = new AtomMap('room documents', documentEntries)
+		this.objects = new AtomMap('room objects', objectEntries)
 		const documentClock = Math.max(maxClockValue, snapshot.documentClock ?? snapshot.clock ?? 0)
 
 		this.documentClock = atom('document clock', documentClock)
@@ -192,9 +192,7 @@ export class InMemorySyncStorage<R extends UnknownRecord> implements TLSyncStora
 			'room tombstones',
 			// If the tombstone history starts now (or we didn't have the
 			// tombstoneHistoryStartsAtClock) then there are no tombstones
-			tombstoneHistoryStartsAtClock === documentClock
-				? []
-				: objectMapEntries(snapshot.tombstones ?? {})
+			tombstoneHistoryStartsAtClock === documentClock ? [] : tombstoneEntries
 		)
 		if (onChange) {
 			this.onChange(onChange)

package/src/lib/TLSocketRoom.ts

@@ -1,7 +1,7 @@
 import type { SerializedSchema, StoreSchema, UnknownRecord } from '@tldraw/store'
 import { createTLSchema, TLInstancePresence, TLStoreSnapshot } from '@tldraw/tlschema'
 import { getOwnProperty, hasOwnProperty, isEqual, structuredClone } from '@tldraw/utils'
-import { JsonChunkAssembler } from './chunk'
+import { JsonChunkAssembler, MessageTooLargeError } from './chunk'
 import { DEFAULT_INITIAL_SNAPSHOT, InMemorySyncStorage } from './InMemorySyncStorage'
 import { TLObjectStoreAccess, TLSocketServerSentEvent } from './protocol'
 import { RoomSessionState } from './RoomSession'
@@ -477,6 +477,11 @@ export class TLSocketRoom<R extends UnknownRecord = UnknownRecord, SessionMeta =
 				})
 				this.room.pruneSessions()
 				this.scheduleDebouncedSnapshot(sessionId)
+			} else if (res.error instanceof MessageTooLargeError) {
+				this.log?.error?.('Error assembling message', res.error)
+				// Resetting the connection would have the client reconnect and re-send the same
+				// oversized message, so reject the session instead and let the client surface it.
+				this.room.rejectSession(sessionId, TLSyncErrorCloseEventReason.MESSAGE_TOO_LARGE)
 			} else {
 				this.log?.error?.('Error assembling message', res.error)
 				// close the socket to reset the connection
@@ -909,7 +914,8 @@ export class TLSocketRoom<R extends UnknownRecord = UnknownRecord, SessionMeta =
 	 * The client will attempt to reconnect automatically unless a fatal reason is provided.
 	 *
 	 * @param sessionId - Session identifier to remove
-	 * @param fatalReason - Optional fatal error reason that prevents reconnection
+	 * @param fatalReason - Optional fatal error reason that prevents reconnection. WebSocket close
+	 * reasons are capped at 123 UTF-8 bytes, so a longer reason is truncated and ends with `... (+N bytes)`.
 	 *
 	 * @example
 	 * ```ts

package/src/lib/TLSyncClient.ts

@@ -122,6 +122,8 @@ export const TLSyncErrorCloseEventReason = {
 	RATE_LIMITED: 'RATE_LIMITED',
 	/** Room has reached maximum capacity */
 	ROOM_FULL: 'ROOM_FULL',
+	/** A single message exceeded the size the server will assemble */
+	MESSAGE_TOO_LARGE: 'MESSAGE_TOO_LARGE',
 } as const
 
 /**

package/src/lib/TLSyncRoom.ts

@@ -88,6 +88,29 @@ export const DATA_MESSAGE_DEBOUNCE_INTERVAL = 1000 / 60
 
 const timeSince = (time: number) => Date.now() - time
 
+// WebSocket close frames cap the reason at 123 UTF-8 bytes; a longer one makes `close()` throw
+const MAX_CLOSE_REASON_BYTES = 123
+const closeReasonEncoder = new TextEncoder()
+
+const truncatedSuffix = (droppedBytes: number) => `... (+${droppedBytes} bytes)`
+
+function truncateCloseReason(reason: string) {
+	const totalBytes = closeReasonEncoder.encode(reason).length
+	if (totalBytes <= MAX_CLOSE_REASON_BYTES) return reason
+	// the dropped count can't have more digits than the total, so reserving room for it always fits
+	const budget = MAX_CLOSE_REASON_BYTES - truncatedSuffix(totalBytes).length
+	let out = ''
+	let outBytes = 0
+	// step by code point so a multi-byte character is never cut in half
+	for (const char of reason) {
+		const charBytes = closeReasonEncoder.encode(char).length
+		if (outBytes + charBytes > budget) break
+		out += char
+		outBytes += charBytes
+	}
+	return out + truncatedSuffix(totalBytes - outBytes)
+}
+
 /**
  * Snapshot of a room's complete state that can be persisted and restored.
  * Contains all documents, tombstones, and metadata needed to reconstruct the room.
@@ -559,7 +582,8 @@ export class TLSyncRoom<R extends UnknownRecord, SessionMeta> {
 
 		try {
 			if (fatalReason) {
-				session.socket.close(TLSyncErrorCloseEventCode, fatalReason)
+				// the session is already gone, so a throwing close() would leave its socket open
+				session.socket.close(TLSyncErrorCloseEventCode, truncateCloseReason(fatalReason))
 			} else {
 				session.socket.close()
 			}
@@ -941,7 +965,8 @@ export class TLSyncRoom<R extends UnknownRecord, SessionMeta> {
 	 * Sends appropriate error messages before closing the connection.
 	 *
 	 * @param sessionId - The session to reject
-	 * @param fatalReason - The reason for rejection (optional)
+	 * @param fatalReason - The reason for rejection (optional). WebSocket close reasons are capped
+	 * at 123 UTF-8 bytes, so a longer reason is truncated and ends with `... (+N bytes)`.
 	 * @example
 	 * ```ts
 	 * // Reject due to version mismatch

package/src/lib/chunk.test.ts

@@ -1,5 +1,11 @@
 import { assert } from 'tldraw'
-import { JsonChunkAssembler, chunk } from './chunk'
+import {
+	JsonChunkAssembler,
+	MAX_ASSEMBLED_MESSAGE_CHARS,
+	MAX_CHUNK_COUNT,
+	MessageTooLargeError,
+	chunk,
+} from './chunk'
 
 describe('chunk (CH1–CH3)', () => {
 	describe('size boundary (CH1)', () => {
@@ -403,4 +409,51 @@ describe('JsonChunkAssembler (CH4–CH8)', () => {
 			})
 		})
 	})
+
+	describe('assembly bounds (CH10)', () => {
+		it('[CH10] rejects a first chunk declaring more than MAX_CHUNK_COUNT chunks', () => {
+			const unchunker = new JsonChunkAssembler()
+
+			const result = unchunker.handleMessage(`${MAX_CHUNK_COUNT}_{"a":`)
+			expect(result).toMatchObject({
+				error: expect.objectContaining({ message: `Too many chunks: ${MAX_CHUNK_COUNT + 1}` }),
+			})
+			expect(unchunker.state).toBe('idle')
+		})
+
+		it('[CH10] accepts a declared count at the limit', () => {
+			const unchunker = new JsonChunkAssembler()
+
+			expect(unchunker.handleMessage(`${MAX_CHUNK_COUNT - 1}_{"a":`)).toBeNull()
+			expect(unchunker.state).not.toBe('idle')
+		})
+
+		it('[CH10] rejects an assembly once its accumulated bodies pass MAX_ASSEMBLED_MESSAGE_CHARS', () => {
+			const unchunker = new JsonChunkAssembler()
+			// A chunk stream that never sends its final chunk. Before CH10 this accumulated
+			// without limit for as long as the sender kept sending.
+			const body = 'x'.repeat(1024 * 1024)
+			const chunksNeeded = Math.ceil(MAX_ASSEMBLED_MESSAGE_CHARS / body.length) + 1
+
+			let result = null
+			for (let i = 0; i < chunksNeeded; i++) {
+				result = unchunker.handleMessage(`${chunksNeeded - i}_${body}`)
+				if (result) break
+			}
+
+			expect(result).toMatchObject({ error: expect.any(MessageTooLargeError) })
+			expect(unchunker.state).toBe('idle')
+		})
+
+		it('[CH10] leaves the assembler usable after rejecting an oversized assembly', () => {
+			const unchunker = new JsonChunkAssembler()
+			const body = 'x'.repeat(1024 * 1024)
+			const chunksNeeded = Math.ceil(MAX_ASSEMBLED_MESSAGE_CHARS / body.length) + 1
+			for (let i = 0; i < chunksNeeded; i++) {
+				if (unchunker.handleMessage(`${chunksNeeded - i}_${body}`)) break
+			}
+
+			expect(unchunker.handleMessage('{"test": true}')).toMatchObject({ data: { test: true } })
+		})
+	})
 })

package/src/lib/chunk.ts

@@ -8,6 +8,35 @@ const MAX_BYTES_PER_CHAR = 4
 // in the (admittedly impossible) worst case, the max size is 1/4 of a megabyte
 const MAX_SAFE_MESSAGE_SIZE = MAX_CLIENT_SENT_MESSAGE_SIZE_BYTES / MAX_BYTES_PER_CHAR
 
+// JsonChunkAssembler holds every chunk of a message until the last one arrives, and the sender
+// decides how many that is — the count is just a prefix on the first chunk. Nothing in the
+// protocol obliges a sender to ever send the last chunk, so without a ceiling a single connection
+// can make the receiver hold an arbitrarily large buffer for as long as the socket stays open.
+//
+// The cap counts UTF-16 code units, which is what a JS string costs in memory: 16M of them is the
+// 32MB of retained string data a room can afford to hold per session and stay inside a Durable
+// Object's 128MB instance limit with room for the document itself. It is far above anything
+// chunk() produces for a real document push. The character cap is the real bound; the chunk-count
+// cap only rejects an absurd declared count on arrival instead of waiting for the characters to
+// accumulate, and is loose enough not to constrain a sender that chunks more finely than the
+// default.
+export const MAX_ASSEMBLED_MESSAGE_CHARS = 16 * 1024 * 1024
+export const MAX_CHUNK_COUNT = 100_000
+
+/**
+ * Thrown when a message would exceed MAX_ASSEMBLED_MESSAGE_CHARS, on either side of the wire.
+ * Distinct from the other assembly errors because it is fatal: retrying sends the same
+ * oversized message, so the session is rejected rather than reset.
+ *
+ * @internal
+ */
+export class MessageTooLargeError extends Error {
+	constructor(chars: number) {
+		super(`Message too large: ${chars} characters, max ${MAX_ASSEMBLED_MESSAGE_CHARS}`)
+		this.name = 'MessageTooLargeError'
+	}
+}
+
 /**
  * Splits a string into smaller chunks suitable for transmission over WebSockets.
  * This function ensures messages don't exceed size limits imposed by platforms like Cloudflare Workers (1MB max).
@@ -100,6 +129,7 @@ export class JsonChunkAssembler {
 		| {
 				chunksReceived: string[]
 				totalChunks: number
+				charsReceived: number
 		  } = 'idle'
 
 	/**
@@ -110,7 +140,8 @@ export class JsonChunkAssembler {
 	 * @param msg - The message to process, either JSON or chunk format
 	 * @returns Result object with data/stringified on success, error object on failure, or null for incomplete chunks
 	 * 	- `\{ data: object, stringified: string \}` - Successfully parsed complete message
-	 * 	- `\{ error: Error \}` - Parse error or invalid chunk sequence
+	 * 	- `\{ error: Error \}` - Parse error, invalid chunk sequence, or a message exceeding
+	 * 	  MAX_ASSEMBLED_MESSAGE_CHARS / MAX_CHUNK_COUNT
 	 * 	- `null` - Chunk received but more chunks expected
 	 *
 	 * @example
@@ -144,17 +175,30 @@ export class JsonChunkAssembler {
 			const data = match[2]
 
 			if (this.state === 'idle') {
+				const totalChunks = numChunksRemaining + 1
+				if (totalChunks > MAX_CHUNK_COUNT) {
+					// Nothing has been buffered yet, so the state is already idle.
+					return { error: new Error(`Too many chunks: ${totalChunks}`) }
+				}
 				this.state = {
 					chunksReceived: [data],
-					totalChunks: numChunksRemaining + 1,
+					totalChunks,
+					charsReceived: data.length,
 				}
 			} else {
 				this.state.chunksReceived.push(data)
+				this.state.charsReceived += data.length
 				if (numChunksRemaining !== this.state.totalChunks - this.state.chunksReceived.length) {
 					this.state = 'idle'
 					return { error: new Error(`Chunks received in wrong order`) }
 				}
 			}
+			if (this.state.charsReceived > MAX_ASSEMBLED_MESSAGE_CHARS) {
+				// Drop what was buffered rather than keeping it until the socket closes.
+				const charsReceived = this.state.charsReceived
+				this.state = 'idle'
+				return { error: new MessageTooLargeError(charsReceived) }
+			}
 			if (this.state.chunksReceived.length === this.state.totalChunks) {
 				try {
 					const stringified = this.state.chunksReceived.join('')

package/src/test/TLSocketRoom.test.ts

@@ -21,6 +21,7 @@ import {
 	ZERO_INDEX_KEY,
 } from 'tldraw'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { MAX_ASSEMBLED_MESSAGE_CHARS } from '../lib/chunk'
 import { RecordOpType } from '../lib/diff'
 import { DEFAULT_INITIAL_SNAPSHOT, InMemorySyncStorage } from '../lib/InMemorySyncStorage'
 import { getTlsyncProtocolVersion } from '../lib/protocol'
@@ -392,6 +393,29 @@ describe('28. TLSocketRoom (SR)', () => {
 			expect(room.getSessions()[0].isConnected).toBe(false)
 		})
 
+		it('[SR5] rejects the session when an assembly exceeds the size cap', () => {
+			const log: TLSyncLog = { warn: vi.fn(), error: vi.fn() }
+			const room = new TLSocketRoom({ log })
+			const socket = createMockSocket()
+			connectSession(room, 'test-session', socket)
+
+			// A chunk stream that never sends its last chunk. Unlike the other assembly errors
+			// this one is fatal: a reconnecting client would re-send the same oversized message.
+			const body = 'x'.repeat(1024 * 1024)
+			const chunksNeeded = Math.ceil(MAX_ASSEMBLED_MESSAGE_CHARS / body.length) + 1
+			for (let i = 0; i < chunksNeeded; i++) {
+				room.handleSocketMessage('test-session', `${chunksNeeded - i}_${body}`)
+				if ((log.error as any).mock.calls.length) break
+			}
+
+			expect(log.error).toHaveBeenCalledWith('Error assembling message', expect.anything())
+			expect(socket.close).toHaveBeenCalledWith(
+				TLSyncErrorCloseEventCode,
+				TLSyncErrorCloseEventReason.MESSAGE_TOO_LARGE
+			)
+			expect(room.getSessions()).toHaveLength(0)
+		})
+
 		it('[SR5] rejects the session with UNKNOWN_ERROR when message handling throws', () => {
 			const log: TLSyncLog = { warn: vi.fn(), error: vi.fn() }
 			const room = new TLSocketRoom({ log })

package/src/test/TLSyncRoom.test.ts

@@ -2340,6 +2340,31 @@ describe('26. Session lifecycle (SES)', () => {
 		expect(room.sessions.size).toBe(0)
 	})
 
+	it('[SES4] a reason longer than 123 UTF-8 bytes is truncated on a code-point boundary with a dropped-bytes suffix', () => {
+		const { room } = makeRoom()
+
+		const exact = connectSession(room, 'exact')
+		room.rejectSession('exact', 'x'.repeat(123))
+		expect(exact.close).toHaveBeenCalledWith(TLSyncErrorCloseEventCode, 'x'.repeat(123))
+
+		// '... (+200 bytes)' reserves 16 bytes, leaving 107 for the reason
+		const ascii = connectSession(room, 'ascii')
+		room.rejectSession('ascii', 'x'.repeat(200))
+		expect(ascii.close).toHaveBeenCalledWith(
+			TLSyncErrorCloseEventCode,
+			'x'.repeat(107) + '... (+93 bytes)'
+		)
+
+		// 161 bytes in, 107 to fill: 1 + 26 × 4 = 105, and a 27th 4-byte character is dropped whole
+		const emoji = connectSession(room, 'emoji')
+		room.rejectSession('emoji', 'x' + '\u{1F600}'.repeat(40))
+		expect(emoji.close).toHaveBeenCalledWith(
+			TLSyncErrorCloseEventCode,
+			'x' + '\u{1F600}'.repeat(26) + '... (+56 bytes)'
+		)
+		expect(room.sessions.size).toBe(0)
+	})
+
 	it('[HS2][SES5] sets supportsStringAppend to false for protocol version 7', () => {
 		const { room } = makeRoom()
 		connectSession(room, 'v7-session', { protocolVersion: 7 })