Checkpoint/plugins/proxy.js

import { registerPlugin, loadConfig } from '../index.js';
import * as logs from '../utils/logs.js';

const proxyConfig = {};
await loadConfig('proxy', proxyConfig);

// Map configuration to internal structure
const enabled = proxyConfig.Core.Enabled;
const wsTimeout = proxyConfig.Timeouts.WebSocketTimeoutMs;
const upstreamTimeout = proxyConfig.Timeouts.UpstreamTimeoutMs;

// Build proxy mappings from array format
const proxyMappings = {};
proxyConfig.Mapping.forEach(mapping => {
  proxyMappings[mapping.Host] = mapping.Target;
});

logs.plugin('proxy', `Proxy mappings loaded: ${JSON.stringify(proxyMappings)}`);

const HOP_BY_HOP_HEADERS = [
  'connection',
  'keep-alive',
  'proxy-authenticate',
  'proxy-authorization',
  'te',
  'trailer',
  'transfer-encoding',
  'upgrade',
];

// Connect to upstream WebSocket with handshake timeout
async function connectUpstreamWebSocket(url, headers) {
  return new Promise((resolve, reject) => {
    const ws = new WebSocket(url, { headers });
    const timer = setTimeout(() => {
      ws.close();
      reject(new Error('timeout'));
    }, wsTimeout);
    ws.onopen = () => {
      clearTimeout(timer);
      resolve(ws);
    };
    ws.onerror = (err) => {
      clearTimeout(timer);
      reject(err);
    };
    ws.onclose = () => {
      clearTimeout(timer);
      reject(new Error('closed'));
    };
  });
}

async function createProxyResponse(targetURL, request) {
  try {
    const url = new URL(request.url);
    const targetPathAndQuery = url.pathname + url.search;
    const fullTargetURL = new URL(targetPathAndQuery, targetURL).toString();

    const startTime = Date.now();

    const outgoingHeaders = new Headers(request.headers);
    outgoingHeaders.delete('host');
    
    // Set proper host header for the target
    const targetHost = new URL(targetURL).host;
    outgoingHeaders.set('host', targetHost);
    
    // Forward the original host as X-Forwarded-Host for applications that need it
    outgoingHeaders.set('x-forwarded-host', request.headers.get('host'));
    outgoingHeaders.set('x-forwarded-proto', url.protocol.replace(':', ''));
    
    const options = {
      method: request.method,
      headers: outgoingHeaders,
      // Always use manual redirect to let client handle it
      redirect: 'manual',
      // Don't decode compressed responses - let the client handle it
      decompress: false,
    };

    // Handle request body
    if (request.body && ['POST', 'PUT', 'PATCH', 'DELETE'].includes(request.method)) {
      options.body = request.body;
    }

    // Add timeout
    const controller = new AbortController();
    const timeoutId = setTimeout(() => {
      logs.warn('proxy', `Upstream request to ${fullTargetURL} timed out after ${upstreamTimeout}ms`);
      controller.abort();
    }, upstreamTimeout);

    let response;
    try {
      response = await fetch(fullTargetURL, { ...options, signal: controller.signal });
    } catch (fetchErr) {
      clearTimeout(timeoutId);
      if (fetchErr.name === 'AbortError') {
        logs.error('proxy', `Upstream fetch aborted for ${fullTargetURL} (likely due to timeout)`);
        return new Response('Gateway Timeout', { status: 504 });
      }
      logs.error('proxy', `Fetch error: ${fetchErr.message}`);
      return new Response('Bad Gateway', { status: 502 });
    }
    clearTimeout(timeoutId);

    const latency = Date.now() - startTime;
    logs.plugin('proxy', `Proxied request to: ${fullTargetURL} (${response.status} ${response.statusText}) (${latency}ms)`);
    
    const responseHeaders = new Headers(response.headers);
    
    // Remove hop-by-hop headers
    HOP_BY_HOP_HEADERS.forEach((h) => responseHeaders.delete(h));
    
    // IMPORTANT: Don't remove content-encoding or modify the body
    // Let the response stream through as-is for SSE compatibility
    
    // Add proxy information
    responseHeaders.set('X-Proxy-Latency', `${latency}ms`);
    
    // Handle Set-Cookie headers - rewrite domain if needed
    const setCookieHeaders = response.headers.getSetCookie ? response.headers.getSetCookie() : [];
    if (setCookieHeaders.length > 0) {
      responseHeaders.delete('set-cookie');
      
      setCookieHeaders.forEach(cookieStr => {
        let modifiedCookie = cookieStr;
        
        // Remove domain restrictions
        modifiedCookie = modifiedCookie.replace(/;\s*domain=[^;]*/gi, '');
        
        // Handle SameSite for local development
        if (url.protocol === 'http:' && modifiedCookie.match(/samesite\s*=\s*none/i)) {
          modifiedCookie = modifiedCookie.replace(/;\s*samesite=[^;]*/gi, '; SameSite=Lax');
          modifiedCookie = modifiedCookie.replace(/;\s*secure/gi, '');
        }
        
        responseHeaders.append('set-cookie', modifiedCookie);
      });
    }

    // Return response with original body stream
    return new Response(response.body, {
      status: response.status,
      statusText: response.statusText,
      headers: responseHeaders,
    });
  } catch (err) {
    logs.error('proxy', `Proxy error processing ${request.method} ${request.url}: ${err.message}`);
    logs.error('proxy', `Full error details: ${err.stack}`);
    return new Response('Bad Gateway', { status: 502 });
  }
}

