proxy-handler.js

const httpProxy = require('http-proxy');
const http = require('http');
const https = require('https');
const config = require('./config');

class ProxyHandler {
    constructor() {
        // Connection pooling agents for better performance
        this.httpAgent = new http.Agent({
            keepAlive: true,
            keepAliveMsecs: 30000,
            maxSockets: 128,
            maxFreeSockets: 32,
            timeout: 60000
        });
        this.httpsAgent = new https.Agent({
            keepAlive: true,
            keepAliveMsecs: 30000,
            maxSockets: 128,
            maxFreeSockets: 32,
            timeout: 60000,
            rejectUnauthorized: false
        });

        this.proxy = httpProxy.createProxyServer({
            changeOrigin: true,
            ws: true,
            followRedirects: true,
            agent: this.httpsAgent,
            xfwd: true,
            preserveHeaderKeyCase: true
        });

        this.rewriters = [];
        this.responseModifiers = [];

        // Circuit breaker state per target host
        this._circuits = new Map();
        this._circuitThreshold = 5;    // failures before opening circuit
        this._circuitResetMs = 30000;  // time before half-open retry

        this.setupDefaultHandlers();
    }

    setupDefaultHandlers() {
        // Response rewriting
        this.proxy.on('proxyRes', (proxyRes, req, res) => {
            // Record success for circuit breaker
            if (req._target) this._recordSuccess(req._target);
            this.handleProxyResponse(proxyRes, req, res);
        });

        // Proxy error handling with categorization
        this.proxy.on('error', (err, req, res) => {
            const category = this._categorizeError(err);
            console.error(`[Proxy ${category}]`, err.code || err.message, req?.url);

            // Record failure for circuit breaker
            if (req?._target) this._recordFailure(req._target);

            if (res && !res.headersSent && typeof res.writeHead === 'function') {
                const status = category === 'timeout' ? 504 : 502;
                res.writeHead(status, { 'Content-Type': 'application/json' });
                res.end(JSON.stringify({
                    error: category === 'timeout' ? 'Gateway timeout' : 'Bad gateway',
                    code: err.code,
                    category
                }));
            }
        });

        // ECONNRESET handling
        this.proxy.on('econnreset', (err, req, res) => {
            console.error('[Proxy ECONNRESET]', err.message, req?.url);
            if (req?._target) this._recordFailure(req._target);
            if (res && !res.headersSent && typeof res.writeHead === 'function') {
                res.writeHead(504, { 'Content-Type': 'application/json' });
                res.end(JSON.stringify({ error: 'Connection reset', code: 'ECONNRESET' }));
            }
        });
    }

    // Categorize proxy errors for better diagnostics
    _categorizeError(err) {
        const code = err.code || '';
        if (['ETIMEDOUT', 'ESOCKETTIMEDOUT', 'ECONNABORTED'].includes(code)) return 'timeout';
        if (['ECONNREFUSED', 'ENOTFOUND', 'EAI_AGAIN'].includes(code)) return 'unreachable';
        if (['ECONNRESET', 'EPIPE', 'ERR_SOCKET_CLOSED'].includes(code)) return 'connection_lost';
        if (code === 'HPE_INVALID_HEADER_TOKEN') return 'protocol';
        return 'unknown';
    }

    // Circuit breaker - record success
    _recordSuccess(target) {
        try {
            const host = new URL(target).host;
            const circuit = this._circuits.get(host);
            if (circuit) {
                circuit.failures = 0;
                circuit.state = 'closed';
            }
        } catch {}
    }

    // Circuit breaker - record failure
    _recordFailure(target) {
        try {
            const host = new URL(target).host;
            let circuit = this._circuits.get(host);
            if (!circuit) {
                circuit = { failures: 0, state: 'closed', lastFailure: 0 };
                this._circuits.set(host, circuit);
            }
            circuit.failures++;
            circuit.lastFailure = Date.now();
            if (circuit.failures >= this._circuitThreshold) {
                circuit.state = 'open';
                console.warn(`[Proxy] Circuit OPEN for ${host} (${circuit.failures} failures)`);
            }
        } catch {}
    }

    // Check if circuit is allowing requests
    _isCircuitOpen(target) {
        try {
            const host = new URL(target).host;
            const circuit = this._circuits.get(host);
            if (!circuit || circuit.state === 'closed') return false;
            // Half-open: allow retry after reset period
            if (Date.now() - circuit.lastFailure > this._circuitResetMs) {
                circuit.state = 'half-open';
                return false;
            }
            return true;
        } catch { return false; }
    }

    // Get circuit breaker status for all targets
    getCircuitStatus() {
        const status = {};
        for (const [host, circuit] of this._circuits) {
            status[host] = { ...circuit };
        }
        return status;
    }

    // Register a URL rewriter
    addRewriter(pattern, replacement) {
        this.rewriters.push({ pattern, replacement });
        return this;
    }

    // Register a response modifier function
    addResponseModifier(modifier) {
        this.responseModifiers.push(modifier);
        return this;
    }

