From 08c77838d111944c8f91d10315ed2b5d9ab9baef Mon Sep 17 00:00:00 2001 From: Christopher Toth Date: Thu, 30 Jul 2026 00:56:56 -0600 Subject: [PATCH] Fix file transfer reliability --- src/FileTransferManager.test.ts | 100 ++++++++++++++++++++++++++++++++ src/FileTransferManager.ts | 29 ++++----- src/FileTransferStore.test.ts | 50 ++++++++++++++++ src/FileTransferStore.ts | 32 +++++----- src/WebRTCService.test.ts | 24 ++++++++ src/WebRTCService.ts | 28 +++++++-- 6 files changed, 230 insertions(+), 33 deletions(-) create mode 100644 src/FileTransferStore.test.ts diff --git a/src/FileTransferManager.test.ts b/src/FileTransferManager.test.ts index 44187451..1d91e235 100644 --- a/src/FileTransferManager.test.ts +++ b/src/FileTransferManager.test.ts @@ -360,6 +360,106 @@ describe('FileTransferManager receive consent + integrity', () => { return (manager as unknown as { incomingTransfers: Map }).incomingTransfers; } + it('records accepted consent before waiting for the data channel to open', async () => { + vi.useFakeTimers(); + const { manager, gmcpFileTransfer, webRTCService } = createManager(); + webRTCService.isDataChannelOpen.mockReturnValue(false); + const offer = { + sender: 'Bob', + hash: 'accepted-hash', + filename: 'accepted.txt', + filesize: 12, + offerSdp: '{}', + }; + gmcpFileTransfer.emit('offer', offer); + const acceptPromise = manager.acceptTransfer(offer.sender, offer.hash); + + try { + await vi.advanceTimersByTimeAsync(0); + + expect(gmcpFileTransfer.sendAccept).toHaveBeenCalledTimes(1); + const acceptedOffers = ( + manager as unknown as { + acceptedOffers: Map< + string, + { filename: string; hash: string; sender: string; filesize: number } + >; + } + ).acceptedOffers; + expect(acceptedOffers.get(offer.hash)).toEqual({ + filename: offer.filename, + hash: offer.hash, + sender: offer.sender, + filesize: offer.filesize, + }); + } finally { + webRTCService.isDataChannelOpen.mockReturnValue(true); + await vi.advanceTimersByTimeAsync(100); + await acceptPromise.catch(() => undefined); + manager.cleanup(); + vi.useRealTimers(); + } + }); + + it('does not double-count a duplicated chunk and completes after all unique chunks arrive', async () => { + const { manager, gmcpFileTransfer, webRTCService } = createManager(); + webRTCService.isDataChannelOpen.mockReturnValue(true); + const capture = installDownloadCapture(); + try { + const bytes = new TextEncoder().encode('abcd'); + const hash = await sha256Hex(bytes); + await acceptOffer(manager, gmcpFileTransfer, { + sender: 'Bob', + hash, + filename: 'duplicate-safe.txt', + filesize: bytes.byteLength, + }); + const onComplete = vi.fn(); + const onError = vi.fn(); + manager.on('fileReceiveComplete', onComplete); + manager.on('fileTransferError', onError); + + const firstChunk = bytes.slice(0, 2); + const secondChunk = bytes.slice(2); + const firstFrame = frameChunk( + { + hash, + filename: 'duplicate-safe.txt', + chunkIndex: 0, + totalChunks: 2, + chunkSize: firstChunk.byteLength, + totalSize: bytes.byteLength, + }, + firstChunk, + ); + const secondFrame = frameChunk( + { + hash, + filename: 'duplicate-safe.txt', + chunkIndex: 1, + totalChunks: 2, + chunkSize: secondChunk.byteLength, + totalSize: bytes.byteLength, + }, + secondChunk, + ); + + webRTCService.emit('dataChannelMessage', firstFrame); + await flush(); + webRTCService.emit('dataChannelMessage', firstFrame); + await flush(); + webRTCService.emit('dataChannelMessage', secondFrame); + await flush(); + + expect(onError).not.toHaveBeenCalled(); + expect(onComplete).toHaveBeenCalledTimes(1); + expect(capture.downloads).toEqual(['duplicate-safe.txt']); + } finally { + capture.restore(); + manager.cleanup(); + } + }); + it('C1: drops chunks for a hash that was never accepted (no download, no completion)', async () => { const { manager, webRTCService } = createManager(); webRTCService.isDataChannelOpen.mockReturnValue(true); diff --git a/src/FileTransferManager.ts b/src/FileTransferManager.ts index e96c0327..cb6793b7 100644 --- a/src/FileTransferManager.ts +++ b/src/FileTransferManager.ts @@ -515,12 +515,15 @@ export default class FileTransferManager extends EventEmitter { } } - transfer.chunks[header.chunkIndex] = chunkData; - transfer.receivedSize += chunkData.byteLength; + const isNewChunk = transfer.chunks[header.chunkIndex] === undefined; + if (isNewChunk) { + transfer.chunks[header.chunkIndex] = chunkData; + transfer.receivedSize += chunkData.byteLength; + } transfer.lastActivityTimestamp = Date.now(); // Persist chunk to IndexedDB for resumable transfers - if (this.storeInitialized) { + if (isNewChunk && this.storeInitialized) { await this.store.saveChunk({ hash: header.hash, index: header.chunkIndex, @@ -830,6 +833,16 @@ export default class FileTransferManager extends EventEmitter { // Send accept only if we're still in a valid state if (this.pendingOffers.has(hash)) { + // Record consent before notifying the sender, so bytes arriving as soon as the + // data channel opens cannot be mistaken for an unsolicited transfer. + this.acceptedOffers.set(hash, { + filename: offer.filename, + hash: offer.hash, + sender: offer.sender, + filesize: offer.filesize, + }); + this.pendingOffers.delete(hash); + await this.gmcpFileTransfer.sendAccept({ sender, hash, @@ -840,16 +853,6 @@ export default class FileTransferManager extends EventEmitter { // Wait for the data channel to open await this.waitForDataChannel(hash); console.log('[FileTransferManager] Data channel ready for incoming transfer'); - - // Record consent BEFORE deleting the offer: this is the only durable marker the - // byte path can consult to know the user agreed to receive this exact file. - this.acceptedOffers.set(hash, { - filename: offer.filename, - hash: offer.hash, - sender: offer.sender, - filesize: offer.filesize, - }); - this.pendingOffers.delete(hash); } else { throw new Error('Transfer was cancelled during setup'); } diff --git a/src/FileTransferStore.test.ts b/src/FileTransferStore.test.ts new file mode 100644 index 00000000..44d17749 --- /dev/null +++ b/src/FileTransferStore.test.ts @@ -0,0 +1,50 @@ +import 'fake-indexeddb/auto'; + +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; + +import { FileTransferStore } from './FileTransferStore'; + +describe('FileTransferStore', () => { + let store: FileTransferStore; + + beforeEach(async () => { + store = new FileTransferStore(); + await store.initialize(); + await store.clearAll(); + }); + + afterEach(async () => { + await store.clearAll(); + await store.close(); + }); + + it('atomically records every concurrently saved chunk in the file metadata', async () => { + const totalChunks = 16; + await store.saveFileMetadata({ + hash: 'concurrent-file', + filename: 'concurrent.bin', + totalSize: totalChunks, + totalChunks, + receivedChunks: [], + direction: 'incoming', + sender: 'Bob', + lastActivityTimestamp: 0, + }); + + await Promise.all( + Array.from({ length: totalChunks }, (_, index) => + store.saveChunk({ + hash: 'concurrent-file', + index, + data: Uint8Array.of(index).buffer, + }), + ), + ); + + const metadata = await store.getFileMetadata('concurrent-file'); + expect(metadata?.receivedChunks.toSorted((a, b) => a - b)).toEqual( + Array.from({ length: totalChunks }, (_, index) => index), + ); + await expect(store.isTransferComplete('concurrent-file')).resolves.toBe(true); + }); +}); diff --git a/src/FileTransferStore.ts b/src/FileTransferStore.ts index d95b9f5b..5fae76e3 100644 --- a/src/FileTransferStore.ts +++ b/src/FileTransferStore.ts @@ -41,20 +41,24 @@ export class FileTransferStore { }); } - async saveChunk(chunk: FileChunk): Promise { - if (!this.db) await this.initialize(); - await this.db!.put('chunks', chunk); - - // Update metadata to track received chunks - const metadata = await this.getFileMetadata(chunk.hash); - if (metadata) { - if (!metadata.receivedChunks.includes(chunk.index)) { - metadata.receivedChunks.push(chunk.index); - metadata.lastActivityTimestamp = Date.now(); - await this.updateFileMetadata(metadata); - } - } - } + async saveChunk(chunk: FileChunk): Promise { + if (!this.db) await this.initialize(); + + const tx = this.db!.transaction(['chunks', 'metadata'], 'readwrite'); + await tx.objectStore('chunks').put(chunk); + + const metadata = await tx.objectStore('metadata').get(chunk.hash); + if (metadata) { + if (!metadata.receivedChunks.includes(chunk.index)) { + await tx.objectStore('metadata').put({ + ...metadata, + receivedChunks: [...metadata.receivedChunks, chunk.index], + lastActivityTimestamp: Date.now(), + }); + } + } + await tx.done; + } async getChunk(hash: string, index: number): Promise { if (!this.db) await this.initialize(); diff --git a/src/WebRTCService.test.ts b/src/WebRTCService.test.ts index 7ca0357a..962e97ff 100644 --- a/src/WebRTCService.test.ts +++ b/src/WebRTCService.test.ts @@ -295,6 +295,30 @@ describe('WebRTCService', () => { expect((webRTCService as any).dataChannel.send).toHaveBeenCalledWith(data); }); + it.each([ + ['close', 'closed'], + ['error', 'error'], + ])( + 'should reject a buffered send when the data channel emits %s', + async (eventName, expectedError) => { + await webRTCService.createPeerConnection(); + const dataChannel = (webRTCService as any).dataChannel; + dataChannel.readyState = 'open'; + dataChannel.bufferedAmount = 2000000; + const addEventListener = vi.spyOn(dataChannel, 'addEventListener'); + + const sendPromise = webRTCService.sendData(new ArrayBuffer(10)); + const eventHandler = addEventListener.mock.calls.find( + ([registeredEvent]) => registeredEvent === eventName, + )?.[1] as EventListener | undefined; + + expect(eventHandler).toBeDefined(); + eventHandler?.(new Event(eventName)); + await expect(sendPromise).rejects.toThrow(expectedError); + expect(dataChannel.send).not.toHaveBeenCalled(); + }, + ); + it('should handle errors in send operation', async () => { await webRTCService.createPeerConnection(); (webRTCService as any).dataChannel.readyState = 'open'; diff --git a/src/WebRTCService.ts b/src/WebRTCService.ts index 72341ba4..a08c138a 100644 --- a/src/WebRTCService.ts +++ b/src/WebRTCService.ts @@ -201,18 +201,34 @@ export class WebRTCService extends EventEmitter { try { // Implement flow control - wait if buffer is getting full const maxBufferSize = 1048576; // 1MB threshold - while (this.dataChannel.bufferedAmount > maxBufferSize) { - await new Promise((resolve) => { + const dataChannel = this.dataChannel; + while (dataChannel.bufferedAmount > maxBufferSize) { + await new Promise((resolve, reject) => { + const cleanup = () => { + dataChannel.removeEventListener('bufferedamountlow', onBufferedAmountLow); + dataChannel.removeEventListener('close', onClose); + dataChannel.removeEventListener('error', onError); + }; const onBufferedAmountLow = () => { - this.dataChannel?.removeEventListener('bufferedamountlow', onBufferedAmountLow); + cleanup(); resolve(); }; - this.dataChannel?.addEventListener('bufferedamountlow', onBufferedAmountLow); - this.dataChannel!.bufferedAmountLowThreshold = maxBufferSize / 2; + const onClose = () => { + cleanup(); + reject(new Error('Data channel closed while waiting to send data')); + }; + const onError = () => { + cleanup(); + reject(new Error('Data channel error while waiting to send data')); + }; + dataChannel.addEventListener('bufferedamountlow', onBufferedAmountLow); + dataChannel.addEventListener('close', onClose); + dataChannel.addEventListener('error', onError); + dataChannel.bufferedAmountLowThreshold = maxBufferSize / 2; }); } - this.dataChannel.send(data); + dataChannel.send(data); } catch (error) { console.error('[WebRTCService] Error sending data:', error); throw error;