CF-Workers 搭建gemini中转+key号池

本地可以随机输入api密钥，自动轮询调用号池内的API KEY。

import { connect } from "cloudflare:sockets";

// 定义API密钥池
const API_KEYS = [
  'AIzaSy..............',
  'AIzaSy..............',
  'AIzaSy..............',
  'AIzaSy..............'
];

// 用于存储当前密钥索引的全局变量
let currentKeyIndex = 0;

// 全局配置，包括目标 URL 和调试模式
const CONFIG = {
  AUTH_TOKEN: "optional-auth-token", // 可选的认证 token，若不需要可留空或删除相关逻辑
  DEFAULT_DST_URL: "https://generativelanguage.googleapis.com", // 目标 API 地址
  DEBUG_MODE: false, // 调试模式，设为 true 可查看日志
};

// 从环境变量更新全局配置（优先使用环境变量值）
function updateConfigFromEnv(env) {
  if (!env) return;
  for (const key of Object.keys(CONFIG)) {
    if (key in env) {
      if (typeof CONFIG[key] === 'boolean') {
        CONFIG[key] = env[key] === 'true';
      } else {
        CONFIG[key] = env[key];
      }
    }
  }
}

// 定义文本编码器和解码器，用于字符串与字节数组的转换
const encoder = new TextEncoder();
const decoder = new TextDecoder();

// 过滤不应该转发的 HTTP 头部（如 host、accept-encoding、cf-*）
const HEADER_FILTER_RE = /^(host|accept-encoding|cf-)/i;

// 根据调试模式定义日志输出函数
const log = CONFIG.DEBUG_MODE
  ? (message, data = "") => console.log(`[DEBUG] ${message}`, data)
  : () => {};

// 拼接多个 Uint8Array
function concatUint8Arrays(...arrays) {
  const total = arrays.reduce((sum, arr) => sum + arr.length, 0);
  const result = new Uint8Array(total);
  let offset = 0;
  for (const arr of arrays) {
    result.set(arr, offset);
    offset += arr.length;
  }
  return result;
}

// 解析 HTTP 响应头部，返回状态码、状态文本、头部和头部结束位置
function parseHttpHeaders(buff) {
  const text = decoder.decode(buff);
  const headerEnd = text.indexOf("\r\n\r\n");
  if (headerEnd === -1) return null;
  const headerSection = text.slice(0, headerEnd).split("\r\n");
  const statusLine = headerSection[0];
  const statusMatch = statusLine.match(/HTTP\/1\.[01] (\d+) (.*)/);
  if (!statusMatch) throw new Error(`Invalid status line: ${statusLine}`);
  const headers = new Headers();
  for (let i = 1; i < headerSection.length; i++) {
    const line = headerSection[i];
    const idx = line.indexOf(": ");
    if (idx !== -1) {
      headers.append(line.slice(0, idx), line.slice(idx + 2));
    }
  }
  return { status: Number(statusMatch[1]), statusText: statusMatch[2], headers, headerEnd };
}

// 读取数据直到遇到双 CRLF（HTTP 头部结束）
async function readUntilDoubleCRLF(reader) {
  let respText = "";
  while (true) {
    const { value, done } = await reader.read();
    if (value) {
      respText += decoder.decode(value, { stream: true });
      if (respText.includes("\r\n\r\n")) break;
    }
    if (done) break;
  }
  return respText;
}

// 异步生成器：读取 chunked 编码的 HTTP 响应数据块
async function* readChunks(reader, buff = new Uint8Array()) {
  while (true) {
    let pos = -1;
    for (let i = 0; i < buff.length - 1; i++) {
      if (buff[i] === 13 && buff[i + 1] === 10) {
        pos = i;
        break;
      }
    }
    if (pos === -1) {
      const { value, done } = await reader.read();
      if (done) break;
      buff = concatUint8Arrays(buff, value);
      continue;
    }
    const size = parseInt(decoder.decode(buff.slice(0, pos)), 16);
    log("Read chunk size", size);
    if (!size) break;
    buff = buff.slice(pos + 2);
    while (buff.length < size + 2) {
      const { value, done } = await reader.read();
      if (done) throw new Error("Unexpected EOF in chunked encoding");
      buff = concatUint8Arrays(buff, value);
    }
    yield buff.slice(0, size);
    buff = buff.slice(size + 2);
  }
}

