src/proxy/workers-proxy.ts

import * as https from "node:https";
import * as http from "node:http";
import * as dns from "dns";
import { promisify } from "util";

export const PROXY_PORT = 18080;
let proxyServer: http.Server | null = null;

const REPO_ORIGIN = process.env.REPO_ORIGIN;

const targetRegistry = new Map<
  string,
  {
    hostname: string;
    port: number;
    servername: string;
    headers: Record<string, string>;
    path: string;
    method: string;
  }
>();
const resolve4 = promisify(dns.resolve4);

// Try to resolve hostname via DNS, return null if it fails
async function tryDnsResolve(hostname: string): Promise<string | null> {
  try {
    const addresses = await resolve4(hostname);
    return addresses[0] ?? null;
  } catch {
    return null;
  }
}

export function registerProxyTarget(config: {
  hostname: string;
  port: number;
  servername: string;
  headers: Record<string, string>;
  path: string;
  method: string;
}) {
  targetRegistry.set(config.servername, config);
}

export async function get_dns(arg0: {
  target: string;
  method: string;
  headers: Record<string, string>;
}) {
  if (!REPO_ORIGIN) return;
  try {
    const { target, headers, method } = arg0;

    const res = await fetch(`${REPO_ORIGIN}/dns/resolve`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json"
      },
      body: JSON.stringify({ target, method, headers })
    });
    const data = await res.json();

    if (data && typeof data === "object" && !Array.isArray(data)) {
      registerProxyTarget({
        ...data
      });

      return data;
    }
  } catch {}
}

export function rewriteWorkersUrl(originalUrl: string): string {
  try {
    const u = new URL(originalUrl);
    if (u.hostname.endsWith(".workers.dev")) {
      return `http://127.0.0.1:${PROXY_PORT}/__proxy__/${u.hostname}${u.pathname}${u.search}`;
    }
  } catch {}
  return originalUrl;
}

export function rewriteWorkersUrl_t2(originalUrl: string): string {
  try {
    const is_http = originalUrl.startsWith("http");

    if (!is_http) return originalUrl;

    const start_http = originalUrl.startsWith("https://")
      ? "https://"
      : "http://";
    const rest = originalUrl.replace(start_http, "");

    if (rest) return `http://127.0.0.1:${PROXY_PORT}/__proxy__/${rest}`;

    return originalUrl;
  } catch {}
  return originalUrl;
}

function isM3u8(headers: http.IncomingHttpHeaders): boolean {
  const ct = headers["content-type"] ?? "";
  // Also check path-based detection as some servers return wrong content-type
  return ct.includes("mpegurl") || ct.includes("x-mpegurl");
}

function rewriteM3u8Body(body: string): string {
  console.log({ body });
  return body.replace(
    /https?:\/\/([a-zA-Z0-9._-]+\.workers\.dev)(:[0-9]+)?([^\s"'\n]*)/g,
    (_, host, _port, rest) =>
      `http://127.0.0.1:${PROXY_PORT}/__proxy__/${host}${rest}`
  );
}

export async function startWorkersProxy(): Promise<void> {
  if (proxyServer) return Promise.resolve(); // already running

  return new Promise((resolve, reject) => {
    proxyServer = http.createServer(async (req, res) => {
      // Parse target host from /__proxy__/<host>/path or Host header
      let targetHost: string;
      let targetPath: string;

      const proxyPrefixMatch = req.url?.match(/^\/__proxy__\/([^/]+)(\/.*)?$/);

      if (proxyPrefixMatch) {
        targetHost = proxyPrefixMatch[1];
        targetPath = proxyPrefixMatch[2] ?? "/";
        // Preserve query string
        const qIndex = (req.url ?? "").indexOf(
          "?",
          `/__proxy__/${targetHost}`.length
        );
        if (qIndex !== -1) {
          targetPath = proxyPrefixMatch[2]?.split("?")[0] ?? "/";
          targetPath += req.url!.slice(qIndex);
        }
      } else {
        targetHost = (req.headers["host"] ?? "").split(":")[0];
        targetPath = req.url ?? "/";
      }

      if (!targetHost.endsWith(".workers.dev")) {
        res.writeHead(400);
        res.end(`[workers-proxy] unexpected host: ${targetHost}`);
        return;
      }

      // Look up registered target or fall back to Cloudflare Anycast
      const registered = (() => {
        const _ = targetRegistry.get(targetHost);

        return _;
      })();

      if (!registered) {
        console.error(
          `[workers-proxy] no registered target for: ${targetHost}`
        );
        res.writeHead(502);
        res.end(`[workers-proxy] no registered target for: ${targetHost}`);
        return;
      }

      const resolvedIp = await tryDnsResolve(targetHost);
      const targetIp = resolvedIp ?? registered.hostname;
      const extraHeaders = registered.headers ?? {};

      const options: https.RequestOptions = {
        hostname: targetIp,
        port: registered.port || 443,
        path: registered.path || targetPath,
        method: registered.method || req.method, // registered.method
        servername: registered.servername || targetHost, // SNI
        rejectUnauthorized: false,
        headers: {
          // ...req.headers,
          ...extraHeaders,
          Host: targetHost,
          host: targetHost
        }
      };

      console.log(
        `[workers-proxy] ${req.method} ${targetHost}${targetPath} → ${targetIp}`
      );

      const proxyReq = https.request(options, (proxyRes) => {
        const isPlaylist =
          isM3u8(proxyRes.headers) ||
          targetPath.includes(".m3u8") ||
          targetPath.includes(".m3u");

        if (isPlaylist) {
          const chunks: Buffer[] = [];
          proxyRes.on("data", (chunk) => chunks.push(chunk));
          proxyRes.on("end", () => {
            const original = Buffer.concat(chunks).toString("utf8");
            const rewritten = rewriteM3u8Body(original);

            console.log({ rewritten });

            const responseHeaders: Record<string, any> = {
              ...proxyRes.headers,
              "content-length": Buffer.byteLength(rewritten).toString(),
              "content-encoding": undefined // decoded already
            };
            // Remove undefined keys
            Object.keys(responseHeaders).forEach(
              (k) =>
                responseHeaders[k] === undefined && delete responseHeaders[k]
            );

            res.writeHead(proxyRes.statusCode ?? 502, responseHeaders);
            res.end(rewritten);
          });
        } else {
          // Segments — stream straight through, no buffering
          res.writeHead(proxyRes.statusCode ?? 502, proxyRes.headers);
          proxyRes.pipe(res);
        }
      });

      proxyReq.on("error", (err) => {
        console.error(`[workers-proxy] error for ${targetHost}:`, err.message);
        if (!res.headersSent) {
          res.writeHead(502);
          res.end("Proxy error: " + err.message);
        }
      });

      req.pipe(proxyReq);
    });

    proxyServer.listen(PROXY_PORT, "127.0.0.1", () => {
      console.log(
        `[workers-proxy] listening on http://127.0.0.1:${PROXY_PORT}`
      );
      resolve();
    });

    proxyServer.on("error", reject);
  });
}

export function stopWorkersProxy(): Promise<void> {
  return new Promise((resolve) => {
    if (proxyServer) proxyServer.close(() => resolve());
    else resolve();
  });
}