packages/docs-core/lib/Services/Websockets/WebsocketService.spec.ts

   1 import { stringToUtf8Array } from '@proton/crypto/lib/utils'
   2 import type { DocumentKeys, NodeMeta } from '@proton/drive-store'
   3 import type { ServerMessageWithDocumentUpdates, ServerMessageWithEvents } from '@proton/docs-proto'
   4 import { DecryptedValue, EventTypeEnum } from '@proton/docs-proto'
   5 import type { LoggerInterface } from '@proton/utils/logs'
   6 import type { GetRealtimeUrlAndToken } from '../../UseCase/CreateRealtimeValetToken'
   7 import type { DecryptMessage } from '../../UseCase/DecryptMessage'
   8 import type { EncryptMessage } from '../../UseCase/EncryptMessage'
   9 import { WebsocketService } from './WebsocketService'
  10 import type { InternalEventBusInterface, WebsocketConnectionInterface } from '@proton/docs-shared'
  11 import { BroadcastSource } from '@proton/docs-shared'
  12 import { Result } from '../../Domain/Result/Result'
  13 import type { EncryptionMetadata } from '../../Types/EncryptionMetadata'
  14 import type { DocumentConnectionRecord } from './DocumentConnectionRecord'
  15 import { WebsocketConnectionEvent } from '../../Realtime/WebsocketEvent/WebsocketConnectionEvent'
  16 import { type UpdateDebouncer } from './Debouncer/UpdateDebouncer'
  17 import { DocumentDebounceMode } from './Debouncer/DocumentDebounceMode'
  18 import type { PrivateKeyReference, SessionKey } from '@proton/crypto'
  19 import type { MetricService } from '../Metrics/MetricService'
  20
  21 const mockOnReadyContentPayload = new TextEncoder().encode(
  22   JSON.stringify({ connectionId: '12345678', clientUpgradeRecommended: true, clientUpgradeRequired: true }),
  23 )
  24
  25 describe('WebsocketService', () => {
  26   let service: WebsocketService
  27   let eventBus: InternalEventBusInterface
  28   let encryptMessage: EncryptMessage
  29   let debouncer: UpdateDebouncer
  30   let connection: WebsocketConnectionInterface
  31   let record: DocumentConnectionRecord
  32   let logger: LoggerInterface
  33   let document: NodeMeta
  34   let keys: DocumentKeys
  35   let metricService: MetricService
  36
  37   const createService = async (mode: DocumentDebounceMode) => {
  38     if (service) {
  39       service.destroy()
  40     }
  41
  42     if (debouncer) {
  43       debouncer.destroy()
  44     }
  45
  46     document = { linkId: 'link-id-123', volumeId: 'volume-id-456' } as NodeMeta
  47
  48     keys = {
  49       documentContentKey: 'key-123' as unknown as SessionKey,
  50       userAddressPrivateKey: 'private-key-123' as unknown as PrivateKeyReference,
  51       userOwnAddress: 'foo',
  52     }
  53
  54     eventBus = {
  55       publish: jest.fn(),
  56     } as unknown as jest.Mocked<InternalEventBusInterface>
  57
  58     encryptMessage = {
  59       execute: jest.fn().mockReturnValue(Result.ok(stringToUtf8Array('123'))),
  60     } as unknown as jest.Mocked<EncryptMessage>
  61
  62     logger = {
  63       info: jest.fn(),
  64       debug: jest.fn(),
  65       error: jest.fn(),
  66     } as unknown as jest.Mocked<LoggerInterface>
  67
  68     service = new WebsocketService(
  69       {} as jest.Mocked<GetRealtimeUrlAndToken>,
  70       encryptMessage,
  71       {
  72         execute: jest.fn().mockReturnValue(Result.ok(stringToUtf8Array('123'))),
  73       } as unknown as jest.Mocked<DecryptMessage>,
  74       logger,
  75       eventBus,
  76       metricService,
  77       '0.0.0.0',
  78     )
  79
  80     service.createConnection(document, keys, {
  81       commitId: () => undefined,
  82     })
  83
  84     record = service.getConnectionRecord('link-id-123')!
  85
  86     debouncer = record.debouncer
  87     debouncer.setMode(mode)
  88     debouncer.markAsReadyToFlush()
  89
  90     connection = record.connection
  91     connection.broadcastMessage = jest.fn()
  92     connection.canBroadcastMessages = jest.fn().mockReturnValue(true)
  93   }
  94
  95   beforeEach(async () => {
  96     await createService(DocumentDebounceMode.Realtime)
  97   })
  98
  99   afterEach(() => {
 100     jest.resetAllMocks()
 101
 102     service.destroy()
 103
 104     debouncer.destroy()
 105   })
 106
 107   describe('createConnection', () => {
 108     it('should post ConnectionEstablishedButNotYetReady when connection is opened', async () => {
 109       connection.callbacks.onOpen()
 110
 111       expect(eventBus.publish).toHaveBeenCalledWith({
 112         type: WebsocketConnectionEvent.ConnectionEstablishedButNotYetReady,
 113         payload: {
 114           document: record.document,
 115         },
 116       })
 117     })
 118   })
 119
 120   describe('sendDocumentUpdateMessage', () => {
 121     it('should add to buffer', async () => {
 122       debouncer.addUpdates = jest.fn()
 123
 124       await service.sendDocumentUpdateMessage(document, new Uint8Array())
 125
 126       expect(debouncer.addUpdates).toHaveBeenCalled()
 127     })
 128   })
 129
 130   describe('handleDocumentUpdateBufferFlush', () => {
 131     it('should encrypt updates', async () => {
 132       const encryptMock = (service.encryptMessage = jest.fn())
 133
 134       await service.handleDocumentUpdateDebouncerFlush(document, new Uint8Array())
 135
 136       expect(encryptMock).toHaveBeenCalled()
 137     })
 138
 139     it('should broadcast message', async () => {
 140       await service.handleDocumentUpdateDebouncerFlush(document, new Uint8Array())
 141
 142       expect(connection.broadcastMessage).toHaveBeenCalled()
 143     })
 144
 145     it('should add message to ack ledger', async () => {
 146       service.ledger.messagePosted = jest.fn()
 147
 148       await service.handleDocumentUpdateDebouncerFlush(document, new Uint8Array())
 149
 150       expect(service.ledger.messagePosted).toHaveBeenCalled()
 151     })
 152   })
 153
 154   describe('onDocumentConnectionOpened', () => {
 155     it('should retry failed messages', async () => {
 156       service.retryFailedDocumentUpdatesForDoc = jest.fn()
 157
 158       service.onDocumentConnectionReadyToBroadcast(record, mockOnReadyContentPayload)
 159
 160       expect(service.retryFailedDocumentUpdatesForDoc).toHaveBeenCalled()
 161     })
 162   })
 163
 164   describe('onDocumentConnectionReadyToBroadcast', () => {
 165     it('should mark connection as ready to broadcast', async () => {
 166       connection.markAsReadyToAcceptMessages = jest.fn()
 167
 168       service.onDocumentConnectionReadyToBroadcast(record, mockOnReadyContentPayload)
 169
 170       expect(connection.markAsReadyToAcceptMessages).toHaveBeenCalled()
 171     })
 172
 173     it('should mark debouncer as ready to flush', () => {
 174       debouncer.markAsReadyToFlush = jest.fn()
 175
 176       service.onDocumentConnectionReadyToBroadcast(record, mockOnReadyContentPayload)
 177
 178       expect(debouncer.markAsReadyToFlush).toHaveBeenCalled()
 179     })
 180
 181     it('should retry failed document updates', () => {
 182       service.retryFailedDocumentUpdatesForDoc = jest.fn()
 183
 184       service.onDocumentConnectionReadyToBroadcast(record, mockOnReadyContentPayload)
 185
 186       expect(service.retryFailedDocumentUpdatesForDoc).toHaveBeenCalled()
 187     })
 188
 189     it('should pass readiness information to eventBus', () => {
 190       service.onDocumentConnectionReadyToBroadcast(record, mockOnReadyContentPayload)
 191       expect(eventBus.publish).toHaveBeenCalledWith({
 192         type: WebsocketConnectionEvent.ConnectedAndReady,
 193         payload: {
 194           document: record.document,
 195           readinessInformation: {
 196             connectionId: '12345678',
 197             clientUpgradeRecommended: true,
 198             clientUpgradeRequired: true,
 199           },
 200         },
 201       })
 202     })
 203
 204     it('should log error and call eventBus if content is not parsable', () => {
 205       service.onDocumentConnectionReadyToBroadcast(record, new TextEncoder().encode('not parsable'))
 206       expect(logger.error).toHaveBeenCalledWith('Unable to parse content from ConnectionReady message')
 207       expect(eventBus.publish).toHaveBeenCalledWith({
 208         type: WebsocketConnectionEvent.ConnectedAndReady,
 209         payload: {
 210           document: record.document,
 211           readinessInformation: undefined,
 212         },
 213       })
 214     })
 215   })
 216
 217   describe('retryAllFailedDocumentUpdates', () => {
 218     it('should get ledger unacknowledged updates', async () => {
 219       service.ledger.getUnacknowledgedUpdates = jest.fn().mockReturnValue([])
 220
 221       service.retryFailedDocumentUpdatesForDoc(document)
 222
 223       expect(service.ledger.getUnacknowledgedUpdates).toHaveBeenCalled()
 224     })
 225   })
 226
 227   describe('handleWindowUnload', () => {
 228     it('should not prevent leaving if no unsaved changes', async () => {
 229       const event = { preventDefault: jest.fn() } as unknown as BeforeUnloadEvent
 230
 231       service.handleWindowUnload(event)
 232
 233       expect(event.preventDefault).not.toHaveBeenCalled()
 234     })
 235
 236     it('should prevent leaving if unsaved changes', async () => {
 237       debouncer.addUpdates([new DecryptedValue(new Uint8Array())])
 238
 239       const event = { preventDefault: jest.fn() } as unknown as BeforeUnloadEvent
 240
 241       service.handleWindowUnload(event)
 242
 243       expect(event.preventDefault).toHaveBeenCalled()
 244     })
 245
 246     it('should immediately flush a buffer that has pending changes', async () => {
 247       debouncer.flush = jest.fn()
 248
 249       debouncer.addUpdates([new DecryptedValue(new Uint8Array())])
 250
 251       const event = { preventDefault: jest.fn() } as unknown as BeforeUnloadEvent
 252
 253       service.handleWindowUnload(event)
 254
 255       expect(debouncer.flush).toHaveBeenCalled()
 256     })
 257
 258     it('should prevent leaving if unacked changes', async () => {
 259       const event = { preventDefault: jest.fn() } as unknown as BeforeUnloadEvent
 260
 261       service.ledger.hasConcerningMessages = jest.fn().mockReturnValue(true)
 262
 263       service.handleWindowUnload(event)
 264
 265       expect(event.preventDefault).toHaveBeenCalled()
 266     })
 267   })
 268
 269   describe('flushPendingUpdates', () => {
 270     it('should immediately flush a buffer that has pending changes', async () => {
 271       debouncer.flush = jest.fn()
 272
 273       debouncer.addUpdates([new DecryptedValue(new Uint8Array())])
 274
 275       service.flushPendingUpdates()
 276
 277       expect(debouncer.flush).toHaveBeenCalled()
 278     })
 279   })
 280
 281   describe('sendEventMessage', () => {
 282     it('should encrypt event message', async () => {
 283       const encryptMock = (service.encryptMessage = jest.fn().mockReturnValue(stringToUtf8Array('123')))
 284
 285       await service.sendEventMessage(
 286         document,
 287         stringToUtf8Array('123'),
 288         EventTypeEnum.ClientHasSentACommentMessage,
 289         BroadcastSource.AwarenessUpdateHandler,
 290       )
 291
 292       expect(encryptMock).toHaveBeenCalled()
 293     })
 294
 295     it('should ignore sending ClientIsBroadcastingItsPresenceState event if not in realtime mode', async () => {
 296       debouncer.getMode = jest.fn().mockReturnValue(DocumentDebounceMode.SinglePlayer)
 297
 298       await service.sendEventMessage(
 299         document,
 300         stringToUtf8Array('123'),
 301         EventTypeEnum.ClientIsBroadcastingItsPresenceState,
 302         BroadcastSource.AwarenessUpdateHandler,
 303       )
 304
 305       expect(connection.broadcastMessage).not.toHaveBeenCalled()
 306     })
 307
 308     it('should ignore sending ClientHasSentACommentMessage event if not in realtime mode', async () => {
 309       debouncer.getMode = jest.fn().mockReturnValue(DocumentDebounceMode.SinglePlayer)
 310
 311       await service.sendEventMessage(
 312         document,
 313         stringToUtf8Array('123'),
 314         EventTypeEnum.ClientHasSentACommentMessage,
 315         BroadcastSource.AwarenessUpdateHandler,
 316       )
 317
 318       expect(connection.broadcastMessage).not.toHaveBeenCalled()
 319     })
 320
 321     it('should send ClientIsBroadcastingItsPresenceState event if in realtime mode', async () => {
 322       Object.defineProperty(debouncer, 'isBufferEnabled', { value: false })
 323
 324       await service.sendEventMessage(
 325         document,
 326         stringToUtf8Array('123'),
 327         EventTypeEnum.ClientIsBroadcastingItsPresenceState,
 328         BroadcastSource.AwarenessUpdateHandler,
 329       )
 330
 331       expect(connection.broadcastMessage).toHaveBeenCalled()
 332     })
 333
 334     it('should not broadcast if connection cannot send messages', async () => {
 335       connection.canBroadcastMessages = jest.fn().mockReturnValue(false)
 336
 337       await service.sendEventMessage(
 338         document,
 339         stringToUtf8Array('123'),
 340         EventTypeEnum.ClientIsBroadcastingItsPresenceState,
 341         BroadcastSource.AwarenessUpdateHandler,
 342       )
 343
 344       expect(connection.broadcastMessage).not.toHaveBeenCalled()
 345     })
 346   })
 347
 348   describe('handleIncomingDocumentUpdatesMessage', () => {
 349     it('should put us into realtime mode if message is not ours', async () => {
 350       const switchToRealtimeMode = (service.switchToRealtimeMode = jest.fn())
 351
 352       await service.handleIncomingDocumentUpdatesMessage(record, {
 353         updates: {
 354           documentUpdates: [
 355             {
 356               authorAddress: 'bar',
 357             },
 358           ],
 359         },
 360       } as unknown as ServerMessageWithDocumentUpdates)
 361
 362       expect(switchToRealtimeMode).toHaveBeenCalled()
 363     })
 364
 365     it('should not put us into realtime mode if message is ours', async () => {
 366       const switchToRealtimeMode = (service.switchToRealtimeMode = jest.fn())
 367
 368       await service.handleIncomingDocumentUpdatesMessage(record, {
 369         updates: {
 370           documentUpdates: [
 371             {
 372               authorAddress: 'foo',
 373             },
 374           ],
 375         },
 376       } as unknown as ServerMessageWithDocumentUpdates)
 377
 378       expect(switchToRealtimeMode).not.toHaveBeenCalled()
 379     })
 380   })
 381
 382   describe('handleIncomingEventsMessage', () => {
 383     it('should switch to realtime mode if event includes ClientIsRequestingOtherClientsToBroadcastTheirState', async () => {
 384       const switchToRealtimeMode = (service.switchToRealtimeMode = jest.fn())
 385
 386       const events = {
 387         events: [{ type: EventTypeEnum.ClientIsRequestingOtherClientsToBroadcastTheirState }],
 388       } as unknown as ServerMessageWithEvents
 389
 390       await service.handleIncomingEventsMessage(record, events)
 391
 392       expect(switchToRealtimeMode).toHaveBeenCalled()
 393     })
 394
 395     it('should switch to realtime mode if event includes ClientIsBroadcastingItsPresenceState', async () => {
 396       const switchToRealtimeMode = (service.switchToRealtimeMode = jest.fn())
 397
 398       const events = {
 399         events: [{ type: EventTypeEnum.ClientIsBroadcastingItsPresenceState }],
 400       } as unknown as ServerMessageWithEvents
 401
 402       await service.handleIncomingEventsMessage(record, events)
 403
 404       expect(switchToRealtimeMode).toHaveBeenCalled()
 405     })
 406
 407     it('should not switch to realtime mode for all other event types', async () => {
 408       const switchToRealtimeMode = (service.switchToRealtimeMode = jest.fn())
 409
 410       const events = {
 411         events: [
 412           { type: EventTypeEnum.ServerIsRequestingClientToBroadcastItsState },
 413           { type: EventTypeEnum.ServerHasMoreOrLessGivenTheClientEverythingItHas },
 414           { type: EventTypeEnum.ServerIsPlacingEmptyActivityIndicatorInStreamToIndicateTheStreamIsStillActive },
 415           { type: EventTypeEnum.ClientIsDebugRequestingServerToPerformCommit },
 416           { type: EventTypeEnum.ClientHasSentACommentMessage },
 417           { type: EventTypeEnum.ServerIsInformingClientThatTheDocumentCommitHasBeenUpdated },
 418         ],
 419       } as unknown as ServerMessageWithEvents
 420
 421       await service.handleIncomingEventsMessage(record, events)
 422
 423       expect(switchToRealtimeMode).not.toHaveBeenCalled()
 424     })
 425
 426     it('should markAsReadyToAcceptMessages on ServerIsReadyToAcceptClientMessages', async () => {
 427       const events = {
 428         events: [{ type: EventTypeEnum.ServerIsReadyToAcceptClientMessages }],
 429       } as unknown as ServerMessageWithEvents
 430
 431       connection.markAsReadyToAcceptMessages = jest.fn()
 432
 433       await service.handleIncomingEventsMessage(record, events)
 434
 435       expect(connection.markAsReadyToAcceptMessages).toHaveBeenCalled()
 436     })
 437   })
 438
 439   describe('handleLedgerStatusChangeCallback', () => {
 440     it('should post AckStatusChange event', () => {
 441       service.handleLedgerStatusChangeCallback()
 442
 443       expect(eventBus.publish).toHaveBeenCalledWith({
 444         type: WebsocketConnectionEvent.AckStatusChange,
 445         payload: expect.anything(),
 446       })
 447     })
 448   })
 449
 450   describe('encryptMessage', () => {
 451     it('should publish encryption error event if failed to encrypt', async () => {
 452       encryptMessage.execute = jest.fn().mockReturnValue(Result.fail('error'))
 453
 454       const spy = (eventBus.publish = jest.fn())
 455
 456       try {
 457         await service.encryptMessage(
 458           stringToUtf8Array('123'),
 459           {} as EncryptionMetadata,
 460           document,
 461           keys,
 462           BroadcastSource.AwarenessUpdateHandler,
 463         )
 464       } catch (error) {}
 465
 466       expect(spy).toHaveBeenCalledWith({
 467         type: WebsocketConnectionEvent.EncryptionError,
 468         payload: {
 469           document,
 470           error: expect.any(String),
 471         },
 472       })
 473     })
 474   })
 475 })