// ── Port Proxy tunnel handler (server side) ─────────────────────────────────
// Routes HTTP/WS proxy requests from browsers to the agent via the existing
// encrypted WebSocket channel.  Also caches per-agent proxy config.
import { randomUUID } from 'crypto';
import { WebSocket } from 'ws';
import { sessions } from './session-manager.js';
import { encryptAndSend } from './encryption.js';
// ── State ───────────────────────────────────────────────────────────────────
// Agent proxy configs (agentId → config), authoritative copy lives on agent
const agentProxyConfigs = new Map();
// Pending HTTP tunnel requests (tunnelId → promise resolver)
const pendingRequests = new Map();
// Active WebSocket tunnels (tunnelId → browser ws + agentId)
const activeTunnelWs = new Map();
const HTTP_TIMEOUT = 30_000;
const MAX_CONCURRENT = 50;
// Track concurrent requests per session
const sessionConcurrency = new Map();
// ── Exported API ────────────────────────────────────────────────────────────
export function updateProxyConfig(agentId, config) {
    agentProxyConfigs.set(agentId, config);
    console.log(`[Tunnel] Config updated for agent ${agentId}: enabled=${config.enabled}, ports=[${config.ports.map(p => `${p.port}(${p.enabled ? 'on' : 'off'})`).join(', ')}]`);
}
export function getProxyConfig(agentId) {
    return agentProxyConfigs.get(agentId);
}
export function removeProxyConfig(agentId) {
    agentProxyConfigs.delete(agentId);
}
function isPortAllowed(agentId, port) {
    const config = agentProxyConfigs.get(agentId);
    if (!config || !config.enabled)
        return false;
    const entry = config.ports.find(p => p.port === port);
    return !!entry && entry.enabled;
}
// ── Agent message handler ───────────────────────────────────────────────────
/**
 * Returns true if the message was handled by the tunnel module (should NOT
 * be forwarded to web clients).
 */
