60 lines
2.3 KiB
JavaScript
60 lines
2.3 KiB
JavaScript
import { WebSocketServer } from 'ws';
|
|
import { Client } from './client.js';
|
|
import { logger } from '../logger/logger.js';
|
|
|
|
/**
|
|
* Attaches the WebSocket server upgrade hook to the given HTTP server instance.
|
|
* It intercepts standard HTTP GET requests for WebSocket protocols and handles upgrades.
|
|
* @param {object} server The native http.Server instance.
|
|
* @param {object} hub The central Hub coordinator instance.
|
|
* @param {object} limiter The IP-based RateLimiter instance.
|
|
*/
|
|
export function setupWebSocket(server, hub, limiter) {
|
|
// Initialize ws server in noServer mode since we manually intercept upgrade requests.
|
|
const wss = new WebSocketServer({ noServer: true });
|
|
|
|
/**
|
|
* Listens for the HTTP 'upgrade' event to switch protocol from HTTP to WebSocket.
|
|
*/
|
|
server.on('upgrade', (request, socket, head) => {
|
|
const ip = getClientIP(request);
|
|
|
|
// Enforce IP-based rate limiting to prevent spam connections.
|
|
if (!limiter.allow(ip)) {
|
|
logger.warn('rate_limit_exceeded', { remote_ip: ip });
|
|
// Send 429 Too Many Requests response to client and close connection.
|
|
socket.write('HTTP/1.1 429 Too Many Requests\r\nConnection: close\r\n\r\nrate_limited');
|
|
socket.destroy();
|
|
return;
|
|
}
|
|
|
|
// Hand over control to 'ws' library to complete the upgrade protocol handshake.
|
|
wss.handleUpgrade(request, socket, head, (wsConn) => {
|
|
// Set remote address on socket mock for later client IP queries.
|
|
// This ensures we can resolve client IP even after connection is upgraded.
|
|
wsConn._socket = wsConn._socket || {};
|
|
wsConn._socket.remoteAddress = ip;
|
|
|
|
// Instantiate a new Client wrapper.
|
|
// Event listeners are automatically attached in the Client constructor.
|
|
new Client(wsConn, hub);
|
|
});
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Resolves the client's original IP address, supporting reverse proxy setups
|
|
* (like Nginx) that forward the real client IP via X-Forwarded-For header.
|
|
* @param {object} request The http.IncomingMessage request object.
|
|
* @returns {string} The resolved IP address.
|
|
*/
|
|
function getClientIP(request) {
|
|
const xff = request.headers['x-forwarded-for'];
|
|
if (xff) {
|
|
// If proxied multiple times, the first element is the client's real IP.
|
|
const ips = xff.split(',');
|
|
return ips[0].trim();
|
|
}
|
|
return request.socket.remoteAddress || 'unknown';
|
|
}
|