feat(snjs): add retrieving web sockets token for establishing connection (#1602)

* feat(snjs): add retrieving web sockets token for establishing connection

* fix(snjs): imports

* fix(snjs): use only web sockets connection token

* fix(snjs): linter issue

packages/snjs/lib/Application/Application.ts

@@ -10,6 +10,10 @@ import {
   UserRegistrationResponseBody,
   UserServer,
   UserServerInterface,
+  WebSocketApiService,
+  WebSocketApiServiceInterface,
+  WebSocketServer,
+  WebSocketServerInterface,
 } from '@standardnotes/api'
 import * as Common from '@standardnotes/common'
 import * as ExternalServices from '@standardnotes/services'
@@ -108,6 +112,8 @@ export class SNApplication
   private declare subscriptionApiService: SubscriptionApiServiceInterface
   private declare subscriptionServer: SubscriptionServerInterface
   private declare subscriptionManager: SubscriptionClientInterface
+  private declare webSocketApiService: WebSocketApiServiceInterface
+  private declare webSocketServer: WebSocketServerInterface
   private sessionManager!: InternalServices.SNSessionManager
   private syncService!: InternalServices.SNSyncService
   private challengeService!: InternalServices.ChallengeService
@@ -329,7 +335,7 @@ export class SNApplication
 
     this.apiService.loadHost()
     this.webSocketsService.loadWebSocketUrl()
-    this.sessionManager.initializeFromDisk()
+    await this.sessionManager.initializeFromDisk()
 
     this.settingsService.initializeFromDisk()
 
@@ -1094,6 +1100,8 @@ export class SNApplication
     this.createUserApiService()
     this.createSubscriptionServer()
     this.createSubscriptionApiService()
+    this.createWebSocketServer()
+    this.createWebSocketApiService()
     this.createSubscriptionManager()
     this.createWebSocketsService()
     this.createSessionManager()
@@ -1135,6 +1143,8 @@ export class SNApplication
     ;(this.userServer as unknown) = undefined
     ;(this.subscriptionApiService as unknown) = undefined
     ;(this.subscriptionServer as unknown) = undefined
+    ;(this.webSocketApiService as unknown) = undefined
+    ;(this.webSocketServer as unknown) = undefined
     ;(this.subscriptionManager as unknown) = undefined
     ;(this.sessionManager as unknown) = undefined
     ;(this.syncService as unknown) = undefined
@@ -1252,6 +1262,7 @@ export class SNApplication
     this.webSocketsService = new InternalServices.SNWebSocketsService(
       this.diskStorageService,
       this.options.webSocketUrl,
+      this.webSocketApiService,
       this.internalEventBus,
     )
     this.services.push(this.webSocketsService)
@@ -1337,6 +1348,14 @@ export class SNApplication
     this.subscriptionServer = new SubscriptionServer(this.httpService)
   }
 
+  private createWebSocketApiService() {
+    this.webSocketApiService = new WebSocketApiService(this.webSocketServer)
+  }
+
+  private createWebSocketServer() {
+    this.webSocketServer = new WebSocketServer(this.httpService)
+  }
+
   private createSubscriptionManager() {
     this.subscriptionManager = new SubscriptionManager(this.subscriptionApiService, this.internalEventBus)
   }

packages/snjs/lib/Services/Api/WebsocketsService.spec.ts

@@ -1,4 +1,6 @@
 import { InternalEventBusInterface } from '@standardnotes/services'
+import { WebSocketApiServiceInterface } from '@standardnotes/api'
+
 import { StorageKey, DiskStorageService } from '@Lib/index'
 import { SNWebSocketsService } from './WebsocketsService'
 