    handleProxyResponse(proxyRes, req, res) {
        const target = req._target || config.DISCORD;
        const isDiscord = target.startsWith('https://discord.com');

        // Fix redirects
        if (proxyRes.headers.location) {
            let location = proxyRes.headers.location;
            location = this.rewriteUrl(location);
            proxyRes.headers.location = location;
            console.log(`[REDIRECT FIXED] → ${location}`);
        }

        // Remove blocking headers (match original working proxy)
        delete proxyRes.headers['content-security-policy'];
        delete proxyRes.headers['content-security-policy-report-only'];
        delete proxyRes.headers['strict-transport-security'];

        // Mirror the request origin for CORS (required for credentials)
        proxyRes.headers['access-control-allow-origin'] = req.headers.origin || '*';
        proxyRes.headers['access-control-allow-credentials'] = 'true';

        // For non-Discord targets, pipe through directly (no rewriting)
        if (!isDiscord) {
            // Timeout for non-Discord pipe
            const pipeTimeout = setTimeout(() => {
                console.error('[Proxy] Pipe timeout for:', req.url);
                proxyRes.destroy();
                if (!res.headersSent) {
                    res.writeHead(504, { 'Content-Type': 'text/plain' });
                    res.end('Gateway timeout');
                } else if (!res.writableEnded) {
                    try { res.end(); } catch (e) {}
                }
            }, 30000);

            // Abort on client disconnect
            const onAbort = () => {
                clearTimeout(pipeTimeout);
                proxyRes.destroy();
            };
            req.on('close', onAbort);

            proxyRes.on('error', (err) => {
                clearTimeout(pipeTimeout);
                console.error('[Proxy] Pipe error:', err.message, req.url);
                if (!res.headersSent) {
                    res.writeHead(502, { 'Content-Type': 'text/plain' });
                    res.end('Bad gateway');
                } else if (!res.writableEnded) {
                    try { res.end(); } catch (e) {}
                }
            });

            proxyRes.on('end', () => {
                clearTimeout(pipeTimeout);
                req.removeListener('close', onAbort);
            });

            res.writeHead(proxyRes.statusCode, proxyRes.headers);
            proxyRes.pipe(res);
            return;
        }

        // Handle client disconnect
        const abortHandler = () => {
            proxyRes.destroy();
        };
        res.on('close', abortHandler);
        req.on('close', abortHandler);

        // Response timeout - prevent hanging if 'end' never fires
        const responseTimeout = setTimeout(() => {
            console.error('[Proxy] Response timeout for:', req.url);
            proxyRes.destroy();
            if (!res.headersSent) {
                res.writeHead(504);
                res.end('Gateway timeout');
            }
        }, 30000);

        // Handle HEAD requests and empty responses immediately
        if (req.method === 'HEAD' || proxyRes.statusCode === 204 || proxyRes.headers['content-length'] === '0') {
            clearTimeout(responseTimeout);
            res.removeListener('close', abortHandler);
            req.removeListener('close', abortHandler);
            res.writeHead(proxyRes.statusCode, proxyRes.headers);
            res.end();
            return;
        }

        // Collect body from Discord responses for rewriting
        const bodyBuffers = [];
        let totalSize = 0;
        const MAX_SIZE = 50 * 1024 * 1024; // 50MB limit
        let ended = false;

        proxyRes.on('data', chunk => {
            totalSize += chunk.length;
            if (totalSize > MAX_SIZE) {
                console.error('[Proxy] Response too large, aborting:', req.url);
                proxyRes.destroy();
                if (!res.headersSent) {
                    res.writeHead(413);
                    res.end('Response too large');
                }
                return;
            }
            bodyBuffers.push(chunk);
        });

        proxyRes.on('end', () => {
            if (ended) return;
            ended = true;
            clearTimeout(responseTimeout);
            res.removeListener('close', abortHandler);
            req.removeListener('close', abortHandler);

            const originalBody = Buffer.concat(bodyBuffers);

            // Skip modification for non-text responses (binary, compressed)
            const contentEncoding = proxyRes.headers['content-encoding'];
            const contentType = proxyRes.headers['content-type'] || '';

            if (contentEncoding || (!contentType.includes('text') && !contentType.includes('json') && !contentType.includes('javascript') && !contentType.includes('html'))) {
                if (proxyRes.headers['content-length']) {
                    proxyRes.headers['content-length'] = String(originalBody.length);
                }
                res.writeHead(proxyRes.statusCode, proxyRes.headers);
                res.end(originalBody);
                return;
            }

            let bodyStr = originalBody.toString('utf8');

            // Apply all URL rewriters
            bodyStr = this.rewriteBody(bodyStr);

            // Apply custom response modifiers
            for (const modifier of this.responseModifiers) {
                try {
                    bodyStr = modifier(bodyStr, proxyRes, req) || bodyStr;
                } catch (err) {
                    console.error('[Proxy] Modifier error:', err.message);
                }
            }

            const modifiedBody = Buffer.from(bodyStr, 'utf8');

            if (proxyRes.headers['content-length']) {
                proxyRes.headers['content-length'] = String(modifiedBody.length);
            }

            res.writeHead(proxyRes.statusCode, proxyRes.headers);
            res.end(modifiedBody);
        });

        // Handle errors
        proxyRes.on('error', (err) => {
            if (ended) return;
            ended = true;
            clearTimeout(responseTimeout);
            console.error('[Proxy] Response error:', err.message);
            if (!res.headersSent) {
                res.writeHead(502);
                res.end('Bad gateway');
            }
        });

        // Handle connection close (upstream dropped before 'end' fired)
        proxyRes.on('close', () => {
            if (ended) return;
            ended = true;
            clearTimeout(responseTimeout);
            res.removeListener('close', abortHandler);
            req.removeListener('close', abortHandler);
            if (!res.headersSent) {
                res.writeHead(502, { 'Content-Type': 'text/plain' });
                res.end('Connection closed');
            } else if (!res.writableEnded) {
                // Headers sent but body never completed — end it now
                try { res.end(); } catch (e) {}
            }
        });
    }

