Просмотр исходного кода

fix(proxy): preserve User-Agent for gemini providers

fix #515
Ding 1 месяц назад
Родитель
Сommit
a4797187
2 измененных файлов с 218 добавлено и 21 удалено
  1. 55 21
      src/app/v1/_lib/proxy/forwarder.ts
  2. 163 0
      tests/unit/proxy/proxy-forwarder.test.ts

+ 55 - 21
src/app/v1/_lib/proxy/forwarder.ts

@@ -848,26 +848,6 @@ export class ProxyForwarder {
       const accessToken = await GeminiAuth.getAccessToken(provider.key);
       const isApiKey = GeminiAuth.isApiKey(provider.key);
 
-      const headers = new Headers();
-      headers.set("Content-Type", "application/json");
-
-      // ⭐ 统一禁用 gzip 压缩(不仅限于流式请求)
-      // 原因:undici(Node.js fetch)在连接提前关闭时会对不完整的 gzip 流抛出 "TypeError: terminated"
-      // Bun 的 fetch 实现更宽松,不会报错,这导致 bun dev 正常但 Docker 构建后失败
-      // 参考:Gunzip.emit → emitErrorNT → emitErrorCloseNT 错误链
-      headers.set("accept-encoding", "identity");
-
-      if (isApiKey) {
-        headers.set(GEMINI_PROTOCOL.HEADERS.API_KEY, accessToken);
-      } else {
-        headers.set("Authorization", `Bearer ${accessToken}`);
-      }
-
-      // CLI specific headers
-      if (provider.providerType === "gemini-cli") {
-        headers.set(GEMINI_PROTOCOL.HEADERS.API_CLIENT, "GeminiCLI/1.0");
-      }
-
       // 3. 直接透传:使用 buildProxyUrl() 拼接原始路径和查询参数
       const baseUrl =
         provider.url ||
@@ -876,7 +856,16 @@ export class ProxyForwarder {
           : GEMINI_PROTOCOL.CLI_ENDPOINT);
 
       proxyUrl = buildProxyUrl(baseUrl, session.requestUrl);
-      processedHeaders = headers;
+
+      // 4. Headers 处理:默认透传 session.headers(含请求过滤器修改),但移除代理认证头并覆盖上游鉴权
+      // 说明:之前 Gemini 分支使用 new Headers() 重建 headers,会导致 user-agent 丢失且过滤器不生效
+      processedHeaders = ProxyForwarder.buildGeminiHeaders(
+        session,
+        provider,
+        baseUrl,
+        accessToken,
+        isApiKey
+      );
 
       if (session.sessionId) {
         void SessionManager.storeSessionUpstreamRequestMeta(
@@ -1809,6 +1798,51 @@ export class ProxyForwarder {
     return headerProcessor.process(session.headers);
   }
 
+  private static buildGeminiHeaders(
+    session: ProxySession,
+    provider: NonNullable<typeof session.provider>,
+    baseUrl: string,
+    accessToken: string,
+    isApiKey: boolean
+  ): Headers {
+    const preserveClientIp = provider.preserveClientIp ?? false;
+    const { clientIp, xForwardedFor } = ProxyForwarder.resolveClientIp(session.headers);
+
+    const overrides: Record<string, string> = {
+      host: HeaderProcessor.extractHost(baseUrl),
+      "content-type": "application/json",
+      "accept-encoding": "identity",
+      "user-agent": session.headers.get("user-agent") ?? session.userAgent ?? "claude-code-hub",
+    };
+
+    if (isApiKey) {
+      overrides[GEMINI_PROTOCOL.HEADERS.API_KEY] = accessToken;
+    } else {
+      overrides.authorization = `Bearer ${accessToken}`;
+    }
+
+    if (provider.providerType === "gemini-cli") {
+      overrides[GEMINI_PROTOCOL.HEADERS.API_CLIENT] = "GeminiCLI/1.0";
+    }
+
+    if (preserveClientIp) {
+      if (xForwardedFor) {
+        overrides["x-forwarded-for"] = xForwardedFor;
+      }
+      if (clientIp) {
+        overrides["x-real-ip"] = clientIp;
+      }
+    }
+
+    const headerProcessor = HeaderProcessor.createForProxy({
+      blacklist: ["content-length", "connection", "x-api-key", GEMINI_PROTOCOL.HEADERS.API_KEY],
+      preserveClientIpHeaders: preserveClientIp,
+      overrides,
+    });
+
+    return headerProcessor.process(session.headers);
+  }
+
   private static resolveClientIp(headers: Headers): {
     clientIp: string | null;
     xForwardedFor: string | null;

+ 163 - 0
tests/unit/proxy/proxy-forwarder.test.ts

@@ -59,6 +59,15 @@ function createCodexProvider(): Provider {
   } as unknown as Provider;
 }
 
+function createGeminiProvider(providerType: "gemini" | "gemini-cli"): Provider {
+  return {
+    providerType,
+    url: "https://generativelanguage.googleapis.com/v1beta",
+    key: "test-outbound-key",
+    preserveClientIp: false,
+  } as unknown as Provider;
+}
+
 describe("ProxyForwarder - buildHeaders User-Agent resolution", () => {
   it("应该优先使用过滤器修改的 user-agent(Codex provider)", () => {
     const session = createSession({
@@ -147,3 +156,157 @@ describe("ProxyForwarder - buildHeaders User-Agent resolution", () => {
     expect(resultHeaders.get("user-agent")).toBe("");
   });
 });
+
+describe("ProxyForwarder - buildGeminiHeaders headers passthrough", () => {
+  it("应该透传 user-agent,并覆盖上游 x-goog-api-key(API key 模式)", () => {
+    const session = createSession({
+      userAgent: "Original-UA/1.0",
+      headers: new Headers([
+        ["user-agent", "Original-UA/1.0"],
+        ["x-api-key", "proxy-user-key-should-not-leak"],
+        ["x-goog-api-key", "proxy-user-key-google-should-not-leak"],
+        ["authorization", "Bearer proxy-user-bearer-should-not-leak"],
+      ]),
+    });
+
+    const provider = createGeminiProvider("gemini");
+    const { buildGeminiHeaders } = ProxyForwarder as unknown as {
+      buildGeminiHeaders: (
+        session: ProxySession,
+        provider: Provider,
+        baseUrl: string,
+        accessToken: string,
+        isApiKey: boolean
+      ) => Headers;
+    };
+    const resultHeaders = buildGeminiHeaders(
+      session,
+      provider,
+      "https://generativelanguage.googleapis.com/v1beta",
+      "upstream-api-key",
+      true
+    );
+
+    expect(resultHeaders.get("user-agent")).toBe("Original-UA/1.0");
+    expect(resultHeaders.get("x-api-key")).toBeNull();
+    expect(resultHeaders.get("authorization")).toBeNull();
+    expect(resultHeaders.get("x-goog-api-key")).toBe("upstream-api-key");
+    expect(resultHeaders.get("accept-encoding")).toBe("identity");
+    expect(resultHeaders.get("content-type")).toBe("application/json");
+    expect(resultHeaders.get("host")).toBe("generativelanguage.googleapis.com");
+  });
+
+  it("应该允许过滤器修改 user-agent(Gemini provider)", () => {
+    const session = createSession({
+      userAgent: "Original-UA/1.0",
+      headers: new Headers([["user-agent", "Filtered-UA/2.0"]]),
+    });
+    (session as any).originalHeaders = new Headers([["user-agent", "Original-UA/1.0"]]);
+
+    const provider = createGeminiProvider("gemini");
+    const { buildGeminiHeaders } = ProxyForwarder as unknown as {
+      buildGeminiHeaders: (
+        session: ProxySession,
+        provider: Provider,
+        baseUrl: string,
+        accessToken: string,
+        isApiKey: boolean
+      ) => Headers;
+    };
+    const resultHeaders = buildGeminiHeaders(
+      session,
+      provider,
+      "https://generativelanguage.googleapis.com/v1beta",
+      "upstream-api-key",
+      true
+    );
+
+    expect(resultHeaders.get("user-agent")).toBe("Filtered-UA/2.0");
+  });
+
+  it("应该在过滤器删除 user-agent 时回退到原始 userAgent(Gemini provider)", () => {
+    const session = createSession({
+      userAgent: "Original-UA/1.0",
+      headers: new Headers(), // user-agent 被删除
+    });
+    (session as any).originalHeaders = new Headers([["user-agent", "Original-UA/1.0"]]);
+
+    const provider = createGeminiProvider("gemini");
+    const { buildGeminiHeaders } = ProxyForwarder as unknown as {
+      buildGeminiHeaders: (
+        session: ProxySession,
+        provider: Provider,
+        baseUrl: string,
+        accessToken: string,
+        isApiKey: boolean
+      ) => Headers;
+    };
+    const resultHeaders = buildGeminiHeaders(
+      session,
+      provider,
+      "https://generativelanguage.googleapis.com/v1beta",
+      "upstream-api-key",
+      true
+    );
+
+    expect(resultHeaders.get("user-agent")).toBe("Original-UA/1.0");
+  });
+
+  it("应该使用 Authorization Bearer(OAuth 模式)并移除 x-goog-api-key", () => {
+    const session = createSession({
+      userAgent: "Original-UA/1.0",
+      headers: new Headers([
+        ["user-agent", "Original-UA/1.0"],
+        ["x-goog-api-key", "proxy-user-key-google-should-not-leak"],
+      ]),
+    });
+
+    const provider = createGeminiProvider("gemini");
+    const { buildGeminiHeaders } = ProxyForwarder as unknown as {
+      buildGeminiHeaders: (
+        session: ProxySession,
+        provider: Provider,
+        baseUrl: string,
+        accessToken: string,
+        isApiKey: boolean
+      ) => Headers;
+    };
+    const resultHeaders = buildGeminiHeaders(
+      session,
+      provider,
+      "https://generativelanguage.googleapis.com/v1beta",
+      "upstream-oauth-token",
+      false
+    );
+
+    expect(resultHeaders.get("authorization")).toBe("Bearer upstream-oauth-token");
+    expect(resultHeaders.get("x-goog-api-key")).toBeNull();
+  });
+
+  it("gemini-cli 应该注入 x-goog-api-client 头", () => {
+    const session = createSession({
+      userAgent: "Original-UA/1.0",
+      headers: new Headers([["user-agent", "Original-UA/1.0"]]),
+    });
+
+    const provider = createGeminiProvider("gemini-cli");
+    const { buildGeminiHeaders } = ProxyForwarder as unknown as {
+      buildGeminiHeaders: (
+        session: ProxySession,
+        provider: Provider,
+        baseUrl: string,
+        accessToken: string,
+        isApiKey: boolean
+      ) => Headers;
+    };
+    const resultHeaders = buildGeminiHeaders(
+      session,
+      provider,
+      "https://cloudcode-pa.googleapis.com/v1internal",
+      "upstream-oauth-token",
+      false
+    );
+
+    expect(resultHeaders.get("x-goog-api-client")).toBe("GeminiCLI/1.0");
+  });
+});