// 解析完整的 HTTP 响应，处理 chunked 或固定长度的响应体
async function parseResponse(reader) {
  let buff = new Uint8Array();
  while (true) {
    const { value, done } = await reader.read();
    if (value) {
      buff = concatUint8Arrays(buff, value);
      const parsed = parseHttpHeaders(buff);
      if (parsed) {
        const { status, statusText, headers, headerEnd } = parsed;
        const isChunked = headers.get("transfer-encoding")?.includes("chunked");
        const contentLength = parseInt(headers.get("content-length") || "0", 10);
        const data = buff.slice(headerEnd + 4);
        return new Response(
          new ReadableStream({
            async start(ctrl) {
              try {
                if (isChunked) {
                  log("Using chunked transfer mode");
                  for await (const chunk of readChunks(reader, data)) {
                    ctrl.enqueue(chunk);
                  }
                } else {
                  log("Using fixed-length transfer mode", { contentLength });
                  let received = data.length;
                  if (data.length) ctrl.enqueue(data);
                  while (received < contentLength) {
                    const { value, done } = await reader.read();
                    if (done) break;
                    received += value.length;
                    ctrl.enqueue(value);
                  }
                }
                ctrl.close();
              } catch (err) {
                log("Error parsing response", err);
                ctrl.error(err);
              }
            },
          }),
          { status, statusText, headers }
        );
      }
    }
    if (done) break;
  }
  throw new Error("Unable to parse response headers");
}

// 处理 HTTP 请求的转发
async function nativeFetch(req, dstUrl) {
  // 清理头部，过滤不需要的头部
  const cleanedHeaders = new Headers();
  for (const [k, v] of req.headers) {
    if (!HEADER_FILTER_RE.test(k)) {
      cleanedHeaders.set(k, v);
    }
  }

  // 从API密钥池中按顺序选择一个密钥
  const apiKey = API_KEYS[currentKeyIndex];
  // 更新索引，循环使用密钥池
  currentKeyIndex = (currentKeyIndex + 1) % API_KEYS.length;

  const targetUrl = new URL(dstUrl);
  // 将API密钥附加到查询参数中
  targetUrl.searchParams.set('key', apiKey);
  
  cleanedHeaders.set("Host", targetUrl.hostname);
  cleanedHeaders.set("accept-encoding", "identity"); // 禁用压缩

  const port = targetUrl.protocol === "https:" ? 443 : 80;
  const socket = await connect(
    { hostname: targetUrl.hostname, port },
    { secureTransport: targetUrl.protocol === "https:" ? "on" : "off" }
  );
  const writer = socket.writable.getWriter();

  // 构造请求行和头部
  const requestLine =
    `${req.method} ${targetUrl.pathname}${targetUrl.search} HTTP/1.1\r\n` +
    Array.from(cleanedHeaders.entries())
      .map(([k, v]) => `${k}: ${v}`)
      .join("\r\n") +
    "\r\n\r\n";
  log("Sending request", requestLine);
  await writer.write(encoder.encode(requestLine));

  // 转发请求体（如果有）
  if (req.body) {
    log("Forwarding request body");
    for await (const chunk of req.body) {
      await writer.write(chunk);
    }
  }

  // 解析并返回目标服务器的响应
  return await parseResponse(socket.readable.getReader());
}

// 请求处理入口
async function handleRequest(req, env) {
  updateConfigFromEnv(env);
  CONFIG.DEBUG_MODE = CONFIG.DEBUG_MODE;

  const url = new URL(req.url);
  const path = url.pathname + url.search; // 获取路径和查询参数
  const targetUrl = `${CONFIG.DEFAULT_DST_URL}${path.startsWith('/') ? path.slice(1) : path}`;

  log("Target URL", targetUrl);
  return await nativeFetch(req, targetUrl);
}

// 导出 Cloudflare Workers 的 fetch 事件处理器
export default { fetch: handleRequest };
export const onRequest = (ctx) => handleRequest(ctx.request, ctx.env);