    rewriteUrl(url) {
        for (const { pattern, replacement } of this.rewriters) {
            url = url.replace(pattern, replacement);
        }
        return url;
    }

    rewriteBody(body) {
        const { LOCAL_URL, HTTPS_PORT } = config;

        // 1. Full HTTP(S) URLs
        body = body.replace(
            /https?:\/\/(?:www\.)?discord(?:app)?\.com/g,
            LOCAL_URL
        );

        // 2. Protocol-relative
        body = body.replace(
            /\/\/(?:www\.)?discord(?:app)?\.com/g,
            `//localhost:${HTTPS_PORT}`
        );

        // 3. Bare domain inside JS
        body = body.replace(
            /(?<![\w\.-])discord(?:app)?\.com/gi,
            `localhost:${HTTPS_PORT}`
        );

        // 4. Gateway WebSocket
        body = body.replace(
            /wss?:\/\/gateway\.discord\.gg/g,
            `wss://localhost:${HTTPS_PORT}/gateway`
        );
        body = body.replace(
            /\/\/gateway\.discord\.gg/g,
            `//localhost:${HTTPS_PORT}/gateway`
        );

        // 5. Apply plugin rewriters (e.g., CDN URLs)
        for (const { pattern, replacement } of this.rewriters) {
            body = body.replace(pattern, replacement);
        }

        return body;
    }

    web(req, res, options, callback, _retryCount = 0) {
        const maxRetries = options.retries || 1;
        const target = options.target;

        // Circuit breaker check
        if (this._isCircuitOpen(target)) {
            const err = new Error('Circuit breaker open');
            err.code = 'ECIRCUIT_OPEN';
            console.warn('[Proxy] Circuit open, rejecting:', target);
            if (callback) return callback(err);
            if (!res.headersSent) {
                res.writeHead(503, { 'Content-Type': 'application/json' });
                res.end(JSON.stringify({ error: 'Service unavailable', code: 'ECIRCUIT_OPEN' }));
            }
            return;
        }

        // Select correct agent based on target protocol
        if (target && !options.agent) {
            options.agent = target.startsWith('https') ? this.httpsAgent : this.httpAgent;
        }

        const proxyReq = this.proxy.web(req, res, options, (err) => {
            if (err) {
                const category = this._categorizeError(err);
                console.error(`[Proxy] Web ${category}:`, err.code || err.message, req.url);

                // Retry on transient errors
                if (_retryCount < maxRetries && ['timeout', 'connection_lost'].includes(category)) {
                    console.log(`[Proxy] Retrying (${_retryCount + 1}/${maxRetries}):`, req.url);
                    return this.web(req, res, options, callback, _retryCount + 1);
                }
            }
            if (callback) callback(err);
        });

        // Handle request-level timeouts
        if (proxyReq && proxyReq.setTimeout) {
            proxyReq.setTimeout(options.timeout || 30000, () => {
                console.error('[Proxy] Request timeout:', req.url);
                proxyReq.destroy();
            });
        }

        return proxyReq;
    }

    ws(req, socket, head, options) {
        // Error handling for WebSocket
        const onError = (err) => {
            console.error('[WS] Error:', err.message);
            if (req._target) this._recordFailure(req._target);
            try {
                socket.destroy();
            } catch (e) {}
        };

        socket.on('error', onError);

        // Select correct agent
        const target = options.target;
        if (target && !options.agent) {
            options.agent = target.startsWith('https') ? this.httpsAgent : this.httpAgent;
        }

        try {
            this.proxy.ws(req, socket, head, {
                ...options,
                changeOrigin: true
            }, (err) => {
                if (err) onError(err);
            });
        } catch (err) {
            onError(err);
        }
    }

    // Cleanup agents on shutdown
    destroy() {
        this.httpAgent.destroy();
        this.httpsAgent.destroy();
    }
}

module.exports = ProxyHandler;