export function handleAgentTunnelMessage(msg) {
    switch (msg.type) {
        case 'tunnel_response': {
            const tunnelId = msg.tunnelId;
            const pending = pendingRequests.get(tunnelId);
            if (pending) {
                clearTimeout(pending.timer);
                pendingRequests.delete(tunnelId);
                pending.resolve(msg);
            }
            return true;
        }
        case 'tunnel_ws_opened': {
            // nothing extra — the browser ws is already connected
            return true;
        }
        case 'tunnel_ws_message': {
            const tunnelId = msg.tunnelId;
            const tunnel = activeTunnelWs.get(tunnelId);
            if (tunnel && tunnel.browserWs.readyState === WebSocket.OPEN) {
                try {
                    const buf = Buffer.from(msg.data, 'base64');
                    tunnel.browserWs.send(msg.binary ? buf : buf.toString('utf8'));
                }
                catch (err) {
                    console.error(`[Tunnel] Failed to send to browser WS ${tunnelId}:`, err.message);
                }
            }
            return true;
        }
        case 'tunnel_ws_close': {
            const tunnelId = msg.tunnelId;
            const tunnel = activeTunnelWs.get(tunnelId);
            if (tunnel) {
                activeTunnelWs.delete(tunnelId);
                if (tunnel.browserWs.readyState === WebSocket.OPEN) {
                    const code = typeof msg.code === 'number' && msg.code >= 1000 && msg.code <= 4999 ? msg.code : 1000;
                    try {
                        tunnel.browserWs.close(code, msg.reason || '');
                    }
                    catch (err) {
                        console.error(`[Tunnel] Failed to close browser WS ${tunnelId}:`, err.message);
                        try {
                            tunnel.browserWs.terminate();
                        }
                        catch { /* swallow */ }
                    }
                }
            }
            return true;
        }
        case 'tunnel_ws_error': {
            const tunnelId = msg.tunnelId;
            const tunnel = activeTunnelWs.get(tunnelId);
            if (tunnel) {
                activeTunnelWs.delete(tunnelId);
                if (tunnel.browserWs.readyState === WebSocket.OPEN) {
                    try {
                        tunnel.browserWs.close(1011, msg.message || 'Tunnel error');
                    }
                    catch (err) {
                        console.error(`[Tunnel] Failed to close browser WS ${tunnelId} on error:`, err.message);
                        try {
                            tunnel.browserWs.terminate();
                        }
                        catch { /* swallow */ }
                    }
                }
            }
            return true;
        }
        default:
            return false;
    }
}
// ── HTTP proxy handler (Express middleware) ──────────────────────────────────
export function httpProxyHandler(req, res) {
    const sessionId = req.params.sessionId;
    const port = parseInt(req.params.port, 10);
    if (!port || port < 1024 || port > 65535) {
        res.status(400).json({ error: 'Invalid port. Must be 1024-65535.' });
        return;
    }
    const agent = sessions.getAgentBySession(sessionId);
    if (!agent) {
        res.status(404).json({ error: 'Agent not found for this session.' });
        return;
    }
    if (!isPortAllowed(agent.agentId, port)) {
        res.status(403).json({ error: `Port ${port} is not enabled for proxy.` });
        return;
    }
    // Concurrency check
    const current = sessionConcurrency.get(sessionId) || 0;
    if (current >= MAX_CONCURRENT) {
        res.status(429).json({ error: 'Too many concurrent proxy requests.' });
        return;
    }
    sessionConcurrency.set(sessionId, current + 1);
    const tunnelId = randomUUID();
    const proxyPath = '/' + (req.params[0] || '');
    const fullPath = proxyPath + (req.url.includes('?') ? '?' + req.url.split('?')[1] : '');
    // Collect request body
    const bodyChunks = [];
    req.on('data', (chunk) => bodyChunks.push(chunk));
    req.on('end', () => {
        const body = bodyChunks.length > 0 ? Buffer.concat(bodyChunks).toString('base64') : undefined;
        // Extract headers (skip hop-by-hop)
        const headers = {};
        for (const [k, v] of Object.entries(req.headers)) {
            if (typeof v !== 'string')
                continue;
            const lower = k.toLowerCase();
            if (lower === 'host' || lower === 'connection' || lower === 'upgrade')
                continue;
            headers[k] = v;
        }
        // Send tunnel_request to agent
        const tunnelReq = {
            type: 'tunnel_request',
            tunnelId,
            port,
            method: req.method,
            path: fullPath,
            headers,
            body,
        };
        const promise = new Promise((resolve, reject) => {
            const timer = setTimeout(() => {
                pendingRequests.delete(tunnelId);
                reject(new Error('Tunnel request timed out'));
            }, HTTP_TIMEOUT);
            pendingRequests.set(tunnelId, { resolve, reject, timer });
        });
        // Pre-check: encryptAndSend swallows internal errors (see encryption.ts),
        // so its returned promise resolves even when the underlying ws.send fails.
        // If the agent socket is already closed, fail fast with 502 here instead
        // of waiting 30s for the timeout.
        if (agent.ws.readyState !== WebSocket.OPEN) {
            const pending = pendingRequests.get(tunnelId);
            if (pending) {
                clearTimeout(pending.timer);
                pendingRequests.delete(tunnelId);
                const sendErr = new Error(`Agent WebSocket not open (readyState=${agent.ws.readyState})`);
                sendErr.status = 502;
                pending.reject(sendErr);
            }
        }
        else {
            encryptAndSend(agent.ws, tunnelReq, agent.sessionKey).catch((err) => {
                console.error(`[Tunnel] Failed to send tunnel_request for ${tunnelId}:`, err.message);
                // Belt-and-suspenders: if encryptAndSend ever does reject in the future,
                // also fail fast here with 502.
                const pending = pendingRequests.get(tunnelId);
                if (pending) {
                    clearTimeout(pending.timer);
                    pendingRequests.delete(tunnelId);
                    const sendErr = new Error(`Failed to send tunnel_request: ${err.message}`);
                    sendErr.status = 502;
                    pending.reject(sendErr);
                }
            });
        }
        promise.then((resp) => {
            const dec = sessionConcurrency.get(sessionId);
            if (dec)
                sessionConcurrency.set(sessionId, dec - 1);
            if (resp.error) {
                const status = typeof resp.status === 'number' ? resp.status : 502;
                res.status(status).json({ error: resp.error });
                return;
            }
            const status = typeof resp.status === 'number' ? resp.status : 200;
            const respHeaders = (resp.headers || {});
            // Set response headers (skip hop-by-hop)
            for (const [k, v] of Object.entries(respHeaders)) {
                const lower = k.toLowerCase();
                if (lower === 'transfer-encoding' || lower === 'connection')
                    continue;
                res.setHeader(k, v);
            }
            const respBody = resp.body ? Buffer.from(resp.body, 'base64') : Buffer.alloc(0);
            // Inject <base> tag into HTML responses so relative asset paths resolve through the proxy prefix
            const contentType = (respHeaders['content-type'] || respHeaders['Content-Type'] || '');
            if (contentType.includes('text/html') && respBody.length > 0) {
                const html = respBody.toString('utf-8');
                const prefix = req.originalUrl?.startsWith('/ms/') ? 'ms' : 's';
                const baseHref = `/${prefix}/${sessionId}/proxy/${port}/`;
                const injected = html.replace('<head>', `<head><base href="${baseHref}">`);
                res.status(status).send(injected);
            }
            else {
                res.status(status).send(respBody);
            }
        }).catch((err) => {
            const dec = sessionConcurrency.get(sessionId);
            if (dec)
                sessionConcurrency.set(sessionId, dec - 1);
            const status = err.status ?? 504;
            res.status(status).json({ error: err.message });
        });
    });
}
// ── WebSocket tunnel upgrade handler ────────────────────────────────────────
/**
 * Returns true if the URL matches a proxy WebSocket path and was handled
 * (or rejected). Returns false if the URL is not a proxy path.
 */
