mihawk
Version:
A tiny & simple mock server tool, support json,js,cjs,ts(typescript).
200 lines (199 loc) • 8.9 kB
JavaScript
;
import https from 'https';
import http from 'http';
import { promisify } from 'util';
import Colors from 'color-cc';
import * as WS from 'ws';
import { PKG_NAME } from '../consts';
import { parseStompMsg } from '../utils/parser';
import { Printer } from '../utils/print';
import { getTimeNowStr } from '../utils/date';
import { getAddrInfoByServer } from '../utils/server';
import { getMyIp, supportLocalHost } from '../utils/net';
const LOGFLAG_WS = `${Colors.cyan('[WS]')}${Colors.gray(':')}`;
export default class WsCtrl {
_useStompMsg;
_host;
_port;
_secure;
_baseServer;
_wsSvr;
_clientIndex = 0;
_resolveFunc;
constructor(options) {
const { host, port, secure = false, server, stomp = false, resolve = _defaultResolveFunc, } = options;
this._useStompMsg = !!stomp;
this._host = host;
this._port = port;
this._secure = secure;
this._wsSvr = null;
this._baseServer = server;
const _resolve = typeof resolve === 'function' ? resolve : _defaultResolveFunc;
this._resolveFunc = function (socket, request, options) {
try {
const funcCtx = socket;
_resolve.call(funcCtx, socket, request, options);
}
catch (error) {
Printer.error(LOGFLAG_WS, 'Failed to exec socket logic resolveFunc:\n', error);
}
};
}
_autoCreateClientId() {
return `Client-${++this._clientIndex}__${getTimeNowStr()}`;
}
async start(server, isRestart) {
if (this._wsSvr) {
Printer.log(LOGFLAG_WS, 'WS server is already running.');
return;
}
const resolveFunc = this._resolveFunc;
const stomp = this._useStompMsg;
if (server && _isHttpOrHttpsServer(server)) {
this._baseServer = server;
}
const { isSecure = this._secure, host = this._host, port = this._port } = _getBaseServerInfo(this._baseServer);
const protocol = isSecure ? 'wss' : 'ws';
const hostFix = host === '0.0.0.0' ? '127.0.0.1' : host;
this._wsSvr = new WS.WebSocketServer({ server: this._baseServer });
this._wsSvr.on('listening', function listening() {
Printer.log(`🚀 ${Colors.green(`${isRestart ? 'Restart' : 'Start'} mock-socket-server success!`)}`);
!isRestart && Printer.log(`Mock Socket Server address:`);
Printer.log(`${Colors.gray('-')} ${Colors.cyan(`${protocol}://${hostFix}:${port}`)}`);
if (supportLocalHost(host)) {
const addr2 = `${protocol}://${getMyIp()}:${port}`;
Printer.log(`${Colors.gray('-')} ${Colors.cyan(addr2)}`);
}
console.log();
});
this._wsSvr.on('headers', function headers(headers, req) {
headers.push(`X-Powered-By: ${PKG_NAME}`);
});
this._wsSvr.on('connection', (socket, request) => {
const { remoteAddress, remotePort } = request?.socket || {};
const clientId = remoteAddress ? `${remoteAddress}:${remotePort || ++this._clientIndex}` : this._autoCreateClientId();
Printer.log(LOGFLAG_WS, 'Socket client connected!', Colors.gray(`clientId=[${clientId}]`), Colors.gray(`time=${getTimeNowStr()}`));
resolveFunc(socket, request, { clientId, stomp });
setTimeout(() => {
const clientId4Log = Colors.gray(`clientId=[${clientId}]`);
const timeNow4Log = Colors.gray(`time=${getTimeNowStr()}`);
if (this._wsSvr) {
try {
socket?.emit('open', () => {
Printer.log(LOGFLAG_WS, "Socket client opened! (socket.emit('open') success!)", clientId4Log, timeNow4Log);
});
}
catch (error) {
Printer.error(LOGFLAG_WS, "Failed to exec socket.emit('open') ! ", clientId4Log, timeNow4Log, '\n', error);
}
}
else {
Printer.warn(LOGFLAG_WS, "Cancel to exec socket.emit('open'), because connection is closed.", clientId4Log, timeNow4Log);
}
}, 0);
});
this._wsSvr.on('error', function error(err) {
Printer.error(LOGFLAG_WS, 'WebSocket server error:', err);
});
this._wsSvr.on('close', function close() {
});
}
async _close(forceClose = false) {
if (!this._wsSvr) {
Printer.log(LOGFLAG_WS, 'WebSocket server is not running.');
return;
}
const clients = this._wsSvr?.clients;
try {
if (clients?.size) {
if (forceClose) {
clients?.forEach(client => client?.terminate());
}
else {
clients?.forEach(client => client?.close());
}
}
}
catch (error) {
Printer.error(LOGFLAG_WS, `Batch ${forceClose ? 'terminate' : 'close'} clients failed!\n`, error);
}
const closeAsync = promisify(this._wsSvr.close).bind(this._wsSvr);
try {
await closeAsync();
this._wsSvr = null;
}
catch (error) {
Printer.error(LOGFLAG_WS, `${forceClose ? 'Destory' : 'Close'} WebSocket server failed!\n`, error);
}
}
async close() {
return await this._close();
}
async destory() {
return await this._close(true);
}
}
function _defaultResolveFunc(socket, request, options) {
const { clientId: cid, stomp } = options || {};
const clientId = cid || request.socket.remoteAddress;
const clientName = `[${clientId}]`;
const logTail = Colors.gray(`from ${clientName}`);
const logName = Colors.gray('socket:');
socket.send(JSON.stringify({ success: true, data: `WsServer: Connection established! Hello, client! ${clientName}` }));
socket.on('message', (message, isBinary) => {
const recived = typeof message === 'string' ? message : Buffer.isBuffer(message) ? message?.toString() : JSON.stringify(message);
let recivedData = recived;
if (stomp) {
recivedData = parseStompMsg(recived);
Printer.log(LOGFLAG_WS, logName, 'Received stomp message', logTail);
Printer.log(LOGFLAG_WS, logName, '<= stompData:', recivedData);
}
else {
Printer.log(LOGFLAG_WS, logName, `Received message <= "${Colors.green(recived)}"`, isBinary ? Colors.gray('binary') : '', logTail);
}
const recivedDataStr = typeof recivedData === 'string' ? recivedData : JSON.stringify(recivedData);
const msgData = {
success: true,
data: `WsServer: I've recived your message(${recivedDataStr})`,
};
Printer.log(LOGFLAG_WS, logName, `Send response to ${Colors.gray(clientName)} =>`, msgData);
socket.send(JSON.stringify(msgData));
console.log('\n');
});
socket.on('open', () => {
Printer.log(LOGFLAG_WS, logName, `Client ${Colors.gray(clientId)} open.`, logTail);
});
socket.on('upgrade', (request) => {
const clientAddr = request.socket.remoteAddress;
Printer.log(LOGFLAG_WS, logName, `Client ${clientAddr} upgraded.`, logTail);
});
socket.on('ping', (data) => {
Printer.log(LOGFLAG_WS, logName, `Received ping => ${data?.toString() || ''}`, logTail);
socket.pong(`Pong: ${data?.toString()}`);
});
socket.on('pong', (data) => {
Printer.log(LOGFLAG_WS, logName, `Received pong => ${data?.toString() || ''}`, logTail);
});
socket.on('error', (err) => {
const errMsg = err?.message || err?.toString() || 'unknow error';
Printer.error(LOGFLAG_WS, logName, `CLient error: ${Colors.red(errMsg)}`, logTail);
Printer.error(LOGFLAG_WS, err, '\n');
});
socket.on('unexpected-response', (request, response) => {
const exceptDetail = `${response.statusCode} ${response.statusMessage}`;
Printer.error(LOGFLAG_WS, logName, `CLient unexpected-response: ${Colors.red(exceptDetail)}`, logTail, '\n');
});
socket.on('close', (code, reason) => {
const closeDetail = `code=${code},reason=${reason.toString() || 'none'}`;
Printer.log(LOGFLAG_WS, logName, `Client close connection.(${Colors.yellow(closeDetail)})`, logTail, '\n');
});
}
function _getBaseServerInfo(server) {
const isSecure = server instanceof https.Server;
const protocol = isSecure ? 'https' : 'http';
const { port, address } = getAddrInfoByServer(server) || {};
return { isSecure, protocol, host: address, port };
}
function _isHttpOrHttpsServer(server) {
return server instanceof http.Server || server instanceof https.Server;
}