王明 / Food-Labeling-Management-Platform

import type { IncomingMessage, ServerResponse } from 'node:http';
import http from 'node:http';
import https from 'node:https';

import type { Connect, Plugin } from 'vite';

/** 仅转发这些请求头，避免重复 Content-Length / Accept-Encoding 导致上游异常断连 */
const FORWARD_HEADER_KEYS = [
  'content-type',
  'accept',
  'accept-language',
  'content-language',
  'clientid',
  'authorization',
  '__tenant',
];

function readRequestBody(req: IncomingMessage): Promise<Buffer> {
  return new Promise((resolve, reject) => {
    const chunks: Buffer[] = [];
    req.on('data', (chunk: Buffer) => chunks.push(chunk));
    req.on('end', () => resolve(Buffer.concat(chunks)));
    req.on('error', reject);
  });
}

function pickForwardHeaders(req: IncomingMessage): http.OutgoingHttpHeaders {
  const out: http.OutgoingHttpHeaders = {};
  for (const key of FORWARD_HEADER_KEYS) {
    const v =
      req.headers[key] ??
      req.headers[key.toUpperCase() as keyof typeof req.headers];
    if (v) {
      out[key] = Array.isArray(v) ? v[0] : v;
    }
  }
  return out;
}

function upstreamRequest(
  targetBase: string,
  pathSuffix: string,
  method: string,
  headers: http.OutgoingHttpHeaders,
  body?: Buffer,
): Promise<{ headers: http.IncomingHttpHeaders; statusCode: number; body: Buffer }> {
  if (!targetBase?.trim()) {
    return Promise.reject(
      new Error('VITE_APP_URL 未配置，请检查 .env.development'),
    );
  }

  const base = targetBase.replace(/\/$/, '');
  const path = pathSuffix.startsWith('/') ? pathSuffix : `/${pathSuffix}`;
  const fullUrl = new URL(`${base}${path}`);
  const lib = fullUrl.protocol === 'https:' ? https : http;

  return new Promise((resolve, reject) => {
    const req = lib.request(
      {
        protocol: fullUrl.protocol,
        hostname: fullUrl.hostname,
        port: fullUrl.port || (fullUrl.protocol === 'https:' ? 443 : 80),
        path: `${fullUrl.pathname}${fullUrl.search}`,
        method,
        headers: {
          ...headers,
          host: fullUrl.host,
          connection: 'close',
          accept: 'application/json',
          ...(body?.length ? { 'content-length': body.length } : {}),
        },
        agent: new lib.Agent({ keepAlive: false }),
      },
      (res) => {
        const chunks: Buffer[] = [];
        let settled = false;

        const finish = (err?: Error) => {
          if (settled) return;
          settled = true;
          const responseBody = Buffer.concat(chunks);
          if (!responseBody.length) {
            reject(err ?? new Error('上游返回空响应'));
            return;
          }
          resolve({
            body: responseBody,
            headers: res.headers,
            statusCode: res.statusCode ?? 502,
          });
        };

        res.on('data', (c: Buffer) => chunks.push(c));
        res.on('end', () => finish());
        // saas-test：chunked 未规范结束，curl 能收齐但 Node 等不到 end，在 close 时用已收到的 body 完成
        res.on('close', () => {
          if (!settled) {
            finish(
              chunks.length
                ? undefined
                : new Error('上游连接提前关闭且无响应体'),
            );
          }
        });
        res.on('error', (e) => {
          if (!settled) {
            finish(chunks.length ? undefined : e);
          }
        });
      },
    );

    req.on('error', (e) => reject(e));
    if (body?.length) {
      req.write(body);
    }
    req.end();
  });
}

function writeBufferedResponse(
  res: ServerResponse,
  statusCode: number,
  upstreamHeaders: http.IncomingHttpHeaders,
  body: Buffer,
) {
  res.statusCode = statusCode;
  const contentType =
    upstreamHeaders['content-type'] ?? 'application/json; charset=utf-8';
  res.setHeader('Content-Type', contentType);
  res.setHeader('Content-Length', String(body.length));
  res.setHeader('Connection', 'close');
  res.setHeader('X-Dev-Api-Buffer', '1');
  res.end(body);
}

function createBufferMiddleware(
  apiTarget: string,
  prefix: string,
): Connect.NextHandleFunction {
  return async (req, res, next) => {
    const url = req.url ?? '';
    if (!url.startsWith(prefix)) {
      next();
      return;
    }

    try {
      const pathSuffix = url.slice(prefix.length) || '/';
      const method = req.method ?? 'GET';
      const body =
        method === 'GET' || method === 'HEAD'
          ? undefined
          : await readRequestBody(req);
      const headers = pickForwardHeaders(req);

      const upstream = await upstreamRequest(
        apiTarget,
        pathSuffix,
        method,
        headers,
        body,
      );

      writeBufferedResponse(
        res,
        upstream.statusCode,
        upstream.headers,
        upstream.body,
      );
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      console.error('[th-dev-api-buffer]', req.method, req.url, message);
      const errBody = Buffer.from(
        JSON.stringify({
          statusCode: 502,
          data: null,
          succeeded: false,
          errors: [{ errorMessage: `dev 代理失败: ${message}` }],
        }),
      );
      writeBufferedResponse(res, 502, {}, errBody);
    }
  };
}

export function thDevApiBufferPlugin(
  apiTarget: string,
  devApiPrefix: string,
): Plugin {
  const prefix = devApiPrefix.startsWith('/')
    ? devApiPrefix
    : `/${devApiPrefix}`;

  return {
    name: 'th-dev-api-buffer',
    enforce: 'pre',
    configureServer(server) {
      if (!apiTarget?.trim()) {
        console.warn(
          '[th-dev-api-buffer] 警告: VITE_APP_URL 为空，/dev-api 请求将返回 502',
        );
      }
      const handler = createBufferMiddleware(apiTarget, prefix);
      return () => {
        const stack = server.middlewares.stack as Array<{
          route: string;
          handle: Connect.NextHandleFunction;
        }>;
        stack.unshift({ route: '', handle: handler });
      };
    },
  };
}