@renegade-fi/core/src/utils/websocket.ts

VanillaJS library for Renegade

import { SIG_EXPIRATION_BUFFER_MS } from '../constants.js'
import type { RenegadeConfig } from '../createConfig.js'
import {
  SocketClosedError,
  WebSocketConnectionError,
  WebSocketRequestError,
} from '../errors/websocket.js'
import { addExpiringAuthToHeaders } from './http.js'

export enum AuthType {
  None = 'None',
  Wallet = 'Wallet',
  Admin = 'Admin',
}

export type RelayerWebsocketParams = {
  config: RenegadeConfig
  topic: string
  authType: AuthType
  onmessage: (this: WebSocket, ev: MessageEvent) => any
  onopenCallback?: (this: WebSocket, ev: Event) => any
  oncloseCallback?: (this: WebSocket, ev: CloseEvent) => any
  onerrorCallback?: (this: WebSocket, ev: Event) => any
}

type SubscriptionMessage = {
  headers?: Record<string, string>
  body: SubscriptionBody
}

type UnsubscriptionMessage = {
  body: UnsubscriptionBody
}

export type SubscriptionBody = {
  method: 'subscribe'
  topic: string
}

export type UnsubscriptionBody = {
  method: 'unsubscribe'
  topic: string
}

export class RelayerWebsocket {
  private config: RenegadeConfig
  private topic: string
  private authType: AuthType
  private onmessage: (this: WebSocket, ev: MessageEvent) => any
  private onopenCallback: ((this: WebSocket, ev: Event) => any) | null
  private oncloseCallback: ((this: WebSocket, ev: CloseEvent) => any) | null
  private onerrorCallback: ((this: WebSocket, ev: Event) => any) | null

  private ws: WebSocket | null = null

  private handleOpen = (event: Event) => {
    if (!this.ws) return
    const message = this.buildSubscriptionMessage()
    this.request(message)

    return this.onopenCallback?.call(this.ws, event)
  }

  private handleClose = (event: CloseEvent) => {
    this.cleanup()
    return this.oncloseCallback?.call(this.ws!, event)
  }

  private handleError = async (event: Event) => {
    this.cleanup()
    return this.onerrorCallback?.call(this.ws!, event)
  }

  constructor(params: RelayerWebsocketParams) {
    this.config = params.config
    this.topic = params.topic
    this.authType = params.authType
    this.onmessage = params.onmessage
    this.onopenCallback = params.onopenCallback ?? null
    this.oncloseCallback = params.oncloseCallback ?? null
    this.onerrorCallback = params.onerrorCallback ?? null
  }

  // --------------
  // | Public API |
  // --------------

  public async connect(): Promise<void> {
    if (this.ws) {
      throw new Error(
        'WebSocket connection attempt aborted: already connected.',
      )
    }

    const WebSocket = await import('isows').then((module) => module.WebSocket)
    const url = this.config.getWebsocketBaseUrl()
    this.ws = new WebSocket(url)

    this.ws.addEventListener('open', this.handleOpen)
    this.ws.addEventListener('message', this.onmessage)
    this.ws.addEventListener('close', this.handleClose)
    this.ws.addEventListener('error', this.handleError)

    // Wait for the socket to open.
    if (this.ws?.readyState === WebSocket.CONNECTING) {
      await new Promise((resolve, reject) => {
        if (!this.ws) return
        this.ws.onopen = (event) => resolve(event)
        this.ws.onerror = (error) =>
          reject(
            new WebSocketConnectionError({
              url,
              cause: error as unknown as Error,
            }),
          )
      })
    }
  }

  public close(): void {
    if (!this.ws) {
      throw new Error('WebSocket connection not open')
    }

    const message = this.buildUnsubscriptionMessage()
    this.request(message)

    this.ws.close()
  }

  // ---------------
  // | Private API |
  // ---------------

  private request(message: SubscriptionMessage | UnsubscriptionMessage): void {
    if (
      this.ws?.readyState === this.ws?.CLOSED ||
      this.ws?.readyState === this.ws?.CLOSING
    ) {
      throw new WebSocketRequestError({
        body: message,
        url: this.ws?.url || '',
        cause: new SocketClosedError({ url: this.ws?.url }),
      })
    }

    this.ws?.send(JSON.stringify(message))
  }

  private buildSubscriptionMessage(): SubscriptionMessage {
    const body = {
      method: 'subscribe' as const,
      topic: this.topic,
    }

    if (this.authType === AuthType.None) {
      return { body }
    }

    const headers = this.buildAuthHeaders(body)
    return {
      headers,
      body,
    }
  }

  private buildUnsubscriptionMessage(): UnsubscriptionMessage {
    return {
      body: {
        method: 'unsubscribe' as const,
        topic: this.topic,
      },
    }
  }

  private buildAuthHeaders(body: SubscriptionBody): Record<string, string> {
    const symmetricKey = this.config.getSymmetricKey(this.authType)

    return addExpiringAuthToHeaders(
      this.config,
      body.topic,
      {}, // Headers
      JSON.stringify(body),
      symmetricKey,
      SIG_EXPIRATION_BUFFER_MS,
    )
  }

  private cleanup(): void {
    // Remove all event listeners before nullifying the reference
    if (this.ws) {
      this.ws.removeEventListener('open', this.handleOpen)
      this.ws.removeEventListener('message', this.onmessage)
      this.ws.removeEventListener('close', this.handleClose)
      this.ws.removeEventListener('error', this.handleError)
    }

    // Nullify the WebSocket instance
    this.ws = null
  }
}