@farmfe/core
Version:
Farm is a extremely fast web build tool written in Rust. Farm can start a project in milliseconds and perform HMR within 10ms, making it much faster than similar tools like webpack and vite.
257 lines • 9.3 kB
JavaScript
import { WebSocket, WebSocketServer as WebSocketServerRaw } from 'ws';
import { Logger, red } from '../index.js';
import { resolveHostname, resolveServerUrls } from '../utils/http.js';
const HMR_HEADER = 'farm_hmr';
const wsServerEvents = [
'connection',
'error',
'headers',
'listening',
'message'
];
export default class WsServer {
constructor(httpServer, config, hmrEngine, logger) {
this.httpServer = httpServer;
this.config = config;
this.hmrEngine = hmrEngine;
this.customListeners = new Map();
this.clientsMap = new WeakMap();
this.bufferedError = null;
this.logger = logger ?? new Logger();
this.hmrOrigins = this.generateHMROrigins(config);
this.createWebSocketServer();
}
generateHMROrigins(config) {
const { protocol, hostname, port } = config;
const origins = [];
// Add localhost with configured port
const urls = resolveServerUrls(this.httpServer, config);
const localUrls = [...(urls.local || []), ...(urls.network || [])];
for (const url of localUrls) {
origins.push(url);
}
// Add non-localhost origin
const configuredOrigin = `${protocol}://${hostname.name}:${port}`;
if (hostname &&
hostname.name &&
localUrls.every((url) => url !== configuredOrigin)) {
origins.push(configuredOrigin);
}
if (this.config.host !== this.config.hmr.host) {
const hmrHostname = resolveHostname(this.config.hmr.host);
origins.push(`${this.config.hmr?.protocol || protocol}://${hmrHostname.name}:${this.config.hmr?.port || this.config.port}`);
}
return origins;
}
createWebSocketServer() {
try {
const WebSocketServer = process.versions.bun
? // @ts-expect-error: Bun defines `import.meta.require`
import.meta.require('ws').WebSocketServer
: WebSocketServerRaw;
if (this.config.host === this.config.hmr.host) {
this.wss = new WebSocketServer({ noServer: true });
this.connection();
this.httpServer.on('upgrade', this.upgradeWsServer.bind(this));
}
else {
const hmrHostname = resolveHostname(this.config.hmr.host);
this.wss = new WebSocketServer({
host: hmrHostname.name,
port: this.config.hmr?.port || this.config.port
});
this.connection();
}
}
catch (err) {
this.handleSocketError(err);
}
}
upgradeWsServer(request, socket, head) {
if (this.isHMRRequest(request)) {
this.handleHMRUpgrade(request, socket, head);
}
else {
this.logger.error(`HMR upgrade failed because of invalid HMR path, header or host. The HMR connection is only allowed on hosts: ${this.hmrOrigins.join(', ')}. You can try set server.host or server.allowedHosts to allow the connection.`);
}
}
listen() {
// TODO alone with use httpServer we need start this function
// Start listening for WebSocket connections
}
// Farm uses the `sendMessage` method in hmr and
// the send method is reserved for migration vite
send(...args) {
let payload;
if (typeof args[0] === 'string') {
payload = {
type: 'custom',
event: args[0],
data: args[1]
};
}
else {
payload = args[0];
}
if (payload.type === 'error' && !this.wss.clients.size) {
this.bufferedError = payload;
return;
}
const stringified = JSON.stringify(payload);
this.wss.clients.forEach((client) => {
// readyState 1 means the connection is open
if (client.readyState === 1) {
client.send(stringified);
}
});
}
isHMRRequest(request) {
return (request.url === this.config.hmr.path &&
request.headers['sec-websocket-protocol'] === HMR_HEADER);
}
handleHMRUpgrade(request, socket, head) {
this.wss.handleUpgrade(request, socket, head, (ws) => {
this.wss.emit('connection', ws, request);
});
}
get clients() {
return new Set(Array.from(this.wss.clients).map(this.getSocketClient.bind(this)));
}
// a custom method defined by farm to send custom events
sendCustomEvent(event, payload) {
// Send a custom event to all clients
this.send({ type: 'custom', event, data: payload });
}
on(event, listener) {
if (wsServerEvents.includes(event)) {
this.wss.on(event, listener);
}
else {
this.addCustomEventListener(event, listener);
}
}
off(event, listener) {
if (wsServerEvents.includes(event)) {
this.wss.off(event, listener);
}
else {
this.removeCustomEventListener(event, listener);
}
}
connection() {
this.wss.on('connection', (socket) => {
socket.on('message', (raw) => {
if (!this.customListeners.size)
return;
let parsed;
try {
parsed = JSON.parse(String(raw));
}
catch {
this.logger.error('Failed to parse WebSocket message: ' + raw);
}
// transform vite js-update to farm update
if (parsed?.type === 'js-update' && parsed?.path) {
this.hmrEngine.hmrUpdate(parsed.path, true);
return;
}
if (!parsed || parsed.type !== 'custom' || !parsed.event)
return;
const listeners = this.customListeners.get(parsed.event);
if (!listeners?.size)
return;
const client = this.getSocketClient(socket);
listeners.forEach((listener) => listener(parsed.data, client));
});
socket.on('error', (err) => {
return this.handleSocketError(err);
});
socket.send(JSON.stringify({ type: 'connected' }));
if (this.bufferedError) {
socket.send(JSON.stringify(this.bufferedError));
this.bufferedError = null;
}
});
}
async close() {
if (this.upgradeWsServer && this.httpServer) {
this.httpServer.off('upgrade', this.upgradeWsServer);
}
await this.terminateAllClients();
await this.closeWebSocketServer();
// TODO if not have httpServer we need close httpServer
}
terminateAllClients() {
const terminatePromises = Array.from(this.wss.clients).map((client) => {
return new Promise((resolve) => {
if (client.readyState === WebSocket.OPEN) {
client.send(JSON.stringify({ type: 'closing' }));
client.close(1000, 'Server shutdown');
}
// Temporarily remove the direct shutdown of ws
// client.terminate();
client.once('close', () => resolve(true));
});
});
return Promise.all(terminatePromises);
}
closeWebSocketServer() {
return new Promise((resolve, reject) => {
this.wss.close((err) => {
if (err) {
reject(err);
}
else {
// TODO if not have httpServer
resolve(true);
}
});
});
}
addCustomEventListener(event, listener) {
if (!this.customListeners.has(event)) {
this.customListeners.set(event, new Set());
}
this.customListeners.get(event).add(listener);
}
removeCustomEventListener(event, listener) {
this.customListeners.get(event)?.delete(listener);
}
getSocketClient(socket) {
if (!this.clientsMap.has(socket)) {
this.clientsMap.set(socket, {
send: (...args) => this.sendMessage(socket, ...args),
socket,
rawSend: (str) => socket.send(str)
});
}
return this.clientsMap.get(socket);
}
sendMessage(socket, ...args) {
let payload;
if (typeof args[0] === 'string') {
payload = {
type: 'custom',
event: args[0],
data: args[1]
};
}
else {
payload = args[0];
}
socket.send(JSON.stringify(payload));
}
handleSocketError(err) {
if (err.code === 'EADDRINUSE') {
this.logger.error(red(`WebSocket server error: Port is already in use`), {
error: err
});
}
else {
this.logger.error(red(`WebSocket server error:\n${err.stack || err.message}`), {
error: err
});
}
}
}
//# sourceMappingURL=ws.js.map