Skip to content

Commit

Permalink
chore(trezor-user-env-link): use @trezor/websocket-client package
Browse files Browse the repository at this point in the history
  • Loading branch information
@szymonlesisz
szymonlesisz committed Jan 30, 2025
1 parent a06a6ab commit f65ec02
Show file tree
Hide file tree
Showing 4 changed files with 48 additions and 217 deletions.
7 changes: 3 additions & 4 deletions packages/trezor-user-env-link/package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.13"
"@types/semver": "^7.5.8"
}
}
250 changes: 40 additions & 210 deletions packages/trezor-user-env-link/src/websocket-client.ts
View file
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
/* eslint-disable no-console */

import fetch from 'cross-fetch';
import WebSocket, { RawData } from 'ws';

import { Deferred, TypedEmitter, createDeferred } 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)
Expand All @@ -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 {
Expand All @@ -43,140 +37,33 @@ interface WebsocketResponse {
response: any;
}

export class WebsocketClient extends TypedEmitter<WebsocketClientEvents> {
private messageID: number;
private options: Options;
private messages: Deferred<any>[];

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,
};
export class WebsocketClient extends WebsocketClientBase<WebsocketClientEvents> {
protected createWebsocket() {
return this.initWebsocket(this.options);
}

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;
}
protected ping() {
return this.send({ type: 'ping' });
}

private setPingTimeout() {
if (this.pingTimeout) {
clearTimeout(this.pingTimeout);
}
this.pingTimeout = setTimeout(
this.onPing.bind(this),
this.options.pingTimeout || DEFAULT_PING_TIMEOUT,
);
}

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();
}
}

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<T extends WebsocketRequest>(params: T): Promise<WebsocketResponse> {
// 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;
return this.sendMessage(params);
}

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();
}

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
Expand All @@ -185,92 +72,35 @@ export class WebsocketClient extends TypedEmitter<WebsocketClientEvents> {
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<void>(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();
});
}
disconnect() {
// TODO: breaking change
// previous implementation `disconnect` acts like `dispose`. does not emit 'disconnected' event
if (this.isConnected()) {
this.removeAllListeners();

public disconnect() {
if (this.ws) {
this.ws.close();
return super.disconnect();
}
this.dispose();
}

private isConnected() {
const { ws } = this;

return ws && ws.readyState === WebSocket.OPEN;
}

private dispose() {
if (this.pingTimeout) {
clearTimeout(this.pingTimeout);
}
if (this.connectionTimeout) {
clearTimeout(this.connectionTimeout);
}

const { ws } = this;
if (this.isConnected()) {
this.disconnect();
}
if (ws) {
ws.removeAllListeners();
}
protected onMessage(message: WebsocketResponseData) {
super.onMessage(message, resp => {
if (resp.type === 'client') {
this.emit('firmwares', resp.firmwares);
} else {
if (!resp.success) {
throw new Error(`websocket_error_message: ${resp.error.message || resp.error}`);
}
}

this.removeAllListeners();
return resp;
});
}

async waitForTrezorUserEnv() {
Expand Down
5 changes: 4 additions & 1 deletion packages/trezor-user-env-link/tsconfig.json
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
{
"extends": "../../tsconfig.base.json",
"compilerOptions": { "outDir": "libDev" },
"references": [{ "path": "../utils" }]
"references": [
{ "path": "../utils" },
{ "path": "../websocket-client" }
]
}
3 changes: 1 addition & 2 deletions yarn.lock
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12873,11 +12873,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.13"
cross-fetch: "npm:^4.0.0"
semver: "npm:^7.6.3"
ws: "npm:^8.18.0"
languageName: unknown
linkType: soft

Expand Down

0 comments on commit f65ec02

Please sign in to comment.