http_upstream.go 53 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565
  1. package repository
  2. import (
  3. "bufio"
  4. "bytes"
  5. "compress/flate"
  6. "compress/gzip"
  7. "context"
  8. "crypto/sha256"
  9. "crypto/tls"
  10. "encoding/json"
  11. "errors"
  12. "fmt"
  13. "io"
  14. "log/slog"
  15. "net"
  16. "net/http"
  17. "net/url"
  18. "os"
  19. "strings"
  20. "sync"
  21. "sync/atomic"
  22. "time"
  23. "github.com/andybalholm/brotli"
  24. "github.com/klauspost/compress/zstd"
  25. "golang.org/x/mod/semver"
  26. "golang.org/x/net/http2"
  27. "github.com/Wei-Shaw/sub2api/internal/config"
  28. "github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
  29. "github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
  30. "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
  31. "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
  32. "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
  33. "github.com/Wei-Shaw/sub2api/internal/service"
  34. "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
  35. )
  36. // 默认配置常量
  37. // 这些值在配置文件未指定时作为回退默认值使用
  38. const (
  39. // directProxyKey: 无代理时的缓存键标识
  40. directProxyKey = "direct"
  41. // defaultMaxIdleConns: 默认最大空闲连接总数
  42. // HTTP/2 场景下,单连接可多路复用,240 足以支撑高并发
  43. defaultMaxIdleConns = 240
  44. // defaultMaxIdleConnsPerHost: 默认每主机最大空闲连接数
  45. defaultMaxIdleConnsPerHost = 120
  46. // defaultMaxConnsPerHost: 默认每主机最大连接数(含活跃连接)
  47. // 达到上限后新请求会等待,而非无限创建连接
  48. defaultMaxConnsPerHost = 240
  49. // defaultIdleConnTimeout: 默认空闲连接超时时间(90秒)
  50. // 超时后连接会被关闭,释放系统资源(建议小于上游 LB 超时)
  51. defaultIdleConnTimeout = 90 * time.Second
  52. // defaultResponseHeaderTimeout: 默认等待响应头超时时间(5分钟)
  53. // LLM 请求可能排队较久,需要较长超时
  54. defaultResponseHeaderTimeout = 300 * time.Second
  55. // defaultUpstreamDialTimeout: 默认 TCP/DNS 建连超时(10秒)
  56. // Transport 不设置 DialContext 时会退化为零值 net.Dialer(无超时),建连阶段
  57. // 只能依赖内核默认 TCP 重传(Linux 约 130 秒)。ResponseHeaderTimeout 只约束
  58. // 连接建立之后等待响应头的阶段,覆盖不到 DNS 解析与 TCP 握手。
  59. // 上游域名被解析到 443 不可达的 IP 时(DNS 污染/路由异常),单个账号就要卡满
  60. // 内核超时;而多账号故障转移是串行的,一次请求会阻塞数分钟且不写中间错误。
  61. defaultUpstreamDialTimeout = 10 * time.Second
  62. // defaultUpstreamDialKeepAlive: TCP keepalive 探测间隔,与 Go 默认值保持一致
  63. defaultUpstreamDialKeepAlive = 30 * time.Second
  64. // defaultUpstreamTLSHandshakeTimeout: TLS 握手超时(10秒)
  65. // 与建连超时同量级,避免 TCP 已连通但对端不推进握手时无限等待
  66. defaultUpstreamTLSHandshakeTimeout = 10 * time.Second
  67. // defaultMaxUpstreamClients: 默认最大客户端缓存数量
  68. // 超出后会淘汰最久未使用的客户端
  69. defaultMaxUpstreamClients = 5000
  70. // defaultClientIdleTTLSeconds: 默认客户端空闲回收阈值(15分钟)
  71. defaultClientIdleTTLSeconds = 900
  72. // OpenAI HTTP/2 代理回退策略默认值
  73. defaultOpenAIHTTP2FallbackErrorThreshold = 2
  74. defaultOpenAIHTTP2FallbackWindow = 60 * time.Second
  75. defaultOpenAIHTTP2FallbackTTL = 10 * time.Minute
  76. // 长流 HTTP/2 连接健康探测:池化连接被代理/NAT
  77. // 静默掐断会成为“死连接”(两端都以为存活),请求落上去会挂到 TCP 重传超时
  78. // (分钟级)。Go 的 http2.Transport 默认 ReadIdleTimeout=0(不发健康 PING),
  79. // 无法检测。启用主动 PING 探测:连接空闲 ReadIdleTimeout 后发 PING,PingTimeout
  80. // 内无响应即判定死连接并关闭,从源头避免请求挂在死连接上。
  81. longStreamHTTP2ReadIdleTimeout = 10 * time.Second
  82. longStreamHTTP2PingTimeout = 5 * time.Second
  83. // The Grok CLI proxy rejects requests that do not identify a supported
  84. // client version. Host/env/version pins live in package xai so service,
  85. // billing, and transport layers advertise the same identity.
  86. grokCLIProxyHost = xai.CLIProxyHost
  87. grokOfficialAPIHost = "api.x.ai"
  88. grokCLIStableVersion = xai.CLIClientVersion // preferred pin (not the minimum floor)
  89. grokCLIVersionOverride = xai.CLIVersionEnv
  90. grokFallbackBodyLimit = 64 << 10
  91. )
  92. const (
  93. upstreamProtocolModeDefault = "default"
  94. upstreamProtocolModeLongStreamH2 = "long_stream_h2"
  95. upstreamProtocolModeOpenAIH1 = "openai_h1"
  96. upstreamProtocolModeOpenAIH2 = "openai_h2"
  97. upstreamProtocolModeOpenAIH1Fallback = "openai_h1_fallback"
  98. upstreamProtocolModeGrok = "grok"
  99. )
  100. var errUpstreamClientLimitReached = errors.New("upstream client cache limit reached")
  101. // poolSettings 连接池配置参数
  102. // 封装 Transport 所需的各项连接池参数
  103. type poolSettings struct {
  104. maxIdleConns int // 最大空闲连接总数
  105. maxIdleConnsPerHost int // 每主机最大空闲连接数
  106. maxConnsPerHost int // 每主机最大连接数(含活跃)
  107. idleConnTimeout time.Duration // 空闲连接超时时间
  108. responseHeaderTimeout time.Duration // 等待响应头超时时间
  109. }
  110. type openAIHTTP2Settings struct {
  111. enabled bool
  112. allowProxyFallbackToHTTP1 bool
  113. fallbackErrorThreshold int
  114. fallbackWindow time.Duration
  115. fallbackTTL time.Duration
  116. }
  117. // upstreamClientEntry 上游客户端缓存条目
  118. // 记录客户端实例及其元数据,用于连接池管理和淘汰策略
  119. type upstreamClientEntry struct {
  120. client *http.Client // HTTP 客户端实例
  121. proxyKey string // 代理标识(用于检测代理变更)
  122. poolKey string // 连接池配置标识(用于检测配置变更)
  123. protocolMode string // 协议模式(default/long_stream_h2/openai_h1/openai_h2/openai_h1_fallback)
  124. lastUsed int64 // 最后使用时间戳(纳秒),用于 LRU 淘汰
  125. inFlight int64 // 当前进行中的请求数,>0 时不可淘汰
  126. }
  127. type openAIHTTP2FallbackState struct {
  128. mu sync.Mutex
  129. windowStart time.Time
  130. errorCount int
  131. fallbackUntil time.Time
  132. }
  133. // httpUpstreamService 通用 HTTP 上游服务
  134. // 用于向任意 HTTP API(Claude、OpenAI 等)发送请求,支持可选代理
  135. //
  136. // 架构设计:
  137. // - 根据隔离策略(proxy/account/account_proxy)缓存客户端实例
  138. // - 每个客户端拥有独立的 Transport 连接池
  139. // - 支持 LRU + 空闲时间双重淘汰策略
  140. //
  141. // 性能优化:
  142. // 1. 根据隔离策略缓存客户端实例,避免频繁创建 http.Client
  143. // 2. 复用 Transport 连接池,减少 TCP 握手和 TLS 协商开销
  144. // 3. 支持账号级隔离与空闲回收,降低连接层关联风险
  145. // 4. 达到最大连接数后等待可用连接,而非无限创建
  146. // 5. 仅回收空闲客户端,避免中断活跃请求
  147. // 6. HTTP/2 多路复用,连接上限不等于并发请求上限
  148. // 7. 代理变更时清空旧连接池,避免复用错误代理
  149. // 8. 账号并发数与连接池上限对应(账号隔离策略下)
  150. type httpUpstreamService struct {
  151. traffic *service.AccountTrafficService
  152. cfg *config.Config // 全局配置
  153. mu sync.RWMutex // 保护 clients map 的读写锁
  154. clients map[string]*upstreamClientEntry // 客户端缓存池,key 由隔离策略决定
  155. // OpenAI 走 HTTP/HTTPS 代理时的 H2->H1 回退状态(key=标准化 proxyKey)
  156. openAIHTTP2Fallbacks sync.Map
  157. }
  158. // NewHTTPUpstream 创建通用 HTTP 上游服务
  159. // 使用配置中的连接池参数构建 Transport
  160. //
  161. // 参数:
  162. // - cfg: 全局配置,包含连接池参数和隔离策略
  163. //
  164. // 返回:
  165. // - service.HTTPUpstream 接口实现
  166. func NewHTTPUpstream(cfg *config.Config) service.HTTPUpstream {
  167. return &httpUpstreamService{
  168. cfg: cfg,
  169. clients: make(map[string]*upstreamClientEntry),
  170. }
  171. }
  172. func NewControlledHTTPUpstream(cfg *config.Config, traffic *service.AccountTrafficService) service.HTTPUpstream {
  173. s := NewHTTPUpstream(cfg).(*httpUpstreamService)
  174. s.traffic = traffic
  175. return s
  176. }
  177. func (s *httpUpstreamService) AccountTrafficController() *service.AccountTrafficService {
  178. return s.traffic
  179. }
  180. // Do 执行 HTTP 请求
  181. // 根据隔离策略获取或创建客户端,并跟踪请求生命周期
  182. //
  183. // 参数:
  184. // - req: HTTP 请求对象
  185. // - proxyURL: 代理地址,空字符串表示直连
  186. // - accountID: 账户 ID,用于账户级隔离
  187. // - accountConcurrency: 账户并发限制,用于动态调整连接池大小
  188. //
  189. // 返回:
  190. // - *http.Response: HTTP 响应(Body 已包装,关闭时自动更新计数)
  191. // - error: 请求错误
  192. //
  193. // 注意:
  194. // - 调用方必须关闭 resp.Body,否则会导致 inFlight 计数泄漏
  195. // - inFlight > 0 的客户端不会被淘汰,确保活跃请求不被中断
  196. func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
  197. return s.traffic.DoHTTP(req, func(controlled *http.Request) (*http.Response, error) {
  198. return s.doUncontrolled(controlled, proxyURL, accountID, accountConcurrency)
  199. })
  200. }
  201. func (s *httpUpstreamService) doUncontrolled(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
  202. applyGrokCLIProxyHeaders(req)
  203. if err := s.validateRequestHost(req); err != nil {
  204. return nil, err
  205. }
  206. profile := service.HTTPUpstreamProfileDefault
  207. if req != nil {
  208. profile = service.HTTPUpstreamProfileFromContext(req.Context())
  209. }
  210. // 获取或创建对应的客户端,并标记请求占用
  211. entry, err := s.acquireClientWithProfile(proxyURL, accountID, accountConcurrency, profile)
  212. if err != nil {
  213. return nil, err
  214. }
  215. // 执行请求
  216. client := s.httpClientForUpstreamRequest(entry.client, req)
  217. client = httpClientWithGrokAccessDeniedFallback(client)
  218. resp, err := servertiming.Do(client, req)
  219. if err != nil {
  220. s.recordOpenAIHTTP2Failure(profile, entry.protocolMode, entry.proxyKey, err)
  221. // 请求失败,立即减少计数
  222. atomic.AddInt64(&entry.inFlight, -1)
  223. atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
  224. return nil, err
  225. }
  226. s.recordOpenAIHTTP2Success(profile, entry.protocolMode, entry.proxyKey)
  227. // 如果上游返回了压缩内容,解压后再交给业务层
  228. decompressResponseBody(resp)
  229. // 包装响应体,在关闭时自动减少计数并更新时间戳
  230. // 这确保了流式响应(如 SSE)在完全读取前不会被淘汰
  231. resp.Body = wrapTrackedBody(resp.Body, func() {
  232. atomic.AddInt64(&entry.inFlight, -1)
  233. atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
  234. })
  235. return resp, nil
  236. }
  237. // DoWithTLS 执行带 TLS 指纹伪装的 HTTP 请求
  238. //
  239. // profile 为 nil 时不启用 TLS 指纹,行为与 Do 方法相同。
  240. // profile 非 nil 时使用指定的 Profile 进行 TLS 指纹伪装。
  241. func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile) (*http.Response, error) {
  242. return s.traffic.DoHTTP(req, func(controlled *http.Request) (*http.Response, error) {
  243. return s.doWithTLSUncontrolled(controlled, proxyURL, accountID, accountConcurrency, profile)
  244. })
  245. }
  246. func (s *httpUpstreamService) doWithTLSUncontrolled(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile) (*http.Response, error) {
  247. if profile == nil {
  248. return s.Do(req, proxyURL, accountID, accountConcurrency)
  249. }
  250. // An explicitly selected TLS profile must not silently take an unprotected
  251. // route, including redirects to HTTP or a plaintext upstream URL.
  252. if req == nil || req.URL == nil || !strings.EqualFold(req.URL.Scheme, "https") {
  253. return nil, fmt.Errorf("TLS fingerprint request requires HTTPS")
  254. }
  255. applyGrokCLIProxyHeaders(req)
  256. upstreamProfile := service.HTTPUpstreamProfileDefault
  257. if req != nil {
  258. upstreamProfile = service.HTTPUpstreamProfileFromContext(req.Context())
  259. }
  260. targetHost := ""
  261. if req != nil && req.URL != nil {
  262. targetHost = req.URL.Host
  263. }
  264. proxyInfo := "direct"
  265. if proxyURL != "" {
  266. if parsed, parseErr := url.Parse(proxyURL); parseErr == nil {
  267. proxyInfo = parsed.Redacted()
  268. } else {
  269. proxyInfo = "configured"
  270. }
  271. }
  272. slog.Debug("tls_fingerprint_enabled", "account_id", accountID, "target", targetHost, "proxy", proxyInfo, "profile", profile.Name)
  273. if err := s.validateRequestHost(req); err != nil {
  274. return nil, err
  275. }
  276. entry, err := s.acquireClientWithTLS(proxyURL, accountID, accountConcurrency, profile, upstreamProfile)
  277. if err != nil {
  278. slog.Debug("tls_fingerprint_acquire_client_failed", "account_id", accountID, "error", err)
  279. return nil, err
  280. }
  281. client := s.httpClientForUpstreamRequest(entry.client, req)
  282. client = httpClientWithGrokAccessDeniedFallback(client)
  283. resp, err := servertiming.Do(client, req)
  284. if err != nil {
  285. atomic.AddInt64(&entry.inFlight, -1)
  286. atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
  287. slog.Debug("tls_fingerprint_request_failed", "account_id", accountID, "error", err)
  288. return nil, err
  289. }
  290. decompressResponseBody(resp)
  291. resp.Body = wrapTrackedBody(resp.Body, func() {
  292. atomic.AddInt64(&entry.inFlight, -1)
  293. atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
  294. })
  295. return resp, nil
  296. }
  297. // httpClientForUpstreamRequest 按请求上下文的标记派生客户端:禁用重定向,或对重定向的每一跳做主机校验。
  298. // 派生的克隆与缓存客户端共享 Transport;未打标记时原样返回。
  299. func (s *httpUpstreamService) httpClientForUpstreamRequest(client *http.Client, req *http.Request) *http.Client {
  300. if client == nil || req == nil {
  301. return client
  302. }
  303. ctx := req.Context()
  304. switch {
  305. case service.HTTPUpstreamRedirectsDisabled(ctx):
  306. clone := *client
  307. clone.CheckRedirect = func(*http.Request, []*http.Request) error {
  308. return http.ErrUseLastResponse
  309. }
  310. return &clone
  311. case service.HTTPUpstreamPublicHostsOnly(ctx) && client.CheckRedirect == nil:
  312. clone := *client
  313. clone.CheckRedirect = s.redirectChecker
  314. return &clone
  315. default:
  316. return client
  317. }
  318. }
  319. // grokAccessDeniedFallbackTransport preserves the subscription CLI proxy as
  320. // the primary OAuth route, but retries a replayable request against api.x.ai
  321. // when the proxy returns its compatibility-specific 403 "Access denied".
  322. // Trial subscriptions can hit this boundary while the same OAuth credential
  323. // remains valid on the official API. Other entitlement failures stay on the
  324. // original response so account scheduling semantics do not change.
  325. type grokAccessDeniedFallbackTransport struct {
  326. base http.RoundTripper
  327. }
  328. func httpClientWithGrokAccessDeniedFallback(client *http.Client) *http.Client {
  329. if client == nil {
  330. return nil
  331. }
  332. clone := *client
  333. base := clone.Transport
  334. if base == nil {
  335. base = http.DefaultTransport
  336. }
  337. clone.Transport = &grokAccessDeniedFallbackTransport{base: base}
  338. return &clone
  339. }
  340. func (t *grokAccessDeniedFallbackTransport) RoundTrip(req *http.Request) (*http.Response, error) {
  341. resp, err := t.base.RoundTrip(req)
  342. if err != nil || !isGrokCLIAccessDeniedFallbackCandidate(req, resp) {
  343. return resp, err
  344. }
  345. body, ok := bufferSmallResponseBody(resp, grokFallbackBodyLimit)
  346. if !ok || !isGrokCLICompatibilityAccessDenied(body) {
  347. return resp, nil
  348. }
  349. fallbackReq, err := newGrokOfficialAPIFallbackRequest(req)
  350. if err != nil {
  351. return resp, nil
  352. }
  353. fallbackResp, fallbackErr := t.base.RoundTrip(fallbackReq)
  354. if fallbackErr != nil {
  355. slog.Debug("grok_cli_access_denied_api_fallback_failed", "path", req.URL.EscapedPath(), "error", fallbackErr)
  356. return resp, nil
  357. }
  358. if fallbackResp.StatusCode < http.StatusOK || fallbackResp.StatusCode >= http.StatusMultipleChoices {
  359. if fallbackResp.Body != nil {
  360. _ = fallbackResp.Body.Close()
  361. }
  362. return resp, nil
  363. }
  364. if resp.Body != nil {
  365. _ = resp.Body.Close()
  366. }
  367. slog.Warn("grok_cli_access_denied_api_fallback_succeeded", "method", req.Method, "path", req.URL.EscapedPath())
  368. return fallbackResp, nil
  369. }
  370. func isGrokCLICompatibilityAccessDenied(body []byte) bool {
  371. lower := bytes.ToLower(body)
  372. if bytes.Contains(lower, []byte("access denied")) {
  373. return true
  374. }
  375. var payload struct {
  376. Code string `json:"code"`
  377. Error string `json:"error"`
  378. }
  379. if err := json.Unmarshal(body, &payload); err != nil || !strings.EqualFold(strings.TrimSpace(payload.Code), "permission_denied") {
  380. return false
  381. }
  382. const chatEndpointDeniedPrefix = "access to the chat endpoint is denied. please ensure you're using the correct credentials. if you believe this is a mistake, please"
  383. return strings.HasPrefix(strings.ToLower(strings.TrimSpace(payload.Error)), chatEndpointDeniedPrefix)
  384. }
  385. func isGrokCLIAccessDeniedFallbackCandidate(req *http.Request, resp *http.Response) bool {
  386. return req != nil && req.URL != nil && req.GetBody != nil && resp != nil &&
  387. resp.StatusCode == http.StatusForbidden &&
  388. strings.EqualFold(strings.TrimSpace(req.URL.Hostname()), grokCLIProxyHost) &&
  389. strings.EqualFold(strings.TrimSpace(req.Header.Get("X-XAI-Token-Auth")), "xai-grok-cli") &&
  390. strings.HasPrefix(strings.ToLower(strings.TrimSpace(req.Header.Get("Authorization"))), "bearer ")
  391. }
  392. func newGrokOfficialAPIFallbackRequest(req *http.Request) (*http.Request, error) {
  393. body, err := req.GetBody()
  394. if err != nil {
  395. return nil, err
  396. }
  397. fallbackReq := req.Clone(req.Context())
  398. fallbackReq.Body = body
  399. fallbackReq.URL = cloneURL(req.URL)
  400. fallbackReq.URL.Scheme = "https"
  401. fallbackReq.URL.Host = grokOfficialAPIHost
  402. fallbackReq.Host = ""
  403. fallbackReq.RequestURI = ""
  404. fallbackReq.Header = req.Header.Clone()
  405. for _, header := range []string{
  406. "X-XAI-Token-Auth",
  407. "X-Grok-Client-Version",
  408. "X-Grok-Client-Surface",
  409. "X-UserID",
  410. "X-Email",
  411. "User-Agent",
  412. } {
  413. fallbackReq.Header.Del(header)
  414. }
  415. return fallbackReq, nil
  416. }
  417. func cloneURL(value *url.URL) *url.URL {
  418. if value == nil {
  419. return nil
  420. }
  421. clone := *value
  422. return &clone
  423. }
  424. func bufferSmallResponseBody(resp *http.Response, limit int64) ([]byte, bool) {
  425. if resp == nil || resp.Body == nil || limit <= 0 {
  426. return nil, false
  427. }
  428. original := resp.Body
  429. body, err := io.ReadAll(io.LimitReader(original, limit+1))
  430. if err != nil || int64(len(body)) > limit {
  431. resp.Body = &prefixedReadCloser{
  432. Reader: io.MultiReader(bytes.NewReader(body), original),
  433. Closer: original,
  434. }
  435. return nil, false
  436. }
  437. _ = original.Close()
  438. resp.Body = io.NopCloser(bytes.NewReader(body))
  439. resp.ContentLength = int64(len(body))
  440. return body, true
  441. }
  442. type prefixedReadCloser struct {
  443. io.Reader
  444. io.Closer
  445. }
  446. // applyGrokCLIProxyHeaders applies the official Grok Build client identity at
  447. // the final shared transport boundary. Keying this behavior to the exact CLI
  448. // proxy host keeps direct api.x.ai traffic unchanged and automatically covers
  449. // Responses, Chat Completions, media, quota probes, and account tests.
  450. //
  451. // Operator overrides must be >= CLIClientVersion (the preferred pin). Package
  452. // xai.IsSupportedCLIVersion uses a lower floor (CLIStableVersion) for general
  453. // validation; transport is stricter so we never silently advertise an older pin
  454. // than the binary default.
  455. func applyGrokCLIProxyHeaders(req *http.Request) {
  456. if req == nil || req.URL == nil || !strings.EqualFold(strings.TrimSpace(req.URL.Hostname()), grokCLIProxyHost) {
  457. return
  458. }
  459. if req.Header == nil {
  460. req.Header = make(http.Header)
  461. }
  462. version := strings.TrimSpace(os.Getenv(grokCLIVersionOverride))
  463. if !isSupportedGrokCLIVersion(version) {
  464. version = grokCLIStableVersion
  465. }
  466. req.Header.Set("X-XAI-Token-Auth", xai.CLITokenAuth)
  467. req.Header.Set("x-grok-client-version", version)
  468. req.Header.Set("x-grok-client-identifier", xai.CLIClientIdentifier)
  469. req.Header.Set("User-Agent", xai.CLIUserAgent(version))
  470. }
  471. func isSupportedGrokCLIVersion(version string) bool {
  472. canonical := "v" + version
  473. minimum := "v" + xai.CLIClientVersion
  474. return semver.IsValid(canonical) &&
  475. semver.Canonical(canonical) == canonical &&
  476. semver.Compare(canonical, minimum) >= 0
  477. }
  478. // acquireClientWithTLS 获取或创建带 TLS 指纹的客户端
  479. func (s *httpUpstreamService) acquireClientWithTLS(proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile, upstreamProfile service.HTTPUpstreamProfile) (*upstreamClientEntry, error) {
  480. return s.getClientEntryWithTLS(proxyURL, accountID, accountConcurrency, profile, upstreamProfile, true, true)
  481. }
  482. // getClientEntryWithTLS 获取或创建带 TLS 指纹的客户端条目
  483. // TLS 指纹客户端使用独立的缓存键,与普通客户端隔离
  484. func (s *httpUpstreamService) getClientEntryWithTLS(proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile, upstreamProfile service.HTTPUpstreamProfile, markInFlight bool, enforceLimit bool) (*upstreamClientEntry, error) {
  485. if profile == nil {
  486. return nil, fmt.Errorf("TLS fingerprint client requires a profile")
  487. }
  488. profile = profile.Clone()
  489. // A fingerprint-enabled account always owns its pool, even when the global
  490. // default allows clients to share a pool by proxy.
  491. isolation := config.ConnectionPoolIsolationAccountProxy
  492. proxyKey, parsedProxy, err := normalizeProxyURL(proxyURL)
  493. if err != nil {
  494. return nil, err
  495. }
  496. settings := s.resolvePoolSettings(isolation, accountConcurrency)
  497. settings = s.applyProfilePoolSettings(settings, upstreamProfile)
  498. proxyDigest := sha256.Sum256([]byte(proxyKey))
  499. profileKey := profile.CacheKey()
  500. cacheKey := fmt.Sprintf("tls:account:%d|proxy:%x|profile:%s", accountID, proxyDigest, profileKey)
  501. poolKey := buildPoolKey(settings, upstreamProtocolModeDefault) + ":tls:" + profileKey
  502. now := time.Now()
  503. nowUnix := now.UnixNano()
  504. // 读锁快速路径
  505. s.mu.RLock()
  506. if entry, ok := s.clients[cacheKey]; ok && s.shouldReuseEntry(entry, isolation, proxyKey, poolKey) {
  507. atomic.StoreInt64(&entry.lastUsed, nowUnix)
  508. if markInFlight {
  509. atomic.AddInt64(&entry.inFlight, 1)
  510. }
  511. s.mu.RUnlock()
  512. slog.Debug("tls_fingerprint_reusing_client", "account_id", accountID, "cache_key", cacheKey)
  513. return entry, nil
  514. }
  515. s.mu.RUnlock()
  516. // 写锁慢路径
  517. s.mu.Lock()
  518. if entry, ok := s.clients[cacheKey]; ok {
  519. if s.shouldReuseEntry(entry, isolation, proxyKey, poolKey) {
  520. atomic.StoreInt64(&entry.lastUsed, nowUnix)
  521. if markInFlight {
  522. atomic.AddInt64(&entry.inFlight, 1)
  523. }
  524. s.mu.Unlock()
  525. slog.Debug("tls_fingerprint_reusing_client", "account_id", accountID, "cache_key", cacheKey)
  526. return entry, nil
  527. }
  528. slog.Debug("tls_fingerprint_evicting_stale_client",
  529. "account_id", accountID,
  530. "cache_key", cacheKey,
  531. "proxy_changed", entry.proxyKey != proxyKey,
  532. "pool_changed", entry.poolKey != poolKey)
  533. s.removeClientLocked(cacheKey, entry)
  534. }
  535. // 超出缓存上限时尝试淘汰
  536. if enforceLimit && s.maxUpstreamClients() > 0 {
  537. s.evictIdleLocked(now)
  538. if len(s.clients) >= s.maxUpstreamClients() {
  539. if !s.evictOldestIdleLocked() {
  540. s.mu.Unlock()
  541. return nil, errUpstreamClientLimitReached
  542. }
  543. }
  544. }
  545. // 创建带 TLS 指纹的 Transport
  546. slog.Debug("tls_fingerprint_creating_new_client", "account_id", accountID, "cache_key", cacheKey)
  547. transport, err := buildUpstreamTransportWithTLSFingerprint(settings, parsedProxy, profile)
  548. if err != nil {
  549. s.mu.Unlock()
  550. return nil, fmt.Errorf("build TLS fingerprint transport: %w", err)
  551. }
  552. client := &http.Client{Transport: transport}
  553. if s.shouldValidateResolvedIP() {
  554. client.CheckRedirect = s.redirectChecker
  555. }
  556. entry := &upstreamClientEntry{
  557. client: client,
  558. proxyKey: proxyKey,
  559. poolKey: poolKey,
  560. }
  561. atomic.StoreInt64(&entry.lastUsed, nowUnix)
  562. if markInFlight {
  563. atomic.StoreInt64(&entry.inFlight, 1)
  564. }
  565. s.clients[cacheKey] = entry
  566. s.evictIdleLocked(now)
  567. s.evictOverLimitLocked()
  568. s.mu.Unlock()
  569. return entry, nil
  570. }
  571. func (s *httpUpstreamService) shouldValidateResolvedIP() bool {
  572. if s.cfg == nil {
  573. return false
  574. }
  575. if !s.cfg.Security.URLAllowlist.Enabled {
  576. return false
  577. }
  578. return !s.cfg.Security.URLAllowlist.AllowPrivateHosts
  579. }
  580. // validateRequestHost 校验请求主机的解析结果不落在回环、私网、链路本地或未指定地址。
  581. // 是否全局启用由 security.url_allowlist 决定;带 WithHTTPUpstreamPublicHostsOnly 标记的请求无论配置如何都校验。
  582. func (s *httpUpstreamService) validateRequestHost(req *http.Request) error {
  583. publicHostsOnly := req != nil && service.HTTPUpstreamPublicHostsOnly(req.Context())
  584. if !s.shouldValidateResolvedIP() && !publicHostsOnly {
  585. return nil
  586. }
  587. if req == nil || req.URL == nil {
  588. return errors.New("request url is nil")
  589. }
  590. host := strings.TrimSpace(req.URL.Hostname())
  591. if host == "" {
  592. return errors.New("request host is empty")
  593. }
  594. if err := urlvalidator.ValidateResolvedIP(host); err != nil {
  595. return err
  596. }
  597. return nil
  598. }
  599. func (s *httpUpstreamService) redirectChecker(req *http.Request, via []*http.Request) error {
  600. if len(via) >= 10 {
  601. return errors.New("stopped after 10 redirects")
  602. }
  603. return s.validateRequestHost(req)
  604. }
  605. // acquireClient 获取或创建客户端,并标记为进行中请求
  606. // 用于请求路径,避免在获取后被淘汰
  607. func (s *httpUpstreamService) acquireClient(proxyURL string, accountID int64, accountConcurrency int) (*upstreamClientEntry, error) {
  608. return s.acquireClientWithProfile(proxyURL, accountID, accountConcurrency, service.HTTPUpstreamProfileDefault)
  609. }
  610. // acquireClientWithProfile 获取或创建客户端,并按请求 profile 选择协议策略。
  611. func (s *httpUpstreamService) acquireClientWithProfile(proxyURL string, accountID int64, accountConcurrency int, profile service.HTTPUpstreamProfile) (*upstreamClientEntry, error) {
  612. return s.getClientEntry(proxyURL, accountID, accountConcurrency, profile, true, true)
  613. }
  614. // getOrCreateClient 获取或创建客户端
  615. // 根据隔离策略和参数决定缓存键,处理代理变更和配置变更
  616. //
  617. // 参数:
  618. // - proxyURL: 代理地址
  619. // - accountID: 账户 ID
  620. // - accountConcurrency: 账户并发限制
  621. //
  622. // 返回:
  623. // - *upstreamClientEntry: 客户端缓存条目
  624. //
  625. // 隔离策略说明:
  626. // - proxy: 按代理地址隔离,同一代理共享客户端
  627. // - account: 按账户隔离,同一账户共享客户端(代理变更时重建)
  628. // - account_proxy: 按账户+代理组合隔离,最细粒度
  629. func (s *httpUpstreamService) getOrCreateClient(proxyURL string, accountID int64, accountConcurrency int) (*upstreamClientEntry, error) {
  630. return s.getClientEntry(proxyURL, accountID, accountConcurrency, service.HTTPUpstreamProfileDefault, false, false)
  631. }
  632. // getClientEntry 获取或创建客户端条目
  633. // markInFlight=true 时会标记进行中请求,用于请求路径防止被淘汰
  634. // enforceLimit=true 时会限制客户端数量,超限且无法淘汰时返回错误
  635. func (s *httpUpstreamService) getClientEntry(proxyURL string, accountID int64, accountConcurrency int, profile service.HTTPUpstreamProfile, markInFlight bool, enforceLimit bool) (*upstreamClientEntry, error) {
  636. // 获取隔离模式
  637. isolation := s.getIsolationMode()
  638. // 标准化代理 URL 并解析
  639. proxyKey, parsedProxy, err := normalizeProxyURL(proxyURL)
  640. if err != nil {
  641. return nil, err
  642. }
  643. // 根据请求 profile(例如 OpenAI)选择协议模式
  644. protocolMode := s.resolveProtocolMode(profile, proxyKey, parsedProxy)
  645. settings := s.resolvePoolSettings(isolation, accountConcurrency)
  646. settings = s.applyProfilePoolSettings(settings, profile)
  647. // 构建缓存键(根据隔离策略不同)
  648. cacheKey := buildCacheKey(isolation, proxyKey, accountID, protocolMode)
  649. // 构建连接池配置键(用于检测配置变更)
  650. poolKey := buildPoolKey(settings, protocolMode)
  651. now := time.Now()
  652. nowUnix := now.UnixNano()
  653. // 读锁快速路径:命中缓存直接返回,减少锁竞争
  654. s.mu.RLock()
  655. if entry, ok := s.clients[cacheKey]; ok && s.shouldReuseEntry(entry, isolation, proxyKey, poolKey) {
  656. atomic.StoreInt64(&entry.lastUsed, nowUnix)
  657. if markInFlight {
  658. atomic.AddInt64(&entry.inFlight, 1)
  659. }
  660. s.mu.RUnlock()
  661. return entry, nil
  662. }
  663. s.mu.RUnlock()
  664. // 写锁慢路径:创建或重建客户端
  665. s.mu.Lock()
  666. if entry, ok := s.clients[cacheKey]; ok {
  667. if s.shouldReuseEntry(entry, isolation, proxyKey, poolKey) {
  668. atomic.StoreInt64(&entry.lastUsed, nowUnix)
  669. if markInFlight {
  670. atomic.AddInt64(&entry.inFlight, 1)
  671. }
  672. s.mu.Unlock()
  673. return entry, nil
  674. }
  675. s.removeClientLocked(cacheKey, entry)
  676. }
  677. // 超出缓存上限时尝试淘汰,无法淘汰则拒绝新建
  678. if enforceLimit && s.maxUpstreamClients() > 0 {
  679. s.evictIdleLocked(now)
  680. if len(s.clients) >= s.maxUpstreamClients() {
  681. if !s.evictOldestIdleLocked() {
  682. s.mu.Unlock()
  683. return nil, errUpstreamClientLimitReached
  684. }
  685. }
  686. }
  687. // 缓存未命中或需要重建,创建新客户端
  688. transport, err := buildUpstreamTransport(settings, parsedProxy, protocolMode)
  689. if err != nil {
  690. s.mu.Unlock()
  691. return nil, fmt.Errorf("build transport: %w", err)
  692. }
  693. client := &http.Client{Transport: transport}
  694. if s.shouldValidateResolvedIP() {
  695. client.CheckRedirect = s.redirectChecker
  696. }
  697. entry := &upstreamClientEntry{
  698. client: client,
  699. proxyKey: proxyKey,
  700. poolKey: poolKey,
  701. protocolMode: protocolMode,
  702. }
  703. atomic.StoreInt64(&entry.lastUsed, nowUnix)
  704. if markInFlight {
  705. atomic.StoreInt64(&entry.inFlight, 1)
  706. }
  707. s.clients[cacheKey] = entry
  708. // 执行淘汰策略:先淘汰空闲超时的,再淘汰超出数量限制的
  709. s.evictIdleLocked(now)
  710. s.evictOverLimitLocked()
  711. s.mu.Unlock()
  712. return entry, nil
  713. }
  714. // shouldReuseEntry 判断缓存条目是否可复用
  715. // 若代理或连接池配置发生变化,则需要重建客户端
  716. func (s *httpUpstreamService) shouldReuseEntry(entry *upstreamClientEntry, isolation, proxyKey, poolKey string) bool {
  717. if entry == nil {
  718. return false
  719. }
  720. if isolation == config.ConnectionPoolIsolationAccount && entry.proxyKey != proxyKey {
  721. return false
  722. }
  723. if entry.poolKey != poolKey {
  724. return false
  725. }
  726. return true
  727. }
  728. // removeClientLocked 移除客户端(需持有锁)
  729. // 从缓存中删除并关闭空闲连接
  730. //
  731. // 参数:
  732. // - key: 缓存键
  733. // - entry: 客户端条目
  734. func (s *httpUpstreamService) removeClientLocked(key string, entry *upstreamClientEntry) {
  735. delete(s.clients, key)
  736. if entry != nil && entry.client != nil {
  737. // 关闭空闲连接,释放系统资源
  738. // 注意:这不会中断活跃连接
  739. entry.client.CloseIdleConnections()
  740. }
  741. }
  742. // evictIdleLocked 淘汰空闲超时的客户端(需持有锁)
  743. // 遍历所有客户端,移除超过 TTL 且无活跃请求的条目
  744. //
  745. // 参数:
  746. // - now: 当前时间
  747. func (s *httpUpstreamService) evictIdleLocked(now time.Time) {
  748. ttl := s.clientIdleTTL()
  749. if ttl <= 0 {
  750. return
  751. }
  752. // 计算淘汰截止时间
  753. cutoff := now.Add(-ttl).UnixNano()
  754. for key, entry := range s.clients {
  755. // 跳过有活跃请求的客户端
  756. if atomic.LoadInt64(&entry.inFlight) != 0 {
  757. continue
  758. }
  759. // 淘汰超时的空闲客户端
  760. if atomic.LoadInt64(&entry.lastUsed) <= cutoff {
  761. s.removeClientLocked(key, entry)
  762. }
  763. }
  764. }
  765. // evictOldestIdleLocked 淘汰最久未使用且无活跃请求的客户端(需持有锁)
  766. func (s *httpUpstreamService) evictOldestIdleLocked() bool {
  767. var (
  768. oldestKey string
  769. oldestEntry *upstreamClientEntry
  770. oldestTime int64
  771. )
  772. // 查找最久未使用且无活跃请求的客户端
  773. for key, entry := range s.clients {
  774. // 跳过有活跃请求的客户端
  775. if atomic.LoadInt64(&entry.inFlight) != 0 {
  776. continue
  777. }
  778. lastUsed := atomic.LoadInt64(&entry.lastUsed)
  779. if oldestEntry == nil || lastUsed < oldestTime {
  780. oldestKey = key
  781. oldestEntry = entry
  782. oldestTime = lastUsed
  783. }
  784. }
  785. // 所有客户端都有活跃请求,无法淘汰
  786. if oldestEntry == nil {
  787. return false
  788. }
  789. s.removeClientLocked(oldestKey, oldestEntry)
  790. return true
  791. }
  792. // evictOverLimitLocked 淘汰超出数量限制的客户端(需持有锁)
  793. // 使用 LRU 策略,优先淘汰最久未使用且无活跃请求的客户端
  794. func (s *httpUpstreamService) evictOverLimitLocked() bool {
  795. maxClients := s.maxUpstreamClients()
  796. if maxClients <= 0 {
  797. return false
  798. }
  799. evicted := false
  800. // 循环淘汰直到满足数量限制
  801. for len(s.clients) > maxClients {
  802. if !s.evictOldestIdleLocked() {
  803. return evicted
  804. }
  805. evicted = true
  806. }
  807. return evicted
  808. }
  809. // getIsolationMode 获取连接池隔离模式
  810. // 从配置中读取,无效值回退到 account_proxy 模式
  811. //
  812. // 返回:
  813. // - string: 隔离模式(proxy/account/account_proxy)
  814. func (s *httpUpstreamService) getIsolationMode() string {
  815. if s.cfg == nil {
  816. return config.ConnectionPoolIsolationAccountProxy
  817. }
  818. mode := strings.ToLower(strings.TrimSpace(s.cfg.Gateway.ConnectionPoolIsolation))
  819. if mode == "" {
  820. return config.ConnectionPoolIsolationAccountProxy
  821. }
  822. switch mode {
  823. case config.ConnectionPoolIsolationProxy, config.ConnectionPoolIsolationAccount, config.ConnectionPoolIsolationAccountProxy:
  824. return mode
  825. default:
  826. return config.ConnectionPoolIsolationAccountProxy
  827. }
  828. }
  829. // maxUpstreamClients 获取最大客户端缓存数量
  830. // 从配置中读取,无效值使用默认值
  831. func (s *httpUpstreamService) maxUpstreamClients() int {
  832. if s.cfg == nil {
  833. return defaultMaxUpstreamClients
  834. }
  835. if s.cfg.Gateway.MaxUpstreamClients > 0 {
  836. return s.cfg.Gateway.MaxUpstreamClients
  837. }
  838. return defaultMaxUpstreamClients
  839. }
  840. // clientIdleTTL 获取客户端空闲回收阈值
  841. // 从配置中读取,无效值使用默认值
  842. func (s *httpUpstreamService) clientIdleTTL() time.Duration {
  843. if s.cfg == nil {
  844. return time.Duration(defaultClientIdleTTLSeconds) * time.Second
  845. }
  846. if s.cfg.Gateway.ClientIdleTTLSeconds > 0 {
  847. return time.Duration(s.cfg.Gateway.ClientIdleTTLSeconds) * time.Second
  848. }
  849. return time.Duration(defaultClientIdleTTLSeconds) * time.Second
  850. }
  851. // resolvePoolSettings 解析连接池配置
  852. // 根据隔离策略和账户并发数动态调整连接池参数
  853. //
  854. // 参数:
  855. // - isolation: 隔离模式
  856. // - accountConcurrency: 账户并发限制
  857. //
  858. // 返回:
  859. // - poolSettings: 连接池配置
  860. //
  861. // 说明:
  862. // - 账户隔离模式下,连接池大小与账户并发数对应
  863. // - 这确保了单账户不会占用过多连接资源
  864. func (s *httpUpstreamService) resolvePoolSettings(isolation string, accountConcurrency int) poolSettings {
  865. settings := defaultPoolSettings(s.cfg)
  866. // 账户隔离模式下,根据账户并发数调整连接池大小
  867. if (isolation == config.ConnectionPoolIsolationAccount || isolation == config.ConnectionPoolIsolationAccountProxy) && accountConcurrency > 0 {
  868. settings.maxIdleConns = accountConcurrency
  869. settings.maxIdleConnsPerHost = accountConcurrency
  870. settings.maxConnsPerHost = accountConcurrency
  871. }
  872. return settings
  873. }
  874. func (s *httpUpstreamService) applyProfilePoolSettings(settings poolSettings, profile service.HTTPUpstreamProfile) poolSettings {
  875. switch profile {
  876. case service.HTTPUpstreamProfileOpenAI:
  877. settings.responseHeaderTimeout = 0
  878. if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIResponseHeaderTimeout > 0 {
  879. settings.responseHeaderTimeout = time.Duration(s.cfg.Gateway.OpenAIResponseHeaderTimeout) * time.Second
  880. }
  881. case service.HTTPUpstreamProfileGrok:
  882. // Grok can stall before its first byte under capacity pressure. Keep the
  883. // generic 600s gateway timeout from turning one request into a 10-minute
  884. // resource hold; streaming after headers is unaffected.
  885. settings.responseHeaderTimeout = 120 * time.Second
  886. if s != nil && s.cfg != nil {
  887. settings.responseHeaderTimeout = time.Duration(s.cfg.Gateway.GrokResponseHeaderTimeout) * time.Second
  888. }
  889. }
  890. return settings
  891. }
  892. // buildPoolKey 构建连接池配置键,用于检测连接池配置变更。
  893. func buildPoolKey(settings poolSettings, protocolMode string) string {
  894. base := fmt.Sprintf(
  895. "idle:%d|idle_host:%d|max:%d|idle_timeout:%s|header_timeout:%s",
  896. settings.maxIdleConns,
  897. settings.maxIdleConnsPerHost,
  898. settings.maxConnsPerHost,
  899. settings.idleConnTimeout,
  900. settings.responseHeaderTimeout,
  901. )
  902. if protocolMode == "" || protocolMode == upstreamProtocolModeDefault {
  903. return base
  904. }
  905. return base + "|proto:" + protocolMode
  906. }
  907. // buildCacheKey 构建客户端缓存键
  908. // 根据隔离策略决定缓存键的组成
  909. //
  910. // 参数:
  911. // - isolation: 隔离模式
  912. // - proxyKey: 代理标识
  913. // - accountID: 账户 ID
  914. //
  915. // 返回:
  916. // - string: 缓存键
  917. //
  918. // 缓存键格式:
  919. // - proxy 模式: "proxy:{proxyKey}"
  920. // - account 模式: "account:{accountID}"
  921. // - account_proxy 模式: "account:{accountID}|proxy:{proxyKey}"
  922. func buildCacheKey(isolation, proxyKey string, accountID int64, protocolMode string) string {
  923. var base string
  924. switch isolation {
  925. case config.ConnectionPoolIsolationAccount:
  926. base = fmt.Sprintf("account:%d", accountID)
  927. case config.ConnectionPoolIsolationAccountProxy:
  928. base = fmt.Sprintf("account:%d|proxy:%s", accountID, proxyKey)
  929. default:
  930. base = fmt.Sprintf("proxy:%s", proxyKey)
  931. }
  932. if protocolMode != "" && protocolMode != upstreamProtocolModeDefault {
  933. base += "|proto:" + protocolMode
  934. }
  935. return base
  936. }
  937. func (s *httpUpstreamService) resolveOpenAIHTTP2Settings() openAIHTTP2Settings {
  938. settings := openAIHTTP2Settings{
  939. enabled: false,
  940. allowProxyFallbackToHTTP1: true,
  941. fallbackErrorThreshold: defaultOpenAIHTTP2FallbackErrorThreshold,
  942. fallbackWindow: defaultOpenAIHTTP2FallbackWindow,
  943. fallbackTTL: defaultOpenAIHTTP2FallbackTTL,
  944. }
  945. if s == nil || s.cfg == nil {
  946. return settings
  947. }
  948. cfg := s.cfg.Gateway.OpenAIHTTP2
  949. settings.enabled = cfg.Enabled
  950. settings.allowProxyFallbackToHTTP1 = cfg.AllowProxyFallbackToHTTP1
  951. if cfg.FallbackErrorThreshold > 0 {
  952. settings.fallbackErrorThreshold = cfg.FallbackErrorThreshold
  953. }
  954. if cfg.FallbackWindowSeconds > 0 {
  955. settings.fallbackWindow = time.Duration(cfg.FallbackWindowSeconds) * time.Second
  956. }
  957. if cfg.FallbackTTLSeconds > 0 {
  958. settings.fallbackTTL = time.Duration(cfg.FallbackTTLSeconds) * time.Second
  959. }
  960. return settings
  961. }
  962. func (s *httpUpstreamService) resolveProtocolMode(profile service.HTTPUpstreamProfile, proxyKey string, parsedProxy *url.URL) string {
  963. if profile == service.HTTPUpstreamProfileLongStream {
  964. return upstreamProtocolModeLongStreamH2
  965. }
  966. if profile == service.HTTPUpstreamProfileGrok {
  967. return upstreamProtocolModeGrok
  968. }
  969. if profile != service.HTTPUpstreamProfileOpenAI {
  970. return upstreamProtocolModeDefault
  971. }
  972. settings := s.resolveOpenAIHTTP2Settings()
  973. if !settings.enabled {
  974. return upstreamProtocolModeOpenAIH1
  975. }
  976. if parsedProxy == nil {
  977. return upstreamProtocolModeOpenAIH2
  978. }
  979. scheme := strings.ToLower(parsedProxy.Scheme)
  980. if scheme != "http" && scheme != "https" {
  981. return upstreamProtocolModeOpenAIH2
  982. }
  983. if settings.allowProxyFallbackToHTTP1 && s.isOpenAIHTTP2FallbackActive(proxyKey) {
  984. return upstreamProtocolModeOpenAIH1Fallback
  985. }
  986. return upstreamProtocolModeOpenAIH2
  987. }
  988. func (s *httpUpstreamService) isOpenAIHTTP2FallbackActive(proxyKey string) bool {
  989. raw, ok := s.openAIHTTP2Fallbacks.Load(proxyKey)
  990. if !ok {
  991. return false
  992. }
  993. state, ok := raw.(*openAIHTTP2FallbackState)
  994. if !ok || state == nil {
  995. return false
  996. }
  997. return state.isFallbackActive(time.Now())
  998. }
  999. func (s *httpUpstreamService) getOrCreateOpenAIHTTP2FallbackState(proxyKey string) *openAIHTTP2FallbackState {
  1000. state := &openAIHTTP2FallbackState{}
  1001. actual, _ := s.openAIHTTP2Fallbacks.LoadOrStore(proxyKey, state)
  1002. cached, ok := actual.(*openAIHTTP2FallbackState)
  1003. if !ok || cached == nil {
  1004. return state
  1005. }
  1006. return cached
  1007. }
  1008. func isHTTPProxyKey(proxyKey string) bool {
  1009. return strings.HasPrefix(proxyKey, "http://") || strings.HasPrefix(proxyKey, "https://")
  1010. }
  1011. func isOpenAIHTTP2CompatibilityError(err error) bool {
  1012. if err == nil {
  1013. return false
  1014. }
  1015. if isUpstreamTimeoutError(err) {
  1016. return false
  1017. }
  1018. msg := strings.ToLower(err.Error())
  1019. if msg == "" {
  1020. return false
  1021. }
  1022. markers := []string{
  1023. "alpn",
  1024. "no application protocol",
  1025. "protocol error",
  1026. "stream error",
  1027. "goaway",
  1028. "refused_stream",
  1029. "frame too large",
  1030. }
  1031. for _, marker := range markers {
  1032. if strings.Contains(msg, marker) {
  1033. return true
  1034. }
  1035. }
  1036. return false
  1037. }
  1038. func isUpstreamTimeoutError(err error) bool {
  1039. if err == nil {
  1040. return false
  1041. }
  1042. if errors.Is(err, context.DeadlineExceeded) {
  1043. return true
  1044. }
  1045. var netErr net.Error
  1046. if errors.As(err, &netErr) && netErr.Timeout() {
  1047. return true
  1048. }
  1049. msg := strings.ToLower(err.Error())
  1050. if msg == "" {
  1051. return false
  1052. }
  1053. timeoutMarkers := []string{
  1054. "timeout awaiting response headers",
  1055. "i/o timeout",
  1056. "context deadline exceeded",
  1057. "client.timeout exceeded while awaiting headers",
  1058. "tls handshake timeout",
  1059. }
  1060. for _, marker := range timeoutMarkers {
  1061. if strings.Contains(msg, marker) {
  1062. return true
  1063. }
  1064. }
  1065. return false
  1066. }
  1067. func (s *httpUpstreamService) recordOpenAIHTTP2Failure(profile service.HTTPUpstreamProfile, protocolMode, proxyKey string, err error) {
  1068. if profile != service.HTTPUpstreamProfileOpenAI || protocolMode != upstreamProtocolModeOpenAIH2 {
  1069. return
  1070. }
  1071. settings := s.resolveOpenAIHTTP2Settings()
  1072. if !settings.enabled || !settings.allowProxyFallbackToHTTP1 {
  1073. return
  1074. }
  1075. if !isHTTPProxyKey(proxyKey) || !isOpenAIHTTP2CompatibilityError(err) {
  1076. return
  1077. }
  1078. state := s.getOrCreateOpenAIHTTP2FallbackState(proxyKey)
  1079. activated, until := state.recordFailure(time.Now(), settings.fallbackErrorThreshold, settings.fallbackWindow, settings.fallbackTTL)
  1080. if activated {
  1081. slog.Warn("openai_http2_proxy_fallback_activated",
  1082. "proxy", proxyKey,
  1083. "fallback_until", until.Format(time.RFC3339))
  1084. }
  1085. }
  1086. func (s *httpUpstreamService) recordOpenAIHTTP2Success(profile service.HTTPUpstreamProfile, protocolMode, proxyKey string) {
  1087. if profile != service.HTTPUpstreamProfileOpenAI || protocolMode != upstreamProtocolModeOpenAIH2 {
  1088. return
  1089. }
  1090. if !isHTTPProxyKey(proxyKey) {
  1091. return
  1092. }
  1093. raw, ok := s.openAIHTTP2Fallbacks.Load(proxyKey)
  1094. if !ok {
  1095. return
  1096. }
  1097. state, ok := raw.(*openAIHTTP2FallbackState)
  1098. if !ok || state == nil {
  1099. return
  1100. }
  1101. state.resetErrorWindow()
  1102. }
  1103. func (s *openAIHTTP2FallbackState) isFallbackActive(now time.Time) bool {
  1104. s.mu.Lock()
  1105. defer s.mu.Unlock()
  1106. if s.fallbackUntil.IsZero() {
  1107. return false
  1108. }
  1109. if now.Before(s.fallbackUntil) {
  1110. return true
  1111. }
  1112. s.fallbackUntil = time.Time{}
  1113. return false
  1114. }
  1115. func (s *openAIHTTP2FallbackState) resetErrorWindow() {
  1116. s.mu.Lock()
  1117. defer s.mu.Unlock()
  1118. s.windowStart = time.Time{}
  1119. s.errorCount = 0
  1120. }
  1121. func (s *openAIHTTP2FallbackState) recordFailure(now time.Time, threshold int, window, ttl time.Duration) (bool, time.Time) {
  1122. if threshold <= 0 {
  1123. threshold = defaultOpenAIHTTP2FallbackErrorThreshold
  1124. }
  1125. if window <= 0 {
  1126. window = defaultOpenAIHTTP2FallbackWindow
  1127. }
  1128. if ttl <= 0 {
  1129. ttl = defaultOpenAIHTTP2FallbackTTL
  1130. }
  1131. s.mu.Lock()
  1132. defer s.mu.Unlock()
  1133. if !s.fallbackUntil.IsZero() && now.Before(s.fallbackUntil) {
  1134. return false, s.fallbackUntil
  1135. }
  1136. if !s.fallbackUntil.IsZero() && !now.Before(s.fallbackUntil) {
  1137. s.fallbackUntil = time.Time{}
  1138. }
  1139. if s.windowStart.IsZero() || now.Sub(s.windowStart) > window {
  1140. s.windowStart = now
  1141. s.errorCount = 0
  1142. }
  1143. s.errorCount++
  1144. if s.errorCount < threshold {
  1145. return false, time.Time{}
  1146. }
  1147. s.fallbackUntil = now.Add(ttl)
  1148. s.windowStart = time.Time{}
  1149. s.errorCount = 0
  1150. return true, s.fallbackUntil
  1151. }
  1152. // normalizeProxyURL 标准化代理 URL
  1153. // 处理空值和解析错误,返回标准化的键和解析后的 URL
  1154. //
  1155. // 参数:
  1156. // - raw: 原始代理 URL 字符串
  1157. //
  1158. // 返回:
  1159. // - string: 标准化的代理键(空返回 "direct")
  1160. // - *url.URL: 解析后的 URL(空返回 nil)
  1161. // - error: 非空代理 URL 解析失败时返回错误(禁止回退到直连)
  1162. func normalizeProxyURL(raw string) (string, *url.URL, error) {
  1163. _, parsed, err := proxyurl.Parse(raw)
  1164. if err != nil {
  1165. var ferr error
  1166. _, parsed, ferr = proxyurl.ParseFlexible(raw, "http")
  1167. if ferr != nil {
  1168. return "", nil, err
  1169. }
  1170. }
  1171. if parsed == nil {
  1172. return directProxyKey, nil, nil
  1173. }
  1174. // 规范化:小写 scheme/host,去除路径和查询参数
  1175. parsed.Scheme = strings.ToLower(parsed.Scheme)
  1176. parsed.Host = strings.ToLower(parsed.Host)
  1177. parsed.Path = ""
  1178. parsed.RawPath = ""
  1179. parsed.RawQuery = ""
  1180. parsed.Fragment = ""
  1181. parsed.ForceQuery = false
  1182. if hostname := parsed.Hostname(); hostname != "" {
  1183. port := parsed.Port()
  1184. if (parsed.Scheme == "http" && port == "80") || (parsed.Scheme == "https" && port == "443") {
  1185. port = ""
  1186. }
  1187. hostname = strings.ToLower(hostname)
  1188. if port != "" {
  1189. parsed.Host = net.JoinHostPort(hostname, port)
  1190. } else {
  1191. parsed.Host = hostname
  1192. }
  1193. }
  1194. return parsed.String(), parsed, nil
  1195. }
  1196. // defaultPoolSettings 获取默认连接池配置
  1197. // 从全局配置中读取,无效值使用常量默认值
  1198. //
  1199. // 参数:
  1200. // - cfg: 全局配置
  1201. //
  1202. // 返回:
  1203. // - poolSettings: 连接池配置
  1204. func defaultPoolSettings(cfg *config.Config) poolSettings {
  1205. maxIdleConns := defaultMaxIdleConns
  1206. maxIdleConnsPerHost := defaultMaxIdleConnsPerHost
  1207. maxConnsPerHost := defaultMaxConnsPerHost
  1208. idleConnTimeout := defaultIdleConnTimeout
  1209. responseHeaderTimeout := defaultResponseHeaderTimeout
  1210. if cfg != nil {
  1211. if cfg.Gateway.MaxIdleConns > 0 {
  1212. maxIdleConns = cfg.Gateway.MaxIdleConns
  1213. }
  1214. if cfg.Gateway.MaxIdleConnsPerHost > 0 {
  1215. maxIdleConnsPerHost = cfg.Gateway.MaxIdleConnsPerHost
  1216. }
  1217. if cfg.Gateway.MaxConnsPerHost >= 0 {
  1218. maxConnsPerHost = cfg.Gateway.MaxConnsPerHost
  1219. }
  1220. if cfg.Gateway.IdleConnTimeoutSeconds > 0 {
  1221. idleConnTimeout = time.Duration(cfg.Gateway.IdleConnTimeoutSeconds) * time.Second
  1222. }
  1223. if cfg.Gateway.ResponseHeaderTimeout >= 0 {
  1224. responseHeaderTimeout = time.Duration(cfg.Gateway.ResponseHeaderTimeout) * time.Second
  1225. }
  1226. }
  1227. return poolSettings{
  1228. maxIdleConns: maxIdleConns,
  1229. maxIdleConnsPerHost: maxIdleConnsPerHost,
  1230. maxConnsPerHost: maxConnsPerHost,
  1231. idleConnTimeout: idleConnTimeout,
  1232. responseHeaderTimeout: responseHeaderTimeout,
  1233. }
  1234. }
  1235. // newUpstreamDialer 构建上游 Transport 的 TCP dialer。
  1236. //
  1237. // 必须显式提供:http.Transport 的 DialContext 为 nil 时使用零值 net.Dialer,
  1238. // 建连没有任何超时上限,只能等内核 TCP 重传耗尽(Linux 约 130 秒)。
  1239. func newUpstreamDialer() *net.Dialer {
  1240. return &net.Dialer{
  1241. Timeout: defaultUpstreamDialTimeout,
  1242. KeepAlive: defaultUpstreamDialKeepAlive,
  1243. }
  1244. }
  1245. // buildUpstreamTransport 构建上游请求的 Transport
  1246. // 使用配置文件中的连接池参数,支持生产环境调优
  1247. //
  1248. // 参数:
  1249. // - settings: 连接池配置
  1250. // - proxyURL: 代理 URL(nil 表示直连)
  1251. //
  1252. // 返回:
  1253. // - *http.Transport: 配置好的 Transport 实例
  1254. // - error: 代理配置错误
  1255. //
  1256. // Transport 参数说明:
  1257. // - DialContext: DNS 解析 + TCP 建连超时(不设置则无上限,退化为内核默认重传)
  1258. // - TLSHandshakeTimeout: TLS 握手超时
  1259. // - MaxIdleConns: 所有主机的最大空闲连接总数
  1260. // - MaxIdleConnsPerHost: 每主机最大空闲连接数(影响连接复用率)
  1261. // - MaxConnsPerHost: 每主机最大连接数(达到后新请求等待)
  1262. // - IdleConnTimeout: 空闲连接超时(超时后关闭)
  1263. // - ResponseHeaderTimeout: 等待响应头超时(不影响流式传输)
  1264. func buildUpstreamTransport(settings poolSettings, proxyURL *url.URL, protocolMode string) (*http.Transport, error) {
  1265. transport := &http.Transport{
  1266. DialContext: newUpstreamDialer().DialContext,
  1267. TLSHandshakeTimeout: defaultUpstreamTLSHandshakeTimeout,
  1268. MaxIdleConns: settings.maxIdleConns,
  1269. MaxIdleConnsPerHost: settings.maxIdleConnsPerHost,
  1270. MaxConnsPerHost: settings.maxConnsPerHost,
  1271. IdleConnTimeout: settings.idleConnTimeout,
  1272. ResponseHeaderTimeout: settings.responseHeaderTimeout,
  1273. }
  1274. switch protocolMode {
  1275. case upstreamProtocolModeLongStreamH2, upstreamProtocolModeOpenAIH2:
  1276. transport.ForceAttemptHTTP2 = true
  1277. // 显式配置 http2 并启用 PING 健康探测,剔除代理/NAT 静默掐断的死连接,
  1278. // 避免请求挂在死连接上直到 TCP 重传超时(分钟级)。
  1279. if _, err := enableHTTP2KeepAlive(transport); err != nil {
  1280. return nil, err
  1281. }
  1282. case upstreamProtocolModeOpenAIH1:
  1283. transport.ForceAttemptHTTP2 = false
  1284. transport.TLSNextProto = make(map[string]func(string, *tls.Conn) http.RoundTripper)
  1285. case upstreamProtocolModeOpenAIH1Fallback:
  1286. // 显式禁用 HTTP/2,确保代理不兼容场景回退到 HTTP/1.1。
  1287. transport.ForceAttemptHTTP2 = false
  1288. transport.TLSNextProto = make(map[string]func(string, *tls.Conn) http.RoundTripper)
  1289. }
  1290. if err := proxyutil.ConfigureTransportProxy(transport, proxyURL); err != nil {
  1291. return nil, err
  1292. }
  1293. return transport, nil
  1294. }
  1295. // enableHTTP2KeepAlive 在 http.Transport 上显式配置 HTTP/2 并启用连接健康探测。
  1296. // Go 默认惰性配置 http2 且 ReadIdleTimeout=0(不发健康 PING),无法检测被代理/NAT
  1297. // 静默掐断的死连接。此处主动设置 ReadIdleTimeout/PingTimeout,让死连接被提前 PING
  1298. // 出并关闭,请求得以重建连接而非挂到 TCP 重传超时。返回底层 *http2.Transport 便于测试。
  1299. func enableHTTP2KeepAlive(transport *http.Transport) (*http2.Transport, error) {
  1300. h2, err := http2.ConfigureTransports(transport)
  1301. if err != nil {
  1302. return nil, err
  1303. }
  1304. if h2 != nil {
  1305. h2.ReadIdleTimeout = longStreamHTTP2ReadIdleTimeout
  1306. h2.PingTimeout = longStreamHTTP2PingTimeout
  1307. }
  1308. return h2, nil
  1309. }
  1310. // buildUpstreamTransportWithTLSFingerprint 构建带 TLS 指纹伪装的 Transport
  1311. // 使用 utls 库模拟 Claude CLI 的 TLS 指纹
  1312. //
  1313. // 参数:
  1314. // - settings: 连接池配置
  1315. // - proxyURL: 代理 URL(nil 表示直连)
  1316. // - profile: TLS 指纹配置
  1317. //
  1318. // 返回:
  1319. // - *http.Transport: 配置好的 Transport 实例
  1320. // - error: 配置错误
  1321. //
  1322. // 代理类型处理:
  1323. // - nil/空: 直连,使用 TLSFingerprintDialer
  1324. // - http/https: HTTP 代理,使用 HTTPProxyDialer(CONNECT 隧道 + utls 握手)
  1325. // - socks5: SOCKS5 代理,使用 SOCKS5ProxyDialer(SOCKS5 隧道 + utls 握手)
  1326. func buildUpstreamTransportWithTLSFingerprint(settings poolSettings, proxyURL *url.URL, profile *tlsfingerprint.Profile) (*http.Transport, error) {
  1327. transport := &http.Transport{
  1328. MaxIdleConns: settings.maxIdleConns,
  1329. MaxIdleConnsPerHost: settings.maxIdleConnsPerHost,
  1330. MaxConnsPerHost: settings.maxConnsPerHost,
  1331. IdleConnTimeout: settings.idleConnTimeout,
  1332. ResponseHeaderTimeout: settings.responseHeaderTimeout,
  1333. // 禁用默认的 TLS,我们使用自定义的 DialTLSContext
  1334. ForceAttemptHTTP2: false,
  1335. }
  1336. if err := tlsfingerprint.ConfigureTransport(transport, profile, proxyURL, tlsfingerprint.TransportOptions{
  1337. DialContext: newUpstreamDialer().DialContext, HandshakeTimeout: defaultUpstreamTLSHandshakeTimeout,
  1338. }); err != nil {
  1339. return nil, err
  1340. }
  1341. return transport, nil
  1342. }
  1343. // trackedBody 带跟踪功能的响应体包装器
  1344. // 在 Close 时执行回调,用于更新请求计数
  1345. type trackedBody struct {
  1346. io.ReadCloser // 原始响应体
  1347. once sync.Once
  1348. onClose func() // 关闭时的回调函数
  1349. }
  1350. // Close 关闭响应体并执行回调
  1351. // 使用 sync.Once 确保回调只执行一次
  1352. func (b *trackedBody) Close() error {
  1353. err := b.ReadCloser.Close()
  1354. if b.onClose != nil {
  1355. b.once.Do(b.onClose)
  1356. }
  1357. return err
  1358. }
  1359. // wrapTrackedBody 包装响应体以跟踪关闭事件
  1360. // 用于在响应体关闭时更新 inFlight 计数
  1361. //
  1362. // 参数:
  1363. // - body: 原始响应体
  1364. // - onClose: 关闭时的回调函数
  1365. //
  1366. // 返回:
  1367. // - io.ReadCloser: 包装后的响应体
  1368. func wrapTrackedBody(body io.ReadCloser, onClose func()) io.ReadCloser {
  1369. if body == nil {
  1370. return body
  1371. }
  1372. return &trackedBody{ReadCloser: body, onClose: onClose}
  1373. }
  1374. // decompressResponseBody 根据 Content-Encoding 解压响应体。
  1375. // 当请求显式设置了 accept-encoding 时,Go 的 Transport 不会自动解压,需要手动处理。
  1376. // 解压成功后会删除 Content-Encoding 和 Content-Length header(长度已不准确)。
  1377. func decompressResponseBody(resp *http.Response) {
  1378. if resp == nil || resp.Body == nil {
  1379. return
  1380. }
  1381. ce := strings.ToLower(strings.TrimSpace(resp.Header.Get("Content-Encoding")))
  1382. if ce == "" {
  1383. return
  1384. }
  1385. originalBody := resp.Body
  1386. var reader io.Reader
  1387. switch ce {
  1388. case "gzip":
  1389. gr, err := gzip.NewReader(resp.Body)
  1390. if err != nil {
  1391. return // 解压失败,保持原样
  1392. }
  1393. reader = gr
  1394. case "br":
  1395. reader = brotli.NewReader(resp.Body)
  1396. case "deflate":
  1397. reader = flate.NewReader(resp.Body)
  1398. case "zstd":
  1399. bufferedBody := bufio.NewReader(resp.Body)
  1400. resp.Body = &decompressedBody{reader: bufferedBody, closer: originalBody}
  1401. headerBytes, _ := bufferedBody.Peek(zstd.HeaderMaxSize)
  1402. var header zstd.Header
  1403. if err := header.Decode(headerBytes); err != nil {
  1404. slog.Warn("zstd_decompress_failed", "error", err)
  1405. return
  1406. }
  1407. zr, err := zstd.NewReader(bufferedBody)
  1408. if err != nil {
  1409. slog.Warn("zstd_decompress_failed", "error", err)
  1410. return
  1411. }
  1412. reader = &zstdResponseReader{ReadCloser: zr.IOReadCloser()}
  1413. default:
  1414. return
  1415. }
  1416. resp.Body = &decompressedBody{reader: reader, closer: originalBody}
  1417. resp.Header.Del("Content-Encoding")
  1418. resp.Header.Del("Content-Length") // 解压后长度不确定
  1419. resp.ContentLength = -1
  1420. }
  1421. type zstdResponseReader struct {
  1422. io.ReadCloser
  1423. warnOnce sync.Once
  1424. }
  1425. func (r *zstdResponseReader) Read(p []byte) (int, error) {
  1426. n, err := r.ReadCloser.Read(p)
  1427. if err != nil && !errors.Is(err, io.EOF) {
  1428. r.warnOnce.Do(func() {
  1429. slog.Warn("zstd_decompress_failed", "error", err)
  1430. })
  1431. }
  1432. return n, err
  1433. }
  1434. // decompressedBody 组合解压 reader 和原始 body 的 close。
  1435. type decompressedBody struct {
  1436. reader io.Reader
  1437. closer io.Closer
  1438. }
  1439. func (d *decompressedBody) Read(p []byte) (int, error) {
  1440. return d.reader.Read(p)
  1441. }
  1442. func (d *decompressedBody) Close() error {
  1443. // 如果 reader 本身也是 Closer(如 gzip.Reader),先关闭它
  1444. if rc, ok := d.reader.(io.Closer); ok {
  1445. _ = rc.Close()
  1446. }
  1447. return d.closer.Close()
  1448. }