export function handleTunnelWsUpgrade(req, socket, head, wss) {
    const url = new URL(req.url || '/', `http://${req.headers.host}`);
    // Match: /s/<sessionId>/proxy/<port>[/path...] or /ms/<sessionId>/proxy/<port>[/path...]
    const match = url.pathname.match(/^\/(s|ms)\/([^/]+)\/proxy\/(\d+)(\/.*)?$/);
    if (!match)
        return false;
    const sessionId = match[2];
    const port = parseInt(match[3], 10);
    const proxyPath = (match[4] || '/') + url.search;
    if (!port || port < 1024 || port > 65535) {
        socket.destroy();
        return true;
    }
    const agent = sessions.getAgentBySession(sessionId);
    if (!agent) {
        socket.destroy();
        return true;
    }
    if (!isPortAllowed(agent.agentId, port)) {
        socket.destroy();
        return true;
    }
    // Accept the WebSocket upgrade from the browser
    wss.handleUpgrade(req, socket, head, (browserWs) => {
        const tunnelId = randomUUID();
        activeTunnelWs.set(tunnelId, { browserWs, agentId: agent.agentId });
        // Extract headers
        const headers = {};
        for (const [k, v] of Object.entries(req.headers)) {
            if (typeof v === 'string')
                headers[k] = v;
        }
        // Send tunnel_ws_open to agent
        encryptAndSend(agent.ws, {
            type: 'tunnel_ws_open',
            tunnelId,
            port,
            path: proxyPath,
            headers,
        }, agent.sessionKey).catch((err) => {
            console.error(`[Tunnel] Failed to send tunnel_ws_open for ${tunnelId}:`, err.message);
        });
        // Relay browser → agent messages
        browserWs.on('message', (data, isBinary) => {
            const buf = Buffer.isBuffer(data) ? data : Buffer.from(data);
            encryptAndSend(agent.ws, {
                type: 'tunnel_ws_message',
                tunnelId,
                data: buf.toString('base64'),
                binary: isBinary,
            }, agent.sessionKey).catch((err) => {
                console.error(`[Tunnel] Failed to relay browser→agent msg for ${tunnelId}:`, err.message);
            });
        });
        browserWs.on('close', (code, reason) => {
            activeTunnelWs.delete(tunnelId);
            const safeCode = code >= 1000 && code <= 4999 ? code : 1000;
            encryptAndSend(agent.ws, {
                type: 'tunnel_ws_close',
                tunnelId,
                code: safeCode,
                reason: reason?.toString(),
            }, agent.sessionKey).catch((err) => {
                console.error(`[Tunnel] Failed to send tunnel_ws_close for ${tunnelId}:`, err.message);
            });
        });
        browserWs.on('error', () => {
            activeTunnelWs.delete(tunnelId);
            encryptAndSend(agent.ws, {
                type: 'tunnel_ws_close',
                tunnelId,
                code: 1011,
                reason: 'Browser WebSocket error',
            }, agent.sessionKey).catch((err) => {
                console.error(`[Tunnel] Failed to send tunnel_ws_close (error) for ${tunnelId}:`, err.message);
            });
        });
    });
    return true;
}
// ── Cleanup on agent disconnect ─────────────────────────────────────────────
export function cleanupAgentTunnels(agentId) {
    removeProxyConfig(agentId);
    // Close any active WebSocket tunnels for this agent
    for (const [tunnelId, tunnel] of activeTunnelWs) {
        if (tunnel.agentId === agentId) {
            activeTunnelWs.delete(tunnelId);
            if (tunnel.browserWs.readyState === WebSocket.OPEN) {
                try {
                    tunnel.browserWs.close(1001, 'Agent disconnected');
                }
                catch { /* swallow */ }
            }
        }
    }
}
//# sourceMappingURL=tunnel.js.map