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(':', ''));

    // Preserve important headers for authentication
    // Don't delete content-length or transfer-encoding here, handle them properly below
    const options = {
      method: request.method,
      headers: outgoingHeaders,
      // Follow redirects automatically for GET; forward redirects for non-GET
      // Absolute requirement: DONT REMOVE
      redirect: request.method === 'GET' ? 'follow' : 'manual',
      credentials: 'include',
    };

    const isChunked = request.headers.get('transfer-encoding')?.toLowerCase() === 'chunked';

    // Define methods that can legitimately have request bodies
    const methodsWithBody = new Set(['POST', 'PUT', 'PATCH', 'DELETE']);

    if (methodsWithBody.has(request.method) && request.body) {
      if (isChunked) {
        logs.plugin('proxy', `De-chunking request body for ${request.method} ${request.url}`);
        try {
          const bodyBuffer = await request.arrayBuffer();
          options.body = bodyBuffer;
          outgoingHeaders.set('content-length', String(bodyBuffer.byteLength));
          outgoingHeaders.delete('transfer-encoding');
        } catch (bufferError) {
          logs.error('proxy', `Error buffering chunked request body: ${bufferError}`);
          return new Response('Error processing chunked request body', { status: 500 });
        }
      } else {
        // For non-chunked bodies, preserve the body stream
        options.body = request.body;
        // Keep the original content-length if it exists
        if (request.headers.has('content-length')) {
          outgoingHeaders.set('content-length', request.headers.get('content-length'));
        }
      }
    }

    // Add a timeout controller for the upstream fetch
    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,
        verbose: true,
      });
    } 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 });
      }
      throw fetchErr;
    }
    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));

    // Remove content-encoding and content-length headers
    // This is necessary because Bun/fetch automatically decompresses the response body
    // but leaves the content-encoding header, causing the browser to try to decompress already decompressed content
    responseHeaders.delete('content-encoding');
    responseHeaders.delete('content-length');

    // 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) => {
        // Parse and potentially rewrite the cookie domain
        let modifiedCookie = cookieStr;

        // Remove domain restrictions that might prevent the cookie from working
        modifiedCookie = modifiedCookie.replace(/;\s*domain=[^;]*/gi, '');

        // If the cookie has SameSite=None, ensure it also has Secure
        if (modifiedCookie.match(/samesite\s*=\s*none/i) && !modifiedCookie.match(/secure/i)) {
          modifiedCookie += '; Secure';
        }

        // For local development, you might need to adjust SameSite
        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 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}${
        err.cause ? ' - Cause: ' + err.cause : ''
      }`,
    );
    let causeDetails = '';
    if (err.cause) {
      causeDetails = typeof err.cause === 'object' ? JSON.stringify(err.cause) : String(err.cause);
    }
    logs.error(
      'proxy',
      `Full error details: ${err.stack}${err.cause ? '\nCause: ' + causeDetails : ''}`,
    );
    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 = new Headers();
      if (request.headers.has('cookie')) wsHeaders.set('Cookie', request.headers.get('cookie'));
      if (request.headers.has('authorization'))
        wsHeaders.set('Authorization', request.headers.get('authorization'));
      if (request.headers.has('origin')) wsHeaders.set('Origin', request.headers.get('origin'));
      if (request.headers.has('sec-websocket-protocol'))
        wsHeaders.set('Sec-WebSocket-Protocol', request.headers.get('sec-websocket-protocol'));
      if (request.headers.has('sec-websocket-extensions'))
        wsHeaders.set('Sec-WebSocket-Extensions', request.headers.get('sec-websocket-extensions'));

      let upstream;
      try {
        // Convert Headers object to a plain object for the WebSocket constructor
        const plainWsHeaders = {};
        for (const [key, value] of wsHeaders) {
          plainWsHeaders[key] = value;
        }
        upstream = await connectUpstreamWebSocket(targetUrl.toString(), plainWsHeaders);
      } 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');
}