@@ -6,10 +8,11 @@ describe('webSocketsService', () => {
   const webSocketUrl = ''
 
   let storageService: DiskStorageService
+  let webSocketApiService: WebSocketApiServiceInterface
   let internalEventBus: InternalEventBusInterface
 
   const createService = () => {
-    return new SNWebSocketsService(storageService, webSocketUrl, internalEventBus)
+    return new SNWebSocketsService(storageService, webSocketUrl, webSocketApiService, internalEventBus)
   }
 
   beforeEach(() => {
@@ -18,6 +21,9 @@ describe('webSocketsService', () => {
 
     internalEventBus = {} as jest.Mocked<InternalEventBusInterface>
     internalEventBus.publish = jest.fn()
+
+    webSocketApiService = {} as jest.Mocked<WebSocketApiServiceInterface>
+    webSocketApiService.createConnectionToken = jest.fn().mockReturnValue({ token: 'foobar' })
   })
 
   describe('setWebSocketUrl()', () => {

packages/snjs/lib/Services/Api/WebsocketsService.ts

@@ -1,6 +1,8 @@
 import { UserRolesChangedEvent } from '@standardnotes/domain-events'
-import { DiskStorageService } from '../Storage/DiskStorageService'
 import { AbstractService, InternalEventBusInterface, StorageKey } from '@standardnotes/services'
+import { WebSocketApiServiceInterface } from '@standardnotes/api'
+
+import { DiskStorageService } from '../Storage/DiskStorageService'
 
 export enum WebSocketsServiceEvent {
   UserRoleMessageReceived = 'WebSocketMessageReceived',
@@ -12,6 +14,7 @@ export class SNWebSocketsService extends AbstractService<WebSocketsServiceEvent,
   constructor(
     private storageService: DiskStorageService,
     private webSocketUrl: string | undefined,
+    private webSocketApiService: WebSocketApiServiceInterface,
     protected override internalEventBus: InternalEventBusInterface,
   ) {
     super(internalEventBus)
@@ -34,10 +37,15 @@ export class SNWebSocketsService extends AbstractService<WebSocketsServiceEvent,
       )._websocket_url
   }
 
-  public startWebSocketConnection(authToken: string): void {
+  async startWebSocketConnection(): Promise<void> {
+    const webSocketConectionToken = await this.createWebSocketConnectionToken()
+    if (webSocketConectionToken === undefined) {
+      return
+    }
+
     if (this.webSocketUrl) {
       try {
-        this.webSocket = new WebSocket(`${this.webSocketUrl}?authToken=Bearer+${authToken}`)
+        this.webSocket = new WebSocket(`${this.webSocketUrl}?authToken=${webSocketConectionToken}`)
         this.webSocket.onmessage = this.onWebSocketMessage.bind(this)
         this.webSocket.onclose = this.onWebSocketClose.bind(this)
       } catch (e) {
@@ -59,9 +67,22 @@ export class SNWebSocketsService extends AbstractService<WebSocketsServiceEvent,
     this.webSocket = undefined
   }
 
+  private async createWebSocketConnectionToken(): Promise<string | undefined> {
+    try {
+      const response = await this.webSocketApiService.createConnectionToken()
+
+      return response.data.token
+    } catch (error) {
+      console.error((error as Error).message)
+
+      return undefined
+    }
+  }
+
   override deinit(): void {
     super.deinit()
     ;(this.storageService as unknown) = undefined
+    ;(this.webSocketApiService as unknown) = undefined
     this.closeWebSocketConnection()
   }
 }

packages/snjs/lib/Services/Session/SessionManager.ts

@@ -107,7 +107,7 @@ export class SNSessionManager extends AbstractService<SessionEvent> implements S
     this.apiService.setUser(user)
   }
 
-  public initializeFromDisk() {
+  async initializeFromDisk() {
     this.setUser(this.diskStorageService.getValue(StorageKey.User))
 
     if (!this.user) {
@@ -121,7 +121,7 @@ export class SNSessionManager extends AbstractService<SessionEvent> implements S
     if (rawSession) {
       const session = SessionFromRawStorageValue(rawSession)
       this.setSession(session, false)
-      this.webSocketsService.startWebSocketConnection(session.authorizationValue)
+      await this.webSocketsService.startWebSocketConnection()
     }
   }
 
@@ -623,9 +623,9 @@ export class SNSessionManager extends AbstractService<SessionEvent> implements S
 
     this.httpService.setHost(host)
 
-    await this.setSession(session)
+    this.setSession(session)
 
-    this.webSocketsService.startWebSocketConnection(session.authorizationValue)
+    await this.webSocketsService.startWebSocketConnection()
   }
 
   private async handleAuthResponse(body: UserRegistrationResponseBody, rootKey: SNRootKey, wrappingKey?: SNRootKey) {