85 lines
2.1 KiB
TypeScript
85 lines
2.1 KiB
TypeScript
import { nanoid } from 'nanoid';
|
|
import { WebSocket } from 'ws';
|
|
import { logger } from '../logger.ts';
|
|
class WsMessage {
|
|
ws: WebSocket;
|
|
user?: string;
|
|
constructor({ ws, user }: WssMessageOptions) {
|
|
this.ws = ws;
|
|
this.user = user;
|
|
}
|
|
async sendData(data: any, opts?: { timeout?: number }) {
|
|
if (this.ws.readyState !== WebSocket.OPEN) {
|
|
return { code: 500, message: 'WebSocket is not open' };
|
|
}
|
|
const timeout = opts?.timeout || 10 * 6 * 1000; // 10 minutes
|
|
const id = nanoid();
|
|
const message = JSON.stringify({
|
|
id,
|
|
type: 'proxy',
|
|
data,
|
|
});
|
|
logger.info('ws-proxy sendData', message);
|
|
this.ws.send(message);
|
|
return new Promise((resolve) => {
|
|
const timer = setTimeout(() => {
|
|
resolve({
|
|
code: 500,
|
|
message: 'timeout',
|
|
});
|
|
}, timeout);
|
|
this.ws.once('message', (event: Buffer) => {
|
|
const eventData = event.toString();
|
|
if (!eventData) {
|
|
return;
|
|
}
|
|
const data = JSON.parse(eventData);
|
|
if (data.id === id) {
|
|
resolve(data.data);
|
|
clearTimeout(timer);
|
|
}
|
|
});
|
|
});
|
|
}
|
|
}
|
|
type WssMessageOptions = {
|
|
ws: WebSocket;
|
|
user?: string;
|
|
};
|
|
export class WsProxyManager {
|
|
wssMap: Map<string, WsMessage> = new Map();
|
|
constructor() {}
|
|
getId(id: string, user?: string) {
|
|
return id + '/' + user;
|
|
}
|
|
register(id: string, opts?: { ws: WebSocket; user: string }) {
|
|
const _id = this.getId(id, opts?.user || '');
|
|
if (this.wssMap.has(_id)) {
|
|
const value = this.wssMap.get(_id);
|
|
if (value) {
|
|
value.ws.close();
|
|
}
|
|
}
|
|
const value = new WsMessage({ ws: opts?.ws, user: opts?.user });
|
|
this.wssMap.set(_id, value);
|
|
}
|
|
unregister(id: string, user?: string) {
|
|
const _id = this.getId(id, user || '');
|
|
const value = this.wssMap.get(_id);
|
|
if (value) {
|
|
value.ws.close();
|
|
}
|
|
this.wssMap.delete(_id);
|
|
}
|
|
getIds() {
|
|
return Array.from(this.wssMap.keys());
|
|
}
|
|
get(id: string, user?: string) {
|
|
if (user) {
|
|
const _id = this.getId(id, user);
|
|
return this.wssMap.get(_id);
|
|
}
|
|
return this.wssMap.get(id);
|
|
}
|
|
}
|