Harden client and WebSocket proxy
This commit is contained in:
+267
-465
@@ -1,491 +1,293 @@
|
||||
import { WebSocketServer } from 'ws';
|
||||
import * as net from 'net';
|
||||
import * as tls from 'tls';
|
||||
import http from 'http';
|
||||
import { parse } from 'url';
|
||||
import { randomBytes } from 'node:crypto';
|
||||
import { lookup } from 'node:dns/promises';
|
||||
import http from 'node:http';
|
||||
import net, { BlockList } from 'node:net';
|
||||
import tls from 'node:tls';
|
||||
import { WebSocket, WebSocketServer } from 'ws';
|
||||
|
||||
// Default configuration for connection persistence (fallback values)
|
||||
const DEFAULT_PERSISTENCE_TIMEOUT = 5 * 60 * 1000; // 5 minutes in milliseconds
|
||||
const DEFAULT_MAX_BUFFER_MESSAGES = 100; // Maximum number of messages to buffer
|
||||
const DEFAULT_MAX_BUFFER_SIZE_KB = 10; // Maximum buffer size in KB
|
||||
const PORT = numberFromEnv('WS_PORT', 3001, 1, 65535);
|
||||
const MAX_SESSIONS_PER_IP = numberFromEnv('PROXY_MAX_SESSIONS_PER_IP', 3, 1, 20);
|
||||
const MAX_GLOBAL_SESSIONS = numberFromEnv('PROXY_MAX_GLOBAL_SESSIONS', 200, 1, 10_000);
|
||||
const NEW_CONNECTION_LIMIT = numberFromEnv('PROXY_NEW_CONNECTIONS_PER_10_MIN', 20, 1, 1_000);
|
||||
const CONNECT_TIMEOUT_MS = 10_000;
|
||||
const PERSISTENCE_TIMEOUT_MS = 5 * 60_000;
|
||||
const MAX_BUFFER_MESSAGES = 250;
|
||||
const MAX_BUFFER_BYTES = 256 * 1024;
|
||||
const MAX_CONTROL_BYTES = 4 * 1024;
|
||||
const MAX_BYTES_PER_SECOND = 2 * 1024 * 1024;
|
||||
const DENIED_PORTS = new Set([25, 465, 587, 2525]);
|
||||
const production = process.env.NODE_ENV === 'production';
|
||||
const allowedOrigins = new Set(
|
||||
(process.env.ALLOWED_ORIGINS || (production ? '' : 'http://localhost:5173,http://127.0.0.1:5173'))
|
||||
.split(',').map((origin) => origin.trim()).filter(Boolean)
|
||||
);
|
||||
const deniedAddresses = createDeniedAddressList();
|
||||
const sessions = new Map();
|
||||
const connectionAttempts = new Map();
|
||||
|
||||
const HEARTBEAT_INTERVAL = 30 * 1000; // 30 seconds
|
||||
|
||||
// Create HTTP server
|
||||
const server = http.createServer();
|
||||
|
||||
// Create WebSocket server
|
||||
const wss = new WebSocketServer({ noServer: true });
|
||||
|
||||
// Active connections and their proxies
|
||||
// Key: connectionId, Value: { ws, socket, sessionId, settings }
|
||||
const connections = new Map();
|
||||
|
||||
// Persistent connections waiting for reconnection
|
||||
// Key: sessionId, Value: { socket, mudHost, mudPort, useSSL, timeoutId, lastActivity, messageBuffer, settings }
|
||||
const persistentConnections = new Map();
|
||||
|
||||
// Parse connection settings from URL parameters with defaults
|
||||
function parseConnectionSettings(urlParts) {
|
||||
const persistenceTimeoutParam = urlParts.searchParams.get('persistenceTimeout');
|
||||
const maxBufferMessagesParam = urlParts.searchParams.get('maxBufferMessages');
|
||||
const maxBufferSizeKBParam = urlParts.searchParams.get('maxBufferSizeKB');
|
||||
|
||||
return {
|
||||
persistenceTimeoutMs: persistenceTimeoutParam ?
|
||||
parseInt(persistenceTimeoutParam) * 60 * 1000 : // Convert minutes to milliseconds
|
||||
DEFAULT_PERSISTENCE_TIMEOUT,
|
||||
maxBufferMessages: maxBufferMessagesParam ?
|
||||
parseInt(maxBufferMessagesParam) :
|
||||
DEFAULT_MAX_BUFFER_MESSAGES,
|
||||
maxBufferSizeKB: maxBufferSizeKBParam ?
|
||||
parseInt(maxBufferSizeKBParam) :
|
||||
DEFAULT_MAX_BUFFER_SIZE_KB
|
||||
};
|
||||
}
|
||||
|
||||
// Generate a unique session ID for persistent connections
|
||||
function generateSessionId() {
|
||||
return `session-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
|
||||
}
|
||||
|
||||
// Buffer a message for a persistent connection
|
||||
function bufferMessage(sessionId, data) {
|
||||
const persistentConn = persistentConnections.get(sessionId);
|
||||
if (!persistentConn) {
|
||||
return; // No persistent connection to buffer for
|
||||
}
|
||||
|
||||
if (!persistentConn.messageBuffer) {
|
||||
persistentConn.messageBuffer = [];
|
||||
persistentConn.bufferSize = 0;
|
||||
}
|
||||
|
||||
// Add timestamp to the message
|
||||
const bufferedMessage = {
|
||||
data: data,
|
||||
timestamp: Date.now()
|
||||
};
|
||||
|
||||
persistentConn.messageBuffer.push(bufferedMessage);
|
||||
persistentConn.bufferSize += data.length;
|
||||
|
||||
// Use this connection's specific settings for buffer limits
|
||||
const settings = persistentConn.settings || {
|
||||
maxBufferMessages: DEFAULT_MAX_BUFFER_MESSAGES,
|
||||
maxBufferSizeKB: DEFAULT_MAX_BUFFER_SIZE_KB
|
||||
};
|
||||
|
||||
// Trim buffer if it gets too large
|
||||
while (persistentConn.messageBuffer.length > settings.maxBufferMessages ||
|
||||
persistentConn.bufferSize > settings.maxBufferSizeKB * 1000) {
|
||||
const removed = persistentConn.messageBuffer.shift();
|
||||
if (removed) {
|
||||
persistentConn.bufferSize -= removed.data.length;
|
||||
}
|
||||
}
|
||||
|
||||
console.log(`Buffered ${data.length} bytes for session ${sessionId} (${persistentConn.messageBuffer.length} messages, ${persistentConn.bufferSize} bytes total, limits: ${settings.maxBufferMessages} msgs/${settings.maxBufferSizeKB}KB)`);
|
||||
}
|
||||
|
||||
// Replay buffered messages to a reconnected client
|
||||
function replayBufferedMessages(ws, sessionId) {
|
||||
const persistentConn = persistentConnections.get(sessionId);
|
||||
if (!persistentConn || !persistentConn.messageBuffer) {
|
||||
return 0; // No messages to replay
|
||||
}
|
||||
|
||||
const messages = persistentConn.messageBuffer;
|
||||
console.log(`Replaying ${messages.length} buffered messages for session ${sessionId}`);
|
||||
|
||||
// Send a notification about message replay
|
||||
const replayNotification = `[SYSTEM]${JSON.stringify({
|
||||
type: 'message_replay_start',
|
||||
messageCount: messages.length,
|
||||
timespan: messages.length > 0 ? Date.now() - messages[0].timestamp : 0
|
||||
})}`;
|
||||
ws.send(replayNotification);
|
||||
|
||||
// Send all buffered messages
|
||||
for (const message of messages) {
|
||||
if (ws.readyState === 1) { // WebSocket.OPEN
|
||||
ws.send(message.data);
|
||||
}
|
||||
}
|
||||
|
||||
// Send replay complete notification
|
||||
const replayComplete = `[SYSTEM]${JSON.stringify({ type: 'message_replay_complete' })}`;
|
||||
ws.send(replayComplete);
|
||||
|
||||
// Clear the buffer since messages have been replayed
|
||||
const messageCount = messages.length;
|
||||
persistentConn.messageBuffer = [];
|
||||
persistentConn.bufferSize = 0;
|
||||
|
||||
return messageCount;
|
||||
}
|
||||
|
||||
// Clean up a persistent connection
|
||||
function cleanupPersistentConnection(sessionId) {
|
||||
const persistentConn = persistentConnections.get(sessionId);
|
||||
if (persistentConn) {
|
||||
console.log(`Cleaning up persistent connection for session ${sessionId}`);
|
||||
|
||||
// Clear timeout
|
||||
if (persistentConn.timeoutId) {
|
||||
clearTimeout(persistentConn.timeoutId);
|
||||
}
|
||||
|
||||
// Close MUD socket
|
||||
if (persistentConn.socket && !persistentConn.socket.destroyed) {
|
||||
persistentConn.socket.end();
|
||||
}
|
||||
|
||||
// Remove from map
|
||||
persistentConnections.delete(sessionId);
|
||||
}
|
||||
}
|
||||
|
||||
// Handle WebSocket connections
|
||||
wss.on('connection', (ws, req, mudHost, mudPort, useSSL) => {
|
||||
console.log(`WebSocket connection established for ${mudHost}:${mudPort} (SSL: ${useSSL})`);
|
||||
|
||||
// Create a unique ID for this connection
|
||||
const connectionId = `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
|
||||
|
||||
// Check for session ID and settings in query parameters
|
||||
const url = req.url || '';
|
||||
const urlParts = new URL(`http://localhost${url}`);
|
||||
const sessionId = urlParts.searchParams.get('sessionId');
|
||||
|
||||
// Parse connection settings for this specific connection
|
||||
const connectionSettings = parseConnectionSettings(urlParts);
|
||||
|
||||
console.log(`Connection settings for ${connectionId}: timeout=${connectionSettings.persistenceTimeoutMs/1000/60}min, maxMessages=${connectionSettings.maxBufferMessages}, maxSizeKB=${connectionSettings.maxBufferSizeKB}`);
|
||||
|
||||
// Special handling for test connections
|
||||
if (mudHost === 'example.com' && mudPort === '23') {
|
||||
console.log('Test connection detected - using echo server mode');
|
||||
|
||||
// Send welcome message
|
||||
ws.send('Hello from WebSocket test server! This is an echo server.');
|
||||
|
||||
// Echo back messages
|
||||
ws.on('message', (message) => {
|
||||
console.log('Test server received:', message.toString());
|
||||
ws.send(`Echo: ${message.toString()}`);
|
||||
});
|
||||
|
||||
// Handle close
|
||||
ws.on('close', () => {
|
||||
console.log('Test connection closed');
|
||||
connections.delete(connectionId);
|
||||
});
|
||||
|
||||
// Store the connection (without a socket)
|
||||
connections.set(connectionId, { ws, testMode: true });
|
||||
const server = http.createServer((request, response) => {
|
||||
if (request.url === '/health') {
|
||||
response.writeHead(200, { 'content-type': 'application/json' });
|
||||
response.end(JSON.stringify({ ok: true, sessions: sessions.size }));
|
||||
return;
|
||||
}
|
||||
response.writeHead(404);
|
||||
response.end();
|
||||
});
|
||||
const wss = new WebSocketServer({ noServer: true, maxPayload: 64 * 1024, perMessageDeflate: false });
|
||||
|
||||
let socket;
|
||||
let currentSessionId = sessionId;
|
||||
|
||||
// Check if this is a reconnection to an existing persistent session
|
||||
if (sessionId && persistentConnections.has(sessionId)) {
|
||||
console.log(`Reconnecting to existing session: ${sessionId}`);
|
||||
|
||||
const persistentConn = persistentConnections.get(sessionId);
|
||||
socket = persistentConn.socket;
|
||||
|
||||
// Clear the timeout since client reconnected
|
||||
if (persistentConn.timeoutId) {
|
||||
clearTimeout(persistentConn.timeoutId);
|
||||
}
|
||||
|
||||
// Replay any buffered messages first
|
||||
const replayedCount = replayBufferedMessages(ws, sessionId);
|
||||
|
||||
// Remove from persistent connections (now active again) - do this after replay
|
||||
persistentConnections.delete(sessionId);
|
||||
|
||||
// Send reconnection notification with session ID in proper JSON format
|
||||
ws.send(`[SYSTEM]${JSON.stringify({
|
||||
type: 'session_resumed',
|
||||
sessionId: sessionId,
|
||||
messagesReplayed: replayedCount
|
||||
})}`);
|
||||
} else {
|
||||
// Create new connection
|
||||
currentSessionId = generateSessionId();
|
||||
console.log(`Creating new session: ${currentSessionId}`);
|
||||
|
||||
try {
|
||||
// Create a TCP socket connection to the MUD server
|
||||
// Use tls for SSL connections, net for regular connections
|
||||
socket = useSSL
|
||||
? tls.connect({ host: mudHost, port: parseInt(mudPort), rejectUnauthorized: false })
|
||||
: net.createConnection({ host: mudHost, port: parseInt(mudPort) });
|
||||
|
||||
// Add error handler
|
||||
socket.on('error', (error) => {
|
||||
console.error(`Socket error for ${mudHost}:${mudPort}:`, error.message);
|
||||
// Send error to client
|
||||
if (ws.readyState === 1) {
|
||||
ws.send(Buffer.from(`ERROR: Connection to MUD server failed: ${error.message}\r\n`));
|
||||
setTimeout(() => {
|
||||
if (ws.readyState === 1) ws.close();
|
||||
}, 1000);
|
||||
}
|
||||
// Remove from connections map
|
||||
connections.delete(connectionId);
|
||||
});
|
||||
|
||||
// Send session ID to client in proper JSON format
|
||||
ws.send(`[SYSTEM]${JSON.stringify({ sessionId: currentSessionId })}`);
|
||||
|
||||
} catch (error) {
|
||||
console.error(`Error creating socket connection: ${error.message}`);
|
||||
if (ws.readyState === 1) {
|
||||
ws.send(Buffer.from(`ERROR: Failed to connect to MUD server: ${error.message}\r\n`));
|
||||
ws.close();
|
||||
}
|
||||
return;
|
||||
}
|
||||
function numberFromEnv(name, fallback, minimum, maximum) {
|
||||
const parsed = Number.parseInt(process.env[name] || '', 10);
|
||||
return Number.isFinite(parsed) && parsed >= minimum && parsed <= maximum ? parsed : fallback;
|
||||
}
|
||||
|
||||
function createDeniedAddressList() {
|
||||
const list = new BlockList();
|
||||
for (const [address, prefix] of [
|
||||
['0.0.0.0', 8], ['10.0.0.0', 8], ['100.64.0.0', 10], ['127.0.0.0', 8],
|
||||
['169.254.0.0', 16], ['172.16.0.0', 12], ['192.0.0.0', 24], ['192.0.2.0', 24],
|
||||
['192.168.0.0', 16], ['198.18.0.0', 15], ['198.51.100.0', 24], ['203.0.113.0', 24],
|
||||
['224.0.0.0', 4], ['240.0.0.0', 4]
|
||||
]) list.addSubnet(address, prefix, 'ipv4');
|
||||
for (const [address, prefix] of [
|
||||
['::', 128], ['::1', 128], ['64:ff9b::', 96], ['2001::', 32],
|
||||
['2001:db8::', 32], ['2002::', 16], ['fc00::', 7], ['fe80::', 10], ['ff00::', 8]
|
||||
]) list.addSubnet(address, prefix, 'ipv6');
|
||||
return list;
|
||||
}
|
||||
|
||||
function getClientIp(request) {
|
||||
if (process.env.TRUST_PROXY === '1') {
|
||||
const forwarded = request.headers['x-forwarded-for'];
|
||||
if (typeof forwarded === 'string' && forwarded.length > 0) return forwarded.split(',')[0].trim();
|
||||
}
|
||||
return request.socket.remoteAddress || 'unknown';
|
||||
}
|
||||
|
||||
// Store the connection with its settings
|
||||
connections.set(connectionId, {
|
||||
ws,
|
||||
socket,
|
||||
sessionId: currentSessionId,
|
||||
settings: connectionSettings
|
||||
function sendControl(ws, payload) {
|
||||
if (ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify(payload));
|
||||
}
|
||||
|
||||
function failUpgrade(socket, status, message) {
|
||||
socket.write(`HTTP/1.1 ${status}\r\nConnection: close\r\nContent-Type: text/plain\r\n\r\n${message}`);
|
||||
socket.destroy();
|
||||
}
|
||||
|
||||
function countSessionsForIp(clientIp) {
|
||||
let count = 0;
|
||||
for (const session of sessions.values()) if (session.clientIp === clientIp) count += 1;
|
||||
return count;
|
||||
}
|
||||
|
||||
function recordConnectionAttempt(clientIp) {
|
||||
const cutoff = Date.now() - 10 * 60_000;
|
||||
const recent = (connectionAttempts.get(clientIp) || []).filter((timestamp) => timestamp > cutoff);
|
||||
if (recent.length >= NEW_CONNECTION_LIMIT) return false;
|
||||
recent.push(Date.now());
|
||||
connectionAttempts.set(clientIp, recent);
|
||||
return true;
|
||||
}
|
||||
|
||||
function validateConnectMessage(value) {
|
||||
if (!value || typeof value !== 'object' || value.type !== 'connect') throw new Error('The first frame must be a connect control message.');
|
||||
if (typeof value.host !== 'string' || value.host.length < 1 || value.host.length > 253) throw new Error('Invalid host.');
|
||||
if (!/^[a-zA-Z0-9._:-]+$/.test(value.host)) throw new Error('Invalid host characters.');
|
||||
if (!Number.isInteger(value.port) || value.port < 1 || value.port > 65535 || DENIED_PORTS.has(value.port)) throw new Error('Invalid or denied port.');
|
||||
if (typeof value.tls !== 'boolean') throw new Error('Invalid TLS setting.');
|
||||
if (value.resumeToken !== undefined && (typeof value.resumeToken !== 'string' || value.resumeToken.length > 128)) throw new Error('Invalid resume token.');
|
||||
return { host: value.host, port: value.port, useTls: value.tls, resumeToken: value.resumeToken };
|
||||
}
|
||||
|
||||
async function resolvePublicTarget(host) {
|
||||
const directFamily = net.isIP(host);
|
||||
const answers = directFamily ? [{ address: host, family: directFamily }] : await lookup(host, { all: true, verbatim: true });
|
||||
if (answers.length === 0) throw new Error('Host did not resolve.');
|
||||
for (const answer of answers) {
|
||||
if (isDeniedAddress(answer.address, answer.family)) throw new Error('Target resolves to a non-public address.');
|
||||
}
|
||||
return answers[0];
|
||||
}
|
||||
|
||||
function isDeniedAddress(address, family = net.isIP(address)) {
|
||||
const mapped = address.match(/^::ffff:(?:(\d+\.\d+\.\d+\.\d+)|([0-9a-f]+):([0-9a-f]+))$/i);
|
||||
if (mapped) {
|
||||
let ipv4 = mapped[1];
|
||||
if (!ipv4) {
|
||||
const high = Number.parseInt(mapped[2], 16);
|
||||
const low = Number.parseInt(mapped[3], 16);
|
||||
ipv4 = `${high >> 8}.${high & 255}.${low >> 8}.${low & 255}`;
|
||||
}
|
||||
return deniedAddresses.check(ipv4, 'ipv4');
|
||||
}
|
||||
return deniedAddresses.check(address, family === 4 ? 'ipv4' : 'ipv6');
|
||||
}
|
||||
|
||||
function consumeTraffic(session, direction, bytes) {
|
||||
if (Date.now() - session.traffic.startedAt >= 1_000) session.traffic = { startedAt: Date.now(), inbound: 0, outbound: 0 };
|
||||
session.traffic[direction] += bytes;
|
||||
return session.traffic[direction] <= MAX_BYTES_PER_SECOND;
|
||||
}
|
||||
|
||||
function bufferMessage(session, data) {
|
||||
session.buffer.push(Buffer.from(data));
|
||||
session.bufferBytes += data.length;
|
||||
while (session.buffer.length > MAX_BUFFER_MESSAGES || session.bufferBytes > MAX_BUFFER_BYTES) {
|
||||
const removed = session.buffer.shift();
|
||||
if (removed) session.bufferBytes -= removed.length;
|
||||
}
|
||||
}
|
||||
|
||||
function destroySession(session, reason = 'closed') {
|
||||
if (session.closed) return;
|
||||
session.closed = true;
|
||||
if (session.persistenceTimer) clearTimeout(session.persistenceTimer);
|
||||
sessions.delete(session.token);
|
||||
const attached = session.ws;
|
||||
session.ws = null;
|
||||
if (attached) sendControl(attached, { type: 'upstream_closed', reason });
|
||||
if (!session.socket.destroyed) session.socket.destroy();
|
||||
}
|
||||
|
||||
function rotateSessionToken(session) {
|
||||
sessions.delete(session.token);
|
||||
session.token = randomBytes(32).toString('base64url');
|
||||
sessions.set(session.token, session);
|
||||
}
|
||||
|
||||
function attachWebSocket(session, ws, resumed) {
|
||||
if (session.persistenceTimer) clearTimeout(session.persistenceTimer);
|
||||
session.persistenceTimer = null;
|
||||
session.ws = ws;
|
||||
rotateSessionToken(session);
|
||||
if (resumed) {
|
||||
sendControl(ws, { type: 'replay_started', messageCount: session.buffer.length });
|
||||
for (const message of session.buffer) if (ws.readyState === WebSocket.OPEN) ws.send(message);
|
||||
const messagesReplayed = session.buffer.length;
|
||||
session.buffer = [];
|
||||
session.bufferBytes = 0;
|
||||
sendControl(ws, { type: 'replay_finished', messagesReplayed });
|
||||
}
|
||||
sendControl(ws, { type: resumed ? 'session_resumed' : 'session_started', resumeToken: session.token });
|
||||
}
|
||||
|
||||
async function createSession(ws, request, target) {
|
||||
const clientIp = getClientIp(request);
|
||||
const origin = request.headers.origin || '';
|
||||
if (sessions.size >= MAX_GLOBAL_SESSIONS) throw new Error('Proxy capacity reached.');
|
||||
if (countSessionsForIp(clientIp) >= MAX_SESSIONS_PER_IP) throw new Error('Per-client connection limit reached.');
|
||||
if (!recordConnectionAttempt(clientIp)) throw new Error('Connection rate limit reached.');
|
||||
const resolved = await resolvePublicTarget(target.host);
|
||||
const socketOptions = { host: resolved.address, port: target.port, family: resolved.family,
|
||||
...(target.useTls ? { servername: net.isIP(target.host) ? undefined : target.host, rejectUnauthorized: true } : {}) };
|
||||
const socket = target.useTls ? tls.connect(socketOptions) : net.createConnection(socketOptions);
|
||||
const session = { token: randomBytes(32).toString('base64url'), clientIp, origin, target, socket, ws: null,
|
||||
buffer: [], bufferBytes: 0, persistenceTimer: null, closed: false,
|
||||
traffic: { startedAt: Date.now(), inbound: 0, outbound: 0 } };
|
||||
sessions.set(session.token, session);
|
||||
const connected = new Promise((resolve, reject) => {
|
||||
const connectedEvent = target.useTls ? 'secureConnect' : 'connect';
|
||||
const fail = (error) => reject(error instanceof Error ? error : new Error('Upstream connection failed.'));
|
||||
socket.once(connectedEvent, resolve);
|
||||
socket.once('error', fail);
|
||||
socket.once('close', () => fail(new Error('Upstream closed before connecting.')));
|
||||
socket.setTimeout(CONNECT_TIMEOUT_MS, () => {
|
||||
fail(new Error('Upstream connection timed out.'));
|
||||
destroySession(session, 'timeout');
|
||||
});
|
||||
});
|
||||
|
||||
// Handle data from the MUD server - only in regular mode, not test mode
|
||||
if (socket) {
|
||||
socket.on('data', (data) => {
|
||||
// Check for GMCP data (IAC SB GMCP) - very basic check for debugging
|
||||
// IAC = 255, SB = 250, GMCP = 201
|
||||
let isGmcp = false;
|
||||
for (let i = 0; i < data.length - 2; i++) {
|
||||
if (data[i] === 255 && data[i+1] === 250 && data[i+2] === 201) {
|
||||
isGmcp = true;
|
||||
console.log('WebSocket server: Detected GMCP data in server response');
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Forward data to the WebSocket client if it's still open
|
||||
if (ws.readyState === 1) { // WebSocket.OPEN
|
||||
ws.send(data);
|
||||
console.log(`WebSocket server: Sent ${data.length} bytes to client${isGmcp ? ' (contains GMCP data)' : ''}`);
|
||||
} else {
|
||||
// WebSocket is not open, buffer the message if we have a session
|
||||
if (currentSessionId) {
|
||||
bufferMessage(currentSessionId, data);
|
||||
}
|
||||
}
|
||||
});
|
||||
socket.on('data', (data) => {
|
||||
if (!consumeTraffic(session, 'inbound', data.length)) return destroySession(session, 'traffic_limit');
|
||||
if (session.ws?.readyState === WebSocket.OPEN) session.ws.send(data);
|
||||
else bufferMessage(session, data);
|
||||
});
|
||||
socket.on('error', (error) => { if (session.ws) sendControl(session.ws, { type: 'error', code: 'UPSTREAM_ERROR', message: error.message }); });
|
||||
socket.on('close', () => destroySession(session, 'upstream_closed'));
|
||||
socket.on('drain', () => session.ws?.resume());
|
||||
try {
|
||||
await connected;
|
||||
} catch (error) {
|
||||
destroySession(session, 'connect_failed');
|
||||
throw error;
|
||||
}
|
||||
|
||||
// Handle socket close from MUD server - this should trigger cleanup
|
||||
if (socket) {
|
||||
socket.on('close', () => {
|
||||
console.log(`MUD connection closed by server for ${mudHost}:${mudPort}`);
|
||||
// Close WebSocket if it's still open
|
||||
if (ws.readyState === 1) {
|
||||
ws.close();
|
||||
}
|
||||
// Remove from connections map
|
||||
connections.delete(connectionId);
|
||||
|
||||
// Also cleanup any persistent connection
|
||||
if (currentSessionId) {
|
||||
cleanupPersistentConnection(currentSessionId);
|
||||
}
|
||||
});
|
||||
socket.setTimeout(0);
|
||||
if (ws.readyState !== WebSocket.OPEN) {
|
||||
destroySession(session, 'client_closed_during_connect');
|
||||
throw new Error('Client closed while the upstream connection was opening.');
|
||||
}
|
||||
attachWebSocket(session, ws, false);
|
||||
return session;
|
||||
}
|
||||
|
||||
// Handle WebSocket messages (data from client to server)
|
||||
ws.on('message', (message) => {
|
||||
function tryResume(ws, request, target) {
|
||||
if (!target.resumeToken) return null;
|
||||
const session = sessions.get(target.resumeToken);
|
||||
const clientIp = getClientIp(request);
|
||||
const origin = request.headers.origin || '';
|
||||
if (!session || session.ws || session.clientIp !== clientIp || session.origin !== origin || session.target.host !== target.host ||
|
||||
session.target.port !== target.port || session.target.useTls !== target.useTls) return null;
|
||||
attachWebSocket(session, ws, true);
|
||||
return session;
|
||||
}
|
||||
|
||||
wss.on('connection', (ws, request) => {
|
||||
let session = null;
|
||||
let explicitDisconnect = false;
|
||||
let initialized = false;
|
||||
const initializationTimer = setTimeout(() => { if (!initialized) ws.close(1008, 'Connect control timeout'); }, 5_000);
|
||||
ws.on('message', async (data, isBinary) => {
|
||||
try {
|
||||
// Skip if this is a test connection (already handled in the test mode section)
|
||||
const conn = connections.get(connectionId);
|
||||
if (conn && conn.testMode) return;
|
||||
|
||||
// Check for system messages
|
||||
const messageStr = message.toString();
|
||||
if (messageStr.startsWith('[SYSTEM]')) {
|
||||
try {
|
||||
const jsonStr = messageStr.substring(8); // Remove "[SYSTEM]"
|
||||
const systemData = JSON.parse(jsonStr);
|
||||
|
||||
if (systemData.type === 'explicit_disconnect') {
|
||||
console.log(`Received explicit disconnect command for session ${currentSessionId}`);
|
||||
// This is an explicit disconnect - don't persist the connection
|
||||
if (socket && socket.writable) {
|
||||
socket.end();
|
||||
}
|
||||
if (ws.readyState === 1) {
|
||||
ws.close();
|
||||
}
|
||||
connections.delete(connectionId);
|
||||
if (currentSessionId) {
|
||||
cleanupPersistentConnection(currentSessionId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Error parsing system message:', error);
|
||||
}
|
||||
// Don't forward system messages to the MUD server
|
||||
if (!initialized) {
|
||||
if (isBinary || data.length > MAX_CONTROL_BYTES) throw new Error('Invalid connect control frame.');
|
||||
const target = validateConnectMessage(JSON.parse(data.toString('utf8')));
|
||||
initialized = true;
|
||||
clearTimeout(initializationTimer);
|
||||
session = tryResume(ws, request, target) || await createSession(ws, request, target);
|
||||
return;
|
||||
}
|
||||
|
||||
// Legacy support for old disconnect command
|
||||
if (messageStr.trim() === '[DISCONNECT]') {
|
||||
console.log(`Received legacy disconnect command for session ${currentSessionId}`);
|
||||
// This is an explicit disconnect - don't persist the connection
|
||||
if (socket && socket.writable) {
|
||||
socket.end();
|
||||
}
|
||||
if (ws.readyState === 1) {
|
||||
ws.close();
|
||||
}
|
||||
connections.delete(connectionId);
|
||||
if (currentSessionId) {
|
||||
cleanupPersistentConnection(currentSessionId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Check for GMCP data (IAC SB GMCP) in client messages
|
||||
let isGmcp = false;
|
||||
if (message instanceof Buffer || message instanceof Uint8Array) {
|
||||
for (let i = 0; i < message.length - 2; i++) {
|
||||
if (message[i] === 255 && message[i+1] === 250 && message[i+2] === 201) {
|
||||
isGmcp = true;
|
||||
console.log('WebSocket server: Detected GMCP data in client message');
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Forward data to the MUD server
|
||||
// The message might be Buffer, ArrayBuffer, or string
|
||||
if (conn && conn.socket && conn.socket.writable) {
|
||||
conn.socket.write(message);
|
||||
console.log(`WebSocket server: Sent ${message.length} bytes to MUD server${isGmcp ? ' (contains GMCP data)' : ''}`);
|
||||
} else {
|
||||
console.error('Socket not writable, cannot send data to MUD server');
|
||||
if (ws.readyState === 1) { // WebSocket.OPEN
|
||||
ws.send(Buffer.from(`ERROR: Cannot send data to MUD server: Socket not connected\r\n`));
|
||||
if (!isBinary) {
|
||||
if (data.length > MAX_CONTROL_BYTES) throw new Error('Control frame too large.');
|
||||
const control = JSON.parse(data.toString('utf8'));
|
||||
if (control.type === 'disconnect') {
|
||||
explicitDisconnect = true;
|
||||
if (session) destroySession(session, 'client_disconnect');
|
||||
ws.close(1000, 'Disconnected');
|
||||
return;
|
||||
}
|
||||
throw new Error('Unknown control frame.');
|
||||
}
|
||||
if (!session || session.ws !== ws || !session.socket.writable) throw new Error('Upstream connection is not writable.');
|
||||
if (!consumeTraffic(session, 'outbound', data.length)) return destroySession(session, 'traffic_limit');
|
||||
if (!session.socket.write(data)) ws.pause();
|
||||
} catch (error) {
|
||||
console.error('Error forwarding message to MUD server:', error);
|
||||
if (ws.readyState === 1) { // WebSocket.OPEN
|
||||
const errorMessage = error instanceof Error ? error.message : String(error);
|
||||
ws.send(Buffer.from(`ERROR: Failed to send data to MUD server: ${errorMessage}\r\n`));
|
||||
}
|
||||
sendControl(ws, { type: 'error', code: initialized ? 'PROTOCOL_ERROR' : 'CONNECT_ERROR', message: error instanceof Error ? error.message : 'Unknown proxy error.' });
|
||||
if (!session) ws.close(1008, 'Connection rejected');
|
||||
}
|
||||
});
|
||||
|
||||
// Handle WebSocket close - THIS IS THE KEY CHANGE FOR PERSISTENCE
|
||||
ws.on('close', () => {
|
||||
console.log(`WebSocket closed for ${mudHost}:${mudPort} (session: ${currentSessionId})`);
|
||||
|
||||
const conn = connections.get(connectionId);
|
||||
if (conn && !conn.testMode && conn.socket && !conn.socket.destroyed) {
|
||||
console.log(`Moving connection to persistent state for ${conn.settings.persistenceTimeoutMs / 1000} seconds`);
|
||||
|
||||
// Move the connection to persistent storage instead of closing it
|
||||
// Use this connection's specific timeout setting
|
||||
const timeoutId = setTimeout(() => {
|
||||
console.log(`Session ${currentSessionId} timed out, closing MUD connection`);
|
||||
cleanupPersistentConnection(currentSessionId);
|
||||
}, conn.settings.persistenceTimeoutMs);
|
||||
|
||||
persistentConnections.set(currentSessionId, {
|
||||
socket: conn.socket,
|
||||
mudHost,
|
||||
mudPort,
|
||||
useSSL,
|
||||
timeoutId,
|
||||
lastActivity: Date.now(),
|
||||
messageBuffer: [],
|
||||
bufferSize: 0,
|
||||
settings: conn.settings // Store the connection's settings
|
||||
});
|
||||
|
||||
console.log(`Session ${currentSessionId} will persist for ${conn.settings.persistenceTimeoutMs / 1000} seconds with settings: ${conn.settings.maxBufferMessages} msgs/${conn.settings.maxBufferSizeKB}KB`);
|
||||
} else if (conn && conn.socket) {
|
||||
// Fallback to immediate cleanup if needed
|
||||
conn.socket.end();
|
||||
}
|
||||
|
||||
// Remove from active connections map
|
||||
connections.delete(connectionId);
|
||||
});
|
||||
|
||||
// Handle WebSocket errors
|
||||
ws.on('error', (error) => {
|
||||
console.error(`WebSocket error for ${mudHost}:${mudPort}:`, error.message);
|
||||
// Close socket on error - but only if it's not going to be persisted
|
||||
const conn = connections.get(connectionId);
|
||||
if (conn && conn.socket) {
|
||||
conn.socket.end();
|
||||
}
|
||||
// Remove from connections map
|
||||
connections.delete(connectionId);
|
||||
clearTimeout(initializationTimer);
|
||||
if (!session || explicitDisconnect || session.ws !== ws) return;
|
||||
session.ws = null;
|
||||
session.persistenceTimer = setTimeout(() => destroySession(session, 'resume_timeout'), PERSISTENCE_TIMEOUT_MS);
|
||||
});
|
||||
ws.on('error', () => ws.close());
|
||||
});
|
||||
|
||||
// Handle HTTP server upgrade (WebSocket handshake)
|
||||
server.on('upgrade', (request, socket, head) => {
|
||||
// Parse URL to get query parameters
|
||||
const { pathname, query } = parse(request.url || '', true);
|
||||
|
||||
// Only handle WebSocket connections to /mud-ws
|
||||
if (pathname === '/mud-ws') {
|
||||
// Extract MUD server details from query parameters
|
||||
const { host, port, useSSL } = query;
|
||||
|
||||
if (!host || !port) {
|
||||
socket.write('HTTP/1.1 400 Bad Request\r\n\r\n');
|
||||
socket.destroy();
|
||||
return;
|
||||
}
|
||||
|
||||
// Handle WebSocket upgrade
|
||||
wss.handleUpgrade(request, socket, head, (ws) => {
|
||||
wss.emit('connection', ws, request, host, port, useSSL === 'true');
|
||||
});
|
||||
} else {
|
||||
// For other upgrades (not to /mud-ws), close the connection
|
||||
socket.destroy();
|
||||
}
|
||||
let pathname;
|
||||
try { pathname = new URL(request.url || '/', 'http://localhost').pathname; } catch { return failUpgrade(socket, '400 Bad Request', 'Invalid URL'); }
|
||||
if (pathname !== '/mud-ws') return failUpgrade(socket, '404 Not Found', 'Not found');
|
||||
const origin = request.headers.origin;
|
||||
if (typeof origin !== 'string' || !allowedOrigins.has(origin)) return failUpgrade(socket, '403 Forbidden', 'Origin not allowed');
|
||||
wss.handleUpgrade(request, socket, head, (client) => wss.emit('connection', client, request));
|
||||
});
|
||||
|
||||
// Periodic cleanup of abandoned persistent connections
|
||||
setInterval(() => {
|
||||
const now = Date.now();
|
||||
for (const [sessionId, persistentConn] of persistentConnections.entries()) {
|
||||
// Clean up connections that have been inactive for too long
|
||||
// Use double the connection's specific timeout or default
|
||||
const timeoutThreshold = (persistentConn.settings?.persistenceTimeoutMs || DEFAULT_PERSISTENCE_TIMEOUT) * 2;
|
||||
if (now - persistentConn.lastActivity > timeoutThreshold) {
|
||||
console.log(`Cleaning up abandoned session: ${sessionId}`);
|
||||
cleanupPersistentConnection(sessionId);
|
||||
}
|
||||
}
|
||||
}, DEFAULT_PERSISTENCE_TIMEOUT); // Run cleanup every default timeout period
|
||||
|
||||
// Start the WebSocket server
|
||||
const PORT = process.env.WS_PORT || 3001;
|
||||
server.listen(PORT, () => {
|
||||
console.log(`WebSocket server is running on port ${PORT}`);
|
||||
console.log(`Default connection persistence timeout: ${DEFAULT_PERSISTENCE_TIMEOUT / 1000} seconds (configurable per connection)`);
|
||||
});
|
||||
|
||||
export default server;
|
||||
function shutdown() {
|
||||
wss.close();
|
||||
for (const session of [...sessions.values()]) destroySession(session, 'server_shutdown');
|
||||
server.close(() => process.exit(0));
|
||||
setTimeout(() => process.exit(1), 5_000).unref();
|
||||
}
|
||||
process.on('SIGINT', shutdown);
|
||||
process.on('SIGTERM', shutdown);
|
||||
if (process.env.NODE_ENV !== 'test') server.listen(PORT, () => console.log(`MUD WebSocket proxy listening on port ${PORT}`));
|
||||
export { server, validateConnectMessage, resolvePublicTarget, isDeniedAddress };
|
||||
|
||||
Reference in New Issue
Block a user