@roochnetwork/rooch-sdk/src/client/wsTransport.ts

Rooch SDK

// Copyright (c) RoochNetwork
// SPDX-License-Identifier: Apache-2.0

import { JsonRpcError } from './error.js'

function getWebsocketUrl(httpUrl: string): string {
  const url = new URL(httpUrl)
  url.protocol = url.protocol.replace('http', 'ws')
  return url.toString()
}

type JsonRpcMessage =
  | {
      id: number
      result: never
      error: {
        code: number
        message: string
      }
    }
  | {
      id: number
      result: unknown
      error: never
    }
  | {
      method: string
      params: NotificationMessageParams
    }

type NotificationMessageParams = {
  subscription?: number
  result: object
}

type SubscriptionRequest<T = any> = {
  method: string
  params: any[]
  onMessage: (event: T) => void
  signal?: AbortSignal
}

/**
 * Configuration options for the websocket connection
 */
export type WebsocketClientOptions = {
  /**
   * Custom WebSocket class to use. Defaults to the global WebSocket class, if available.
   */
  WebSocketConstructor?: typeof WebSocket
  /**
   * Milliseconds before timing out while calling an RPC method
   */
  callTimeout?: number
  /**
   * Milliseconds between attempts to connect
   */
  reconnectTimeout?: number
  /**
   * Maximum number of times to try connecting before giving up
   */
  maxReconnects?: number
}

export const DEFAULT_CLIENT_OPTIONS = {
  // We fudge the typing because we also check for undefined in the constructor:
  WebSocketConstructor: (typeof WebSocket !== 'undefined'
    ? WebSocket
    : undefined) as typeof WebSocket,
  callTimeout: 30000,
  reconnectTimeout: 3000,
  maxReconnects: 5,
} satisfies WebsocketClientOptions

export class WebsocketClient {
  endpoint: string
  options: Required<WebsocketClientOptions>
  #requestId = 0
  #disconnects = 0
  #webSocket: WebSocket | null = null
  #connectionPromise: Promise<WebSocket> | null = null
  #subscriptions = new Set<RpcSubscription>()
  #pendingRequests = new Map<
    number,
    {
      resolve: (result: Extract<JsonRpcMessage, { id: number }>) => void
      reject: (reason: unknown) => void
      timeout: ReturnType<typeof setTimeout>
    }
  >()

  constructor(endpoint: string, options: WebsocketClientOptions = {}) {
    this.endpoint = endpoint
    this.options = { ...DEFAULT_CLIENT_OPTIONS, ...options }

    if (!this.options.WebSocketConstructor) {
      throw new Error('Missing WebSocket constructor')
    }

    if (this.endpoint.startsWith('http')) {
      this.endpoint = getWebsocketUrl(this.endpoint)
    }
  }

  async makeRequest<T>(method: string, params: any[], signal?: AbortSignal): Promise<T> {
    const webSocket = await this.#setupWebSocket()

    return new Promise<Extract<JsonRpcMessage, { id: number }>>((resolve, reject) => {
      this.#requestId += 1
      this.#pendingRequests.set(this.#requestId, {
        resolve: resolve,
        reject,
        timeout: setTimeout(() => {
          this.#pendingRequests.delete(this.#requestId)
          reject(new Error(`Request timeout: ${method}`))
        }, this.options.callTimeout),
      })

      signal?.addEventListener('abort', () => {
        this.#pendingRequests.delete(this.#requestId)
        reject(signal.reason)
      })

      webSocket.send(JSON.stringify({ jsonrpc: '2.0', id: this.#requestId, method, params }))
    }).then(({ error, result }) => {
      if (error) {
        throw new JsonRpcError(error.message, error.code)
      }

      return result as T
    })
  }

  #setupWebSocket() {
    if (this.#connectionPromise) {
      return this.#connectionPromise
    }

    this.#connectionPromise = new Promise<WebSocket>((resolve) => {
      this.#webSocket?.close()
      this.#webSocket = new this.options.WebSocketConstructor(this.endpoint)

      this.#webSocket.addEventListener('open', () => {
        this.#disconnects = 0
        resolve(this.#webSocket!)
      })

      this.#webSocket.addEventListener('close', () => {
        this.#disconnects++
        if (this.#disconnects <= this.options.maxReconnects) {
          setTimeout(() => {
            this.#reconnect()
          }, this.options.reconnectTimeout)
        }
      })

      this.#webSocket.addEventListener('message', ({ data }: { data: string }) => {
        let json: JsonRpcMessage
        try {
          json = JSON.parse(data) as JsonRpcMessage
        } catch (error) {
          console.error(new Error(`Failed to parse RPC message: ${data}`, { cause: error }))
          return
        }

        if ('id' in json && json.id != null && this.#pendingRequests.has(json.id)) {
          const { resolve, timeout } = this.#pendingRequests.get(json.id)!

          clearTimeout(timeout)
          resolve(json)
        } else if ('params' in json) {
          const { params } = json
          this.#subscriptions.forEach((subscription) => {
            if (subscription.subscriptionId === params.subscription)
              if (params.subscription === subscription.subscriptionId) {
                subscription.onMessage(params.result)
              }
          })
        }
      })
    })

    return this.#connectionPromise
  }

  async #reconnect() {
    this.#webSocket?.close()
    this.#connectionPromise = null

    return Promise.allSettled(
      [...this.#subscriptions].map((subscription) => subscription.subscribe(this)),
    )
  }

  async subscribe<T>(input: SubscriptionRequest<T>) {
    const subscription = new RpcSubscription(input)
    this.#subscriptions.add(subscription)
    await subscription.subscribe(this)
    return () => subscription.unsubscribe(this)
  }
}

class RpcSubscription {
  subscriptionId: number | null = null
  input: SubscriptionRequest<any>
  subscribed = false

  constructor(input: SubscriptionRequest) {
    this.input = input
  }

  onMessage(message: unknown) {
    if (this.subscribed) {
      this.input.onMessage(message)
    }
  }

  // TODO
  async unsubscribe(_: WebsocketClient) {
    const { subscriptionId } = this
    this.subscribed = false
    if (subscriptionId == null) return false
    this.subscriptionId = null

    return true

    // return client.makeRequest('rooch_unsubscribe', [subscriptionId])
  }

  async subscribe(client: WebsocketClient) {
    this.subscriptionId = null
    this.subscribed = true
    const newSubscriptionId = await client.makeRequest<number>(
      this.input.method,
      this.input.params,
      this.input.signal,
    )

    if (this.subscribed) {
      this.subscriptionId = newSubscriptionId
    }
  }
}