function proxyMiddleware() {
  return async (request, server) => {
    const url = new URL(request.url);
    const path = url.pathname;

    // Skip checkpoint endpoints
    if (path.startsWith('/api/challenge') || path.startsWith('/api/verify')) return undefined;

    // Skip static assets
    if (path.startsWith('/webfont/') || path.startsWith('/js/')) return undefined;

    // Get the hostname from the request
    const hostname = request.headers.get('host')?.split(':')[0];
    const target = proxyMappings[hostname];
    if (!target) return undefined;

    // Handle WebSocket upgrade requests
    const upgradeHeader = request.headers.get('upgrade')?.toLowerCase();
    if (upgradeHeader === 'websocket') {
      const targetUrl = new URL(url.pathname + url.search, target);
      targetUrl.protocol = targetUrl.protocol.replace(/^http/, 'ws');
      
      // Forward important headers for WebSocket
      const wsHeaders = {};
      ['cookie', 'authorization', 'origin', 'sec-websocket-protocol', 'sec-websocket-extensions']
        .forEach(header => {
          const value = request.headers.get(header);
          if (value) wsHeaders[header] = value;
        });

      let upstream;
      try {
        upstream = await connectUpstreamWebSocket(targetUrl.toString(), wsHeaders);
      } catch (err) {
        logs.error('proxy', `Upstream WebSocket connection failed: ${err}`);
        return new Response('Bad Gateway', { status: 502 });
      }
      
      // Upgrade incoming client connection and attach upstream socket
      const ok = server.upgrade(request, { data: { upstream } });
      if (!ok) {
        logs.error('proxy', 'WebSocket upgrade failed');
        upstream.close();
        return new Response('Bad Gateway', { status: 502 });
      }
      logs.plugin('proxy', `WebSocket proxied to: ${targetUrl.toString()}`);
      return;
    }

    return createProxyResponse(target, request);
  };
}

// WebSocket handlers for proxying messages between client and upstream
export const proxyWebSocketHandler = {
  open(ws) {
    const upstream = ws.data.upstream;
    upstream.onopen = () => logs.plugin('proxy', 'Upstream WebSocket connected');
    // Forward messages from target to client
    upstream.onmessage = (event) => ws.send(event.data);
    upstream.onerror = (err) => {
      logs.error('proxy', `Upstream WebSocket error: ${err}`);
      ws.close(1011, 'Upstream error');
    };
    upstream.onclose = ({ code, reason }) => ws.close(code, reason);
  },
  message(ws, message) {
    const upstream = ws.data.upstream;
    // Forward messages from client to target
    upstream.send(message);
  },
  close(ws, code, reason) {
    const upstream = ws.data.upstream;
    upstream.close(code, reason);
  },
  error(ws, err) {
    logs.error('proxy', `WebSocket proxy error: ${err}`);
    const upstream = ws.data.upstream;
    upstream.close();
  },
};

if (enabled) {
  registerPlugin('proxy', proxyMiddleware());
} else {
  logs.plugin('proxy', 'Proxy plugin disabled via config');
}