smartproxy/ts/proxies/http-proxy/websocket-handler.ts

504 lines
20 KiB
TypeScript

import * as plugins from '../../plugins.js';
import '../../core/models/socket-augmentation.js';
import { type IHttpProxyOptions, type IWebSocketWithHeartbeat, type ILogger, createLogger, type IReverseProxyConfig } from './models/types.js';
import { ConnectionPool } from './connection-pool.js';
import { ProxyRouter, RouteRouter } from '../../routing/router/index.js';
import type { IRouteConfig } from '../smart-proxy/models/route-types.js';
import type { IRouteContext } from '../../core/models/route-context.js';
import { toBaseContext } from '../../core/models/route-context.js';
import { ContextCreator } from './context-creator.js';
import { SecurityManager } from './security-manager.js';
import { TemplateUtils } from '../../core/utils/template-utils.js';
import { getMessageSize, toBuffer } from '../../core/utils/websocket-utils.js';
/**
* Handles WebSocket connections and proxying
*/
export class WebSocketHandler {
private heartbeatInterval: NodeJS.Timeout | null = null;
private wsServer: plugins.ws.WebSocketServer | null = null;
private logger: ILogger;
private contextCreator: ContextCreator = new ContextCreator();
private routeRouter: RouteRouter | null = null;
private securityManager: SecurityManager;
constructor(
private options: IHttpProxyOptions,
private connectionPool: ConnectionPool,
private legacyRouter: ProxyRouter, // Legacy router for backward compatibility
private routes: IRouteConfig[] = [] // Routes for modern router
) {
this.logger = createLogger(options.logLevel || 'info');
this.securityManager = new SecurityManager(this.logger, routes);
// Initialize modern router if we have routes
if (routes.length > 0) {
this.routeRouter = new RouteRouter(routes, this.logger);
}
}
/**
* Set the route configurations
*/
public setRoutes(routes: IRouteConfig[]): void {
this.routes = routes;
// Initialize or update the route router
if (!this.routeRouter) {
this.routeRouter = new RouteRouter(routes, this.logger);
} else {
this.routeRouter.setRoutes(routes);
}
// Update the security manager
this.securityManager.setRoutes(routes);
}
/**
* Initialize WebSocket server on an existing HTTPS server
*/
public initialize(server: plugins.https.Server): void {
// Create WebSocket server
this.wsServer = new plugins.ws.WebSocketServer({
server: server,
clientTracking: true
});
// Handle WebSocket connections
this.wsServer.on('connection', (wsIncoming: IWebSocketWithHeartbeat, req: plugins.http.IncomingMessage) => {
this.handleWebSocketConnection(wsIncoming, req);
});
// Start the heartbeat interval
this.startHeartbeat();
this.logger.info('WebSocket handler initialized');
}
/**
* Start the heartbeat interval to check for inactive WebSocket connections
*/
private startHeartbeat(): void {
// Clean up existing interval if any
if (this.heartbeatInterval) {
clearInterval(this.heartbeatInterval);
}
// Set up the heartbeat interval (check every 30 seconds)
this.heartbeatInterval = setInterval(() => {
if (!this.wsServer || this.wsServer.clients.size === 0) {
return; // Skip if no active connections
}
this.logger.debug(`WebSocket heartbeat check for ${this.wsServer.clients.size} clients`);
this.wsServer.clients.forEach((ws: plugins.wsDefault) => {
const wsWithHeartbeat = ws as IWebSocketWithHeartbeat;
if (wsWithHeartbeat.isAlive === false) {
this.logger.debug('Terminating inactive WebSocket connection');
return wsWithHeartbeat.terminate();
}
wsWithHeartbeat.isAlive = false;
wsWithHeartbeat.ping();
});
}, 30000);
// Make sure the interval doesn't keep the process alive
if (this.heartbeatInterval.unref) {
this.heartbeatInterval.unref();
}
}
/**
* Handle a new WebSocket connection
*/
private handleWebSocketConnection(wsIncoming: IWebSocketWithHeartbeat, req: plugins.http.IncomingMessage): void {
this.logger.debug(`WebSocket connection initiated from ${req.headers.host}`);
try {
// Initialize heartbeat tracking
wsIncoming.isAlive = true;
wsIncoming.lastPong = Date.now();
// Handle pong messages to track liveness
wsIncoming.on('pong', () => {
wsIncoming.isAlive = true;
wsIncoming.lastPong = Date.now();
});
// Create a context for routing
const connectionId = `ws-${Date.now()}-${Math.floor(Math.random() * 10000)}`;
const routeContext = this.contextCreator.createHttpRouteContext(req, {
connectionId,
clientIp: req.socket.remoteAddress?.replace('::ffff:', '') || '0.0.0.0',
serverIp: req.socket.localAddress?.replace('::ffff:', '') || '0.0.0.0',
tlsVersion: req.socket.getTLSVersion?.() || undefined
});
// Try modern router first if available
let route: IRouteConfig | undefined;
if (this.routeRouter) {
route = this.routeRouter.routeReq(req);
}
// Define destination variables
let destination: { host: string; port: number };
// If we found a route with the modern router, use it
if (route && route.action.type === 'forward' && route.action.target) {
this.logger.debug(`Found matching WebSocket route: ${route.name || 'unnamed'}`);
// Check if WebSockets are enabled for this route
if (route.action.websocket?.enabled === false) {
this.logger.debug(`WebSockets are disabled for route: ${route.name || 'unnamed'}`);
wsIncoming.close(1003, 'WebSockets not supported for this route');
return;
}
// Check security restrictions if configured to authenticate WebSocket requests
if (route.action.websocket?.authenticateRequest !== false && route.security) {
if (!this.securityManager.isAllowed(route, toBaseContext(routeContext))) {
this.logger.warn(`WebSocket connection denied by security policy for ${routeContext.clientIp}`);
wsIncoming.close(1008, 'Access denied by security policy');
return;
}
// Check origin restrictions if configured
const origin = req.headers.origin;
if (origin && route.action.websocket?.allowedOrigins && route.action.websocket.allowedOrigins.length > 0) {
const isAllowed = route.action.websocket.allowedOrigins.some(allowedOrigin => {
// Handle wildcards and template variables
if (allowedOrigin.includes('*') || allowedOrigin.includes('{')) {
const pattern = allowedOrigin.replace(/\*/g, '.*');
const resolvedPattern = TemplateUtils.resolveTemplateVariables(pattern, routeContext);
const regex = new RegExp(`^${resolvedPattern}$`);
return regex.test(origin);
}
return allowedOrigin === origin;
});
if (!isAllowed) {
this.logger.warn(`WebSocket origin ${origin} not allowed for route: ${route.name || 'unnamed'}`);
wsIncoming.close(1008, 'Origin not allowed');
return;
}
}
}
// Extract target information, resolving functions if needed
let targetHost: string | string[];
let targetPort: number;
try {
// Resolve host if it's a function
if (typeof route.action.target.host === 'function') {
const resolvedHost = route.action.target.host(toBaseContext(routeContext));
targetHost = resolvedHost;
this.logger.debug(`Resolved function-based host for WebSocket: ${Array.isArray(resolvedHost) ? resolvedHost.join(', ') : resolvedHost}`);
} else {
targetHost = route.action.target.host;
}
// Resolve port if it's a function
if (typeof route.action.target.port === 'function') {
targetPort = route.action.target.port(toBaseContext(routeContext));
this.logger.debug(`Resolved function-based port for WebSocket: ${targetPort}`);
} else {
targetPort = route.action.target.port === 'preserve' ? routeContext.port : route.action.target.port as number;
}
// Select a single host if an array was provided
const selectedHost = Array.isArray(targetHost)
? targetHost[Math.floor(Math.random() * targetHost.length)]
: targetHost;
// Create a destination for the WebSocket connection
destination = {
host: selectedHost,
port: targetPort
};
this.logger.debug(`WebSocket destination resolved: ${selectedHost}:${targetPort}`);
} catch (err) {
this.logger.error(`Error evaluating function-based target for WebSocket: ${err}`);
wsIncoming.close(1011, 'Internal server error');
return;
}
} else {
// Fall back to legacy routing if no matching route found via modern router
const proxyConfig = this.legacyRouter.routeReq(req);
if (!proxyConfig) {
this.logger.warn(`No proxy configuration for WebSocket host: ${req.headers.host}`);
wsIncoming.close(1008, 'No proxy configuration for this host');
return;
}
// Get destination target using round-robin if multiple targets
destination = this.connectionPool.getNextTarget(
proxyConfig.destinationIps,
proxyConfig.destinationPorts[0]
);
}
// Build target URL with potential path rewriting
// Determine protocol based on the target's configuration
// For WebSocket connections, we use ws for HTTP backends and wss for HTTPS backends
const isTargetSecure = destination.port === 443;
const protocol = isTargetSecure ? 'wss' : 'ws';
let targetPath = req.url || '/';
// Apply path rewriting if configured
if (route?.action.websocket?.rewritePath) {
const originalPath = targetPath;
targetPath = TemplateUtils.resolveTemplateVariables(
route.action.websocket.rewritePath,
{...routeContext, path: targetPath}
);
this.logger.debug(`WebSocket path rewritten: ${originalPath} -> ${targetPath}`);
}
const targetUrl = `${protocol}://${destination.host}:${destination.port}${targetPath}`;
this.logger.debug(`WebSocket connection from ${req.socket.remoteAddress} to ${targetUrl}`);
// Create headers for outgoing WebSocket connection
const headers: { [key: string]: string } = {};
// Copy relevant headers from incoming request
for (const [key, value] of Object.entries(req.headers)) {
if (value && typeof value === 'string' &&
key.toLowerCase() !== 'connection' &&
key.toLowerCase() !== 'upgrade' &&
key.toLowerCase() !== 'sec-websocket-key' &&
key.toLowerCase() !== 'sec-websocket-version') {
headers[key] = value;
}
}
// Always rewrite host header for WebSockets for consistency
headers['host'] = `${destination.host}:${destination.port}`;
// Add custom headers from route configuration
if (route?.action.websocket?.customHeaders) {
for (const [key, value] of Object.entries(route.action.websocket.customHeaders)) {
// Skip if header already exists and we're not overriding
if (headers[key.toLowerCase()] && !value.startsWith('!')) {
continue;
}
// Handle special delete directive (!delete)
if (value === '!delete') {
delete headers[key.toLowerCase()];
continue;
}
// Handle forced override (!value)
let finalValue: string;
if (value.startsWith('!') && value !== '!delete') {
// Keep the ! but resolve any templates in the rest
const templateValue = value.substring(1);
finalValue = '!' + TemplateUtils.resolveTemplateVariables(templateValue, routeContext);
} else {
// Resolve templates in the entire value
finalValue = TemplateUtils.resolveTemplateVariables(value, routeContext);
}
// Set the header
headers[key.toLowerCase()] = finalValue;
}
}
// Create WebSocket connection options
const wsOptions: any = {
headers: headers,
followRedirects: true
};
// Add subprotocols if configured
if (route?.action.websocket?.subprotocols && route.action.websocket.subprotocols.length > 0) {
wsOptions.protocols = route.action.websocket.subprotocols;
} else if (req.headers['sec-websocket-protocol']) {
// Pass through client requested protocols
wsOptions.protocols = req.headers['sec-websocket-protocol'].split(',').map(p => p.trim());
}
// Create outgoing WebSocket connection
this.logger.debug(`Creating WebSocket connection to ${targetUrl} with options:`, {
headers: wsOptions.headers,
protocols: wsOptions.protocols
});
const wsOutgoing = new plugins.wsDefault(targetUrl, wsOptions);
this.logger.debug(`WebSocket instance created, waiting for connection...`);
// Handle connection errors
wsOutgoing.on('error', (err) => {
this.logger.error(`WebSocket target connection error: ${err.message}`);
if (wsIncoming.readyState === wsIncoming.OPEN) {
wsIncoming.close(1011, 'Internal server error');
}
});
// Handle outgoing connection open
wsOutgoing.on('open', () => {
this.logger.debug(`WebSocket target connection opened to ${targetUrl}`);
// Set up custom ping interval if configured
let pingInterval: NodeJS.Timeout | null = null;
if (route?.action.websocket?.pingInterval && route.action.websocket.pingInterval > 0) {
pingInterval = setInterval(() => {
if (wsIncoming.readyState === wsIncoming.OPEN) {
wsIncoming.ping();
this.logger.debug(`Sent WebSocket ping to client for route: ${route.name || 'unnamed'}`);
}
}, route.action.websocket.pingInterval);
// Don't keep process alive just for pings
if (pingInterval.unref) pingInterval.unref();
}
// Set up custom ping timeout if configured
let pingTimeout: NodeJS.Timeout | null = null;
const pingTimeoutMs = route?.action.websocket?.pingTimeout || 60000; // Default 60s
// Define timeout function for cleaner code
const resetPingTimeout = () => {
if (pingTimeout) clearTimeout(pingTimeout);
pingTimeout = setTimeout(() => {
this.logger.debug(`WebSocket ping timeout for client connection on route: ${route?.name || 'unnamed'}`);
wsIncoming.terminate();
}, pingTimeoutMs);
// Don't keep process alive just for timeouts
if (pingTimeout.unref) pingTimeout.unref();
};
// Reset timeout on pong
wsIncoming.on('pong', () => {
wsIncoming.isAlive = true;
wsIncoming.lastPong = Date.now();
resetPingTimeout();
});
// Initial ping timeout
resetPingTimeout();
// Handle potential message size limits
const maxSize = route?.action.websocket?.maxPayloadSize || 0;
// Forward incoming messages to outgoing connection
wsIncoming.on('message', (data, isBinary) => {
this.logger.debug(`WebSocket forwarding message from client to target: ${data.toString()}`);
if (wsOutgoing.readyState === wsOutgoing.OPEN) {
// Check message size if limit is set
const messageSize = getMessageSize(data);
if (maxSize > 0 && messageSize > maxSize) {
this.logger.warn(`WebSocket message exceeds max size (${messageSize} > ${maxSize})`);
wsIncoming.close(1009, 'Message too big');
return;
}
wsOutgoing.send(data, { binary: isBinary });
} else {
this.logger.warn(`WebSocket target connection not open (state: ${wsOutgoing.readyState})`);
}
});
// Forward outgoing messages to incoming connection
wsOutgoing.on('message', (data, isBinary) => {
this.logger.debug(`WebSocket forwarding message from target to client: ${data.toString()}`);
if (wsIncoming.readyState === wsIncoming.OPEN) {
wsIncoming.send(data, { binary: isBinary });
} else {
this.logger.warn(`WebSocket client connection not open (state: ${wsIncoming.readyState})`);
}
});
// Handle closing of connections
wsIncoming.on('close', (code, reason) => {
this.logger.debug(`WebSocket client connection closed: ${code} ${reason}`);
if (wsOutgoing.readyState === wsOutgoing.OPEN) {
// Ensure code is a valid WebSocket close code number
const validCode = typeof code === 'number' && code >= 1000 && code <= 4999 ? code : 1000;
try {
const reasonString = reason ? toBuffer(reason).toString() : '';
wsOutgoing.close(validCode, reasonString);
} catch (err) {
this.logger.error('Error closing wsOutgoing:', err);
wsOutgoing.close(validCode);
}
}
// Clean up timers
if (pingInterval) clearInterval(pingInterval);
if (pingTimeout) clearTimeout(pingTimeout);
});
wsOutgoing.on('close', (code, reason) => {
this.logger.debug(`WebSocket target connection closed: ${code} ${reason}`);
if (wsIncoming.readyState === wsIncoming.OPEN) {
// Ensure code is a valid WebSocket close code number
const validCode = typeof code === 'number' && code >= 1000 && code <= 4999 ? code : 1000;
try {
const reasonString = reason ? toBuffer(reason).toString() : '';
wsIncoming.close(validCode, reasonString);
} catch (err) {
this.logger.error('Error closing wsIncoming:', err);
wsIncoming.close(validCode);
}
}
// Clean up timers
if (pingInterval) clearInterval(pingInterval);
if (pingTimeout) clearTimeout(pingTimeout);
});
this.logger.debug(`WebSocket connection established: ${req.headers.host} -> ${destination.host}:${destination.port}`);
});
} catch (error) {
this.logger.error(`Error handling WebSocket connection: ${error.message}`);
if (wsIncoming.readyState === wsIncoming.OPEN) {
wsIncoming.close(1011, 'Internal server error');
}
}
}
/**
* Get information about active WebSocket connections
*/
public getConnectionInfo(): { activeConnections: number } {
return {
activeConnections: this.wsServer ? this.wsServer.clients.size : 0
};
}
/**
* Shutdown the WebSocket handler
*/
public shutdown(): void {
// Stop heartbeat interval
if (this.heartbeatInterval) {
clearInterval(this.heartbeatInterval);
this.heartbeatInterval = null;
}
// Close all WebSocket connections
if (this.wsServer) {
this.logger.info(`Closing ${this.wsServer.clients.size} WebSocket connections`);
for (const client of this.wsServer.clients) {
try {
client.terminate();
} catch (error) {
this.logger.error('Error terminating WebSocket client', error);
}
}
// Close the server
this.wsServer.close();
this.wsServer = null;
}
}
}