// WebSocket Server // For matting-server connection const WebSocket = require('ws'); const WS_PORT = 9527; // WebSocket server instance let wss = null; // Matting server client connection let mattingClient = null; // Pending callbacks map const pendingCallbacks = new Map(); /** * Start WebSocket Server */ function startWebSocketServer() { if (wss) { console.log('[WebSocket] Server already running'); return; } wss = new WebSocket.Server({ port: WS_PORT }); wss.on('listening', () => { console.log(`[WebSocket] Server listening on port: ${WS_PORT}`); }); wss.on('connection', (ws, req) => { const clientIP = req.socket.remoteAddress; console.log(`[WebSocket] New connection: ${clientIP}`); mattingClient = ws; ws.on('message', (data) => { try { // Binary format: 4-byte filename length + filename + zip data if (Buffer.isBuffer(data)) { const filenameLen = data.readUInt32BE(0); const filename = data.slice(4, 4 + filenameLen).toString('utf-8'); const zipData = data.slice(4 + filenameLen); console.log(`[WebSocket] Received: ${filename}, size: ${zipData.length} bytes`); // Find and execute callback const callback = pendingCallbacks.get(filename); if (callback) { pendingCallbacks.delete(filename); callback(null, zipData); } else { console.warn(`[WebSocket] No callback found for: ${filename}`); } } else { console.log(`[WebSocket] Received text: ${data}`); } } catch (error) { console.error('[WebSocket] Message handling error:', error); } }); ws.on('close', () => { console.log('[WebSocket] Connection closed'); if (mattingClient === ws) { mattingClient = null; } }); ws.on('error', (error) => { console.error('[WebSocket] Connection error:', error); }); }); wss.on('error', (error) => { console.error('[WebSocket] Server error:', error); }); } /** * Check if matting-server is connected * @returns {boolean} */ function isMattingServerConnected() { return mattingClient !== null && mattingClient.readyState === WebSocket.OPEN; } /** * Send matting task to matting-server * @param {string} taskId - Task ID (used as filename) * @param {Buffer} zipData - ZIP data buffer * @param {Function} callback - Callback function (error, resultZipData) * @param {number} timeout - Timeout in milliseconds (default 5 minutes) */ function sendMattingTask(taskId, zipData, callback, timeout = 5 * 60 * 1000) { if (!isMattingServerConnected()) { callback(new Error('matting-server not connected')); return; } try { // Binary format: 4-byte filename length + filename + zip data const filenameBytes = Buffer.from(taskId, 'utf-8'); const lenBuffer = Buffer.alloc(4); lenBuffer.writeUInt32BE(filenameBytes.length, 0); const message = Buffer.concat([lenBuffer, filenameBytes, zipData]); // Register callback pendingCallbacks.set(taskId, callback); // Set timeout const timeoutId = setTimeout(() => { if (pendingCallbacks.has(taskId)) { pendingCallbacks.delete(taskId); callback(new Error('VIP matting request timeout')); } }, timeout); // Wrap callback to clear timeout const originalCallback = callback; pendingCallbacks.set(taskId, (error, data) => { clearTimeout(timeoutId); originalCallback(error, data); }); // Send message mattingClient.send(message); console.log(`[WebSocket] Sent task: ${taskId}, size: ${message.length} bytes`); } catch (error) { pendingCallbacks.delete(taskId); callback(error); } } /** * Stop WebSocket Server */ function stopWebSocketServer() { if (wss) { wss.close(() => { console.log('[WebSocket] Server stopped'); }); wss = null; mattingClient = null; } } // Auto-start WebSocket Server startWebSocketServer(); module.exports = { startWebSocketServer, stopWebSocketServer, isMattingServerConnected, sendMattingTask };