diff --git a/packages/trezor-user-env-link/package.json b/packages/trezor-user-env-link/package.json index 4b3e19655d9..821e375df03 100644 --- a/packages/trezor-user-env-link/package.json +++ b/packages/trezor-user-env-link/package.json @@ -10,12 +10,11 @@ }, "dependencies": { "@trezor/utils": "workspace:*", + "@trezor/websocket-client": "workspace:*", "cross-fetch": "^4.0.0", - "semver": "^7.6.3", - "ws": "^8.18.0" + "semver": "^7.6.3" }, "devDependencies": { - "@types/semver": "^7.5.8", - "@types/ws": "^8.5.12" + "@types/semver": "^7.5.8" } } diff --git a/packages/trezor-user-env-link/src/websocket-client.ts b/packages/trezor-user-env-link/src/websocket-client.ts index 8cc1eb8bdec..bc7333fa491 100644 --- a/packages/trezor-user-env-link/src/websocket-client.ts +++ b/packages/trezor-user-env-link/src/websocket-client.ts @@ -1,14 +1,14 @@ /* eslint-disable no-console */ -import WebSocket, { RawData } from 'ws'; import fetch from 'cross-fetch'; -import { createDeferred, Deferred, TypedEmitter } from '@trezor/utils'; +import { + WebsocketClient as WebsocketClientBase, + WebsocketResponse as WebsocketResponseData, +} from '@trezor/websocket-client'; import { Firmwares } from './types'; -const NOT_INITIALIZED = new Error('websocket_not_initialized'); - // Making the timeout high because the controller in trezor-user-env // must synchronously run actions on emulator and they may take a long time // (for example in case of Shamir backup) @@ -22,17 +22,11 @@ const USER_ENV_URL = { DASHBOARD: `http://127.0.0.1:9002`, }; -interface Options { - pingTimeout?: number; - url?: string; - timeout?: number; -} - const delay = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); export type WebsocketClientEvents = { - firmwares: (firmwares: Firmwares) => void; - disconnected: () => void; + firmwares: Firmwares; + disconnected: undefined; }; interface WebsocketRequest { @@ -43,140 +37,33 @@ interface WebsocketResponse { response: any; } -export class WebsocketClient extends TypedEmitter { - private messageID: number; - private options: Options; - private messages: Deferred[]; - - private ws?: WebSocket; - private connectionTimeout?: NodeJS.Timeout; - private pingTimeout?: NodeJS.Timeout; - - constructor(options: Options = {}) { - super(); - this.messageID = 0; - this.messages = []; - this.setMaxListeners(Infinity); - this.options = { - ...options, - url: options.url || USER_ENV_URL.WEBSOCKET, - }; - } - - private setConnectionTimeout() { - this.clearConnectionTimeout(); - this.connectionTimeout = setTimeout( - this.onTimeout.bind(this), - this.options.timeout || DEFAULT_TIMEOUT, - ); - } - - private clearConnectionTimeout() { - if (this.connectionTimeout) { - clearTimeout(this.connectionTimeout); - this.connectionTimeout = undefined; - } - } - - private setPingTimeout() { - if (this.pingTimeout) { - clearTimeout(this.pingTimeout); - } - this.pingTimeout = setTimeout( - this.onPing.bind(this), - this.options.pingTimeout || DEFAULT_PING_TIMEOUT, - ); +export class WebsocketClient extends WebsocketClientBase { + protected createWebsocket() { + return this.initWebsocket(this.options); } - private onTimeout() { - const { ws } = this; - if (!ws) return; - if (ws.listenerCount('open') > 0) { - ws.emit('error', 'Websocket timeout'); - try { - ws.close(); - } catch { - // empty - } - } else { - this.messages.forEach(m => m.reject(new Error('websocket_timeout'))); - ws.close(); - } + protected ping() { + return this.send({ type: 'ping' }); } - private onPing() { - // make sure that connection is alive if there are subscriptions - if (this.ws && this.isConnected()) { - try { - this.ws.close(); - } catch { - // empty - } - } - } - - private onError() { - this.dispose(); + constructor(options: any = {}) { + super({ + ...options, + url: options.url || USER_ENV_URL.WEBSOCKET, + timeout: options.timeout || DEFAULT_TIMEOUT, + pingTimeout: options.pingTimeout || DEFAULT_PING_TIMEOUT, + }); } async send(params: T): Promise { // probably after update to node 18 it started to disconnect after certain // period of inactivity. await this.connect(); - const { ws } = this; - if (!ws) throw NOT_INITIALIZED; - const id = this.messageID; - - const dfd = createDeferred<{ response: any }>(id); - const req = { - id, - ...params, - }; - - this.messageID++; - this.messages.push(dfd); - - this.setConnectionTimeout(); - this.setPingTimeout(); - - ws.send(JSON.stringify(req)); - - return dfd.promise; - } - - private onmessage(message: RawData) { - try { - const resp = JSON.parse(message.toString()); - const { id, success } = resp; - const dfd = this.messages.find(m => m.id === id); - - if (resp.type === 'client') { - const { firmwares } = resp; - - this.emit('firmwares', firmwares); - } - - if (dfd) { - if (!success) { - dfd.reject( - new Error(`websocket_error_message: ${resp.error.message || resp.error}`), - ); - } else { - dfd.resolve(resp); - } - this.messages.splice(this.messages.indexOf(dfd), 1); - } - } catch (error) { - console.error('websocket onmessage error: ', error); - } - if (this.messages.length === 0) { - this.clearConnectionTimeout(); - } - this.setPingTimeout(); + return this.sendMessage(params); } - public async connect() { + async connect() { if (this.isConnected()) return Promise.resolve(); // workaround for karma... proper fix: set allow origin headers in trezor-user-env server. but we are going @@ -185,92 +72,33 @@ export class WebsocketClient extends TypedEmitter { await this.waitForTrezorUserEnv(); } - return new Promise(resolve => { - // url validation - let { url } = this.options; - if (typeof url !== 'string') { - throw new Error('websocket_no_url'); - } - - if (url.startsWith('https')) { - url = url.replace('https', 'wss'); - } - if (url.startsWith('http')) { - url = url.replace('http', 'ws'); - } - - // set connection timeout before WebSocket initialization - // it will be be cancelled by this.init or this.dispose after the error - this.setConnectionTimeout(); - - // initialize connection - const ws = new WebSocket(url); - - ws.once('error', error => { - console.error('websocket error', error); - this.dispose(); - }); - - this.on('firmwares', () => { - resolve(this); - }); - - this.ws = ws; - - ws.on('open', () => { - this.init(); + return new Promise(resolve => { + super.connect().then(() => { + this.once('firmwares', () => resolve()); }); }); } - private init() { - const { ws } = this; - if (!ws || !this.isConnected()) { - throw Error('Websocket init cannot be called'); - } - // clear timeout from this.connect - this.clearConnectionTimeout(); - - // remove previous listeners and add new listeners - ws.removeAllListeners(); - ws.on('error', this.onError.bind(this)); - ws.on('message', this.onmessage.bind(this)); - ws.on('close', () => { - this.emit('disconnected'); - this.dispose(); - }); - } - - public disconnect() { - if (this.ws) { - this.ws.close(); - } - this.dispose(); - } - - private isConnected() { - const { ws } = this; - - return ws && ws.readyState === WebSocket.OPEN; - } + protected onMessage(message: WebsocketResponseData) { + try { + const resp = JSON.parse(message.toString()); + const { id, success } = resp; - private dispose() { - if (this.pingTimeout) { - clearTimeout(this.pingTimeout); - } - if (this.connectionTimeout) { - clearTimeout(this.connectionTimeout); - } + if (resp.type === 'client') { + this.emit('firmwares', resp.firmwares); + } - const { ws } = this; - if (this.isConnected()) { - this.disconnect(); - } - if (ws) { - ws.removeAllListeners(); + if (!success) { + this.messages.reject( + Number(id), + new Error(`websocket_error_message: ${resp.error.message || resp.error}`), + ); + } else { + this.messages.resolve(Number(id), resp); + } + } catch { + // empty } - - this.removeAllListeners(); } async waitForTrezorUserEnv() { diff --git a/packages/trezor-user-env-link/tsconfig.json b/packages/trezor-user-env-link/tsconfig.json index 0ec4519c33d..cb94bd20b74 100644 --- a/packages/trezor-user-env-link/tsconfig.json +++ b/packages/trezor-user-env-link/tsconfig.json @@ -1,5 +1,8 @@ { "extends": "../../tsconfig.base.json", "compilerOptions": { "outDir": "libDev" }, - "references": [{ "path": "../utils" }] + "references": [ + { "path": "../utils" }, + { "path": "../websocket-client" } + ] } diff --git a/yarn.lock b/yarn.lock index 4d3f7dc8918..1b57a9a5a14 100644 --- a/yarn.lock +++ b/yarn.lock @@ -12843,11 +12843,10 @@ __metadata: resolution: "@trezor/trezor-user-env-link@workspace:packages/trezor-user-env-link" dependencies: "@trezor/utils": "workspace:*" + "@trezor/websocket-client": "workspace:*" "@types/semver": "npm:^7.5.8" - "@types/ws": "npm:^8.5.12" cross-fetch: "npm:^4.0.0" semver: "npm:^7.6.3" - ws: "npm:^8.18.0" languageName: unknown linkType: soft