| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565 |
- package repository
- import (
- "bufio"
- "bytes"
- "compress/flate"
- "compress/gzip"
- "context"
- "crypto/sha256"
- "crypto/tls"
- "encoding/json"
- "errors"
- "fmt"
- "io"
- "log/slog"
- "net"
- "net/http"
- "net/url"
- "os"
- "strings"
- "sync"
- "sync/atomic"
- "time"
- "github.com/andybalholm/brotli"
- "github.com/klauspost/compress/zstd"
- "golang.org/x/mod/semver"
- "golang.org/x/net/http2"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
- "github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
- "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
- "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
- "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
- "github.com/Wei-Shaw/sub2api/internal/service"
- "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
- )
- // 默认配置常量
- // 这些值在配置文件未指定时作为回退默认值使用
- const (
- // directProxyKey: 无代理时的缓存键标识
- directProxyKey = "direct"
- // defaultMaxIdleConns: 默认最大空闲连接总数
- // HTTP/2 场景下,单连接可多路复用,240 足以支撑高并发
- defaultMaxIdleConns = 240
- // defaultMaxIdleConnsPerHost: 默认每主机最大空闲连接数
- defaultMaxIdleConnsPerHost = 120
- // defaultMaxConnsPerHost: 默认每主机最大连接数(含活跃连接)
- // 达到上限后新请求会等待,而非无限创建连接
- defaultMaxConnsPerHost = 240
- // defaultIdleConnTimeout: 默认空闲连接超时时间(90秒)
- // 超时后连接会被关闭,释放系统资源(建议小于上游 LB 超时)
- defaultIdleConnTimeout = 90 * time.Second
- // defaultResponseHeaderTimeout: 默认等待响应头超时时间(5分钟)
- // LLM 请求可能排队较久,需要较长超时
- defaultResponseHeaderTimeout = 300 * time.Second
- // defaultUpstreamDialTimeout: 默认 TCP/DNS 建连超时(10秒)
- // Transport 不设置 DialContext 时会退化为零值 net.Dialer(无超时),建连阶段
- // 只能依赖内核默认 TCP 重传(Linux 约 130 秒)。ResponseHeaderTimeout 只约束
- // 连接建立之后等待响应头的阶段,覆盖不到 DNS 解析与 TCP 握手。
- // 上游域名被解析到 443 不可达的 IP 时(DNS 污染/路由异常),单个账号就要卡满
- // 内核超时;而多账号故障转移是串行的,一次请求会阻塞数分钟且不写中间错误。
- defaultUpstreamDialTimeout = 10 * time.Second
- // defaultUpstreamDialKeepAlive: TCP keepalive 探测间隔,与 Go 默认值保持一致
- defaultUpstreamDialKeepAlive = 30 * time.Second
- // defaultUpstreamTLSHandshakeTimeout: TLS 握手超时(10秒)
- // 与建连超时同量级,避免 TCP 已连通但对端不推进握手时无限等待
- defaultUpstreamTLSHandshakeTimeout = 10 * time.Second
- // defaultMaxUpstreamClients: 默认最大客户端缓存数量
- // 超出后会淘汰最久未使用的客户端
- defaultMaxUpstreamClients = 5000
- // defaultClientIdleTTLSeconds: 默认客户端空闲回收阈值(15分钟)
- defaultClientIdleTTLSeconds = 900
- // OpenAI HTTP/2 代理回退策略默认值
- defaultOpenAIHTTP2FallbackErrorThreshold = 2
- defaultOpenAIHTTP2FallbackWindow = 60 * time.Second
- defaultOpenAIHTTP2FallbackTTL = 10 * time.Minute
- // 长流 HTTP/2 连接健康探测:池化连接被代理/NAT
- // 静默掐断会成为“死连接”(两端都以为存活),请求落上去会挂到 TCP 重传超时
- // (分钟级)。Go 的 http2.Transport 默认 ReadIdleTimeout=0(不发健康 PING),
- // 无法检测。启用主动 PING 探测:连接空闲 ReadIdleTimeout 后发 PING,PingTimeout
- // 内无响应即判定死连接并关闭,从源头避免请求挂在死连接上。
- longStreamHTTP2ReadIdleTimeout = 10 * time.Second
- longStreamHTTP2PingTimeout = 5 * time.Second
- // The Grok CLI proxy rejects requests that do not identify a supported
- // client version. Host/env/version pins live in package xai so service,
- // billing, and transport layers advertise the same identity.
- grokCLIProxyHost = xai.CLIProxyHost
- grokOfficialAPIHost = "api.x.ai"
- grokCLIStableVersion = xai.CLIClientVersion // preferred pin (not the minimum floor)
- grokCLIVersionOverride = xai.CLIVersionEnv
- grokFallbackBodyLimit = 64 << 10
- )
- const (
- upstreamProtocolModeDefault = "default"
- upstreamProtocolModeLongStreamH2 = "long_stream_h2"
- upstreamProtocolModeOpenAIH1 = "openai_h1"
- upstreamProtocolModeOpenAIH2 = "openai_h2"
- upstreamProtocolModeOpenAIH1Fallback = "openai_h1_fallback"
- upstreamProtocolModeGrok = "grok"
- )
- var errUpstreamClientLimitReached = errors.New("upstream client cache limit reached")
- // poolSettings 连接池配置参数
- // 封装 Transport 所需的各项连接池参数
- type poolSettings struct {
- maxIdleConns int // 最大空闲连接总数
- maxIdleConnsPerHost int // 每主机最大空闲连接数
- maxConnsPerHost int // 每主机最大连接数(含活跃)
- idleConnTimeout time.Duration // 空闲连接超时时间
- responseHeaderTimeout time.Duration // 等待响应头超时时间
- }
- type openAIHTTP2Settings struct {
- enabled bool
- allowProxyFallbackToHTTP1 bool
- fallbackErrorThreshold int
- fallbackWindow time.Duration
- fallbackTTL time.Duration
- }
- // upstreamClientEntry 上游客户端缓存条目
- // 记录客户端实例及其元数据,用于连接池管理和淘汰策略
- type upstreamClientEntry struct {
- client *http.Client // HTTP 客户端实例
- proxyKey string // 代理标识(用于检测代理变更)
- poolKey string // 连接池配置标识(用于检测配置变更)
- protocolMode string // 协议模式(default/long_stream_h2/openai_h1/openai_h2/openai_h1_fallback)
- lastUsed int64 // 最后使用时间戳(纳秒),用于 LRU 淘汰
- inFlight int64 // 当前进行中的请求数,>0 时不可淘汰
- }
- type openAIHTTP2FallbackState struct {
- mu sync.Mutex
- windowStart time.Time
- errorCount int
- fallbackUntil time.Time
- }
- // httpUpstreamService 通用 HTTP 上游服务
- // 用于向任意 HTTP API(Claude、OpenAI 等)发送请求,支持可选代理
- //
- // 架构设计:
- // - 根据隔离策略(proxy/account/account_proxy)缓存客户端实例
- // - 每个客户端拥有独立的 Transport 连接池
- // - 支持 LRU + 空闲时间双重淘汰策略
- //
- // 性能优化:
- // 1. 根据隔离策略缓存客户端实例,避免频繁创建 http.Client
- // 2. 复用 Transport 连接池,减少 TCP 握手和 TLS 协商开销
- // 3. 支持账号级隔离与空闲回收,降低连接层关联风险
- // 4. 达到最大连接数后等待可用连接,而非无限创建
- // 5. 仅回收空闲客户端,避免中断活跃请求
- // 6. HTTP/2 多路复用,连接上限不等于并发请求上限
- // 7. 代理变更时清空旧连接池,避免复用错误代理
- // 8. 账号并发数与连接池上限对应(账号隔离策略下)
- type httpUpstreamService struct {
- traffic *service.AccountTrafficService
- cfg *config.Config // 全局配置
- mu sync.RWMutex // 保护 clients map 的读写锁
- clients map[string]*upstreamClientEntry // 客户端缓存池,key 由隔离策略决定
- // OpenAI 走 HTTP/HTTPS 代理时的 H2->H1 回退状态(key=标准化 proxyKey)
- openAIHTTP2Fallbacks sync.Map
- }
- // NewHTTPUpstream 创建通用 HTTP 上游服务
- // 使用配置中的连接池参数构建 Transport
- //
- // 参数:
- // - cfg: 全局配置,包含连接池参数和隔离策略
- //
- // 返回:
- // - service.HTTPUpstream 接口实现
- func NewHTTPUpstream(cfg *config.Config) service.HTTPUpstream {
- return &httpUpstreamService{
- cfg: cfg,
- clients: make(map[string]*upstreamClientEntry),
- }
- }
- func NewControlledHTTPUpstream(cfg *config.Config, traffic *service.AccountTrafficService) service.HTTPUpstream {
- s := NewHTTPUpstream(cfg).(*httpUpstreamService)
- s.traffic = traffic
- return s
- }
- func (s *httpUpstreamService) AccountTrafficController() *service.AccountTrafficService {
- return s.traffic
- }
- // Do 执行 HTTP 请求
- // 根据隔离策略获取或创建客户端,并跟踪请求生命周期
- //
- // 参数:
- // - req: HTTP 请求对象
- // - proxyURL: 代理地址,空字符串表示直连
- // - accountID: 账户 ID,用于账户级隔离
- // - accountConcurrency: 账户并发限制,用于动态调整连接池大小
- //
- // 返回:
- // - *http.Response: HTTP 响应(Body 已包装,关闭时自动更新计数)
- // - error: 请求错误
- //
- // 注意:
- // - 调用方必须关闭 resp.Body,否则会导致 inFlight 计数泄漏
- // - inFlight > 0 的客户端不会被淘汰,确保活跃请求不被中断
- func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
- return s.traffic.DoHTTP(req, func(controlled *http.Request) (*http.Response, error) {
- return s.doUncontrolled(controlled, proxyURL, accountID, accountConcurrency)
- })
- }
- func (s *httpUpstreamService) doUncontrolled(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
- applyGrokCLIProxyHeaders(req)
- if err := s.validateRequestHost(req); err != nil {
- return nil, err
- }
- profile := service.HTTPUpstreamProfileDefault
- if req != nil {
- profile = service.HTTPUpstreamProfileFromContext(req.Context())
- }
- // 获取或创建对应的客户端,并标记请求占用
- entry, err := s.acquireClientWithProfile(proxyURL, accountID, accountConcurrency, profile)
- if err != nil {
- return nil, err
- }
- // 执行请求
- client := s.httpClientForUpstreamRequest(entry.client, req)
- client = httpClientWithGrokAccessDeniedFallback(client)
- resp, err := servertiming.Do(client, req)
- if err != nil {
- s.recordOpenAIHTTP2Failure(profile, entry.protocolMode, entry.proxyKey, err)
- // 请求失败,立即减少计数
- atomic.AddInt64(&entry.inFlight, -1)
- atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
- return nil, err
- }
- s.recordOpenAIHTTP2Success(profile, entry.protocolMode, entry.proxyKey)
- // 如果上游返回了压缩内容,解压后再交给业务层
- decompressResponseBody(resp)
- // 包装响应体,在关闭时自动减少计数并更新时间戳
- // 这确保了流式响应(如 SSE)在完全读取前不会被淘汰
- resp.Body = wrapTrackedBody(resp.Body, func() {
- atomic.AddInt64(&entry.inFlight, -1)
- atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
- })
- return resp, nil
- }
- // DoWithTLS 执行带 TLS 指纹伪装的 HTTP 请求
- //
- // profile 为 nil 时不启用 TLS 指纹,行为与 Do 方法相同。
- // profile 非 nil 时使用指定的 Profile 进行 TLS 指纹伪装。
- func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile) (*http.Response, error) {
- return s.traffic.DoHTTP(req, func(controlled *http.Request) (*http.Response, error) {
- return s.doWithTLSUncontrolled(controlled, proxyURL, accountID, accountConcurrency, profile)
- })
- }
- func (s *httpUpstreamService) doWithTLSUncontrolled(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile) (*http.Response, error) {
- if profile == nil {
- return s.Do(req, proxyURL, accountID, accountConcurrency)
- }
- // An explicitly selected TLS profile must not silently take an unprotected
- // route, including redirects to HTTP or a plaintext upstream URL.
- if req == nil || req.URL == nil || !strings.EqualFold(req.URL.Scheme, "https") {
- return nil, fmt.Errorf("TLS fingerprint request requires HTTPS")
- }
- applyGrokCLIProxyHeaders(req)
- upstreamProfile := service.HTTPUpstreamProfileDefault
- if req != nil {
- upstreamProfile = service.HTTPUpstreamProfileFromContext(req.Context())
- }
- targetHost := ""
- if req != nil && req.URL != nil {
- targetHost = req.URL.Host
- }
- proxyInfo := "direct"
- if proxyURL != "" {
- if parsed, parseErr := url.Parse(proxyURL); parseErr == nil {
- proxyInfo = parsed.Redacted()
- } else {
- proxyInfo = "configured"
- }
- }
- slog.Debug("tls_fingerprint_enabled", "account_id", accountID, "target", targetHost, "proxy", proxyInfo, "profile", profile.Name)
- if err := s.validateRequestHost(req); err != nil {
- return nil, err
- }
- entry, err := s.acquireClientWithTLS(proxyURL, accountID, accountConcurrency, profile, upstreamProfile)
- if err != nil {
- slog.Debug("tls_fingerprint_acquire_client_failed", "account_id", accountID, "error", err)
- return nil, err
- }
- client := s.httpClientForUpstreamRequest(entry.client, req)
- client = httpClientWithGrokAccessDeniedFallback(client)
- resp, err := servertiming.Do(client, req)
- if err != nil {
- atomic.AddInt64(&entry.inFlight, -1)
- atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
- slog.Debug("tls_fingerprint_request_failed", "account_id", accountID, "error", err)
- return nil, err
- }
- decompressResponseBody(resp)
- resp.Body = wrapTrackedBody(resp.Body, func() {
- atomic.AddInt64(&entry.inFlight, -1)
- atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
- })
- return resp, nil
- }
- // httpClientForUpstreamRequest 按请求上下文的标记派生客户端:禁用重定向,或对重定向的每一跳做主机校验。
- // 派生的克隆与缓存客户端共享 Transport;未打标记时原样返回。
- func (s *httpUpstreamService) httpClientForUpstreamRequest(client *http.Client, req *http.Request) *http.Client {
- if client == nil || req == nil {
- return client
- }
- ctx := req.Context()
- switch {
- case service.HTTPUpstreamRedirectsDisabled(ctx):
- clone := *client
- clone.CheckRedirect = func(*http.Request, []*http.Request) error {
- return http.ErrUseLastResponse
- }
- return &clone
- case service.HTTPUpstreamPublicHostsOnly(ctx) && client.CheckRedirect == nil:
- clone := *client
- clone.CheckRedirect = s.redirectChecker
- return &clone
- default:
- return client
- }
- }
- // grokAccessDeniedFallbackTransport preserves the subscription CLI proxy as
- // the primary OAuth route, but retries a replayable request against api.x.ai
- // when the proxy returns its compatibility-specific 403 "Access denied".
- // Trial subscriptions can hit this boundary while the same OAuth credential
- // remains valid on the official API. Other entitlement failures stay on the
- // original response so account scheduling semantics do not change.
- type grokAccessDeniedFallbackTransport struct {
- base http.RoundTripper
- }
- func httpClientWithGrokAccessDeniedFallback(client *http.Client) *http.Client {
- if client == nil {
- return nil
- }
- clone := *client
- base := clone.Transport
- if base == nil {
- base = http.DefaultTransport
- }
- clone.Transport = &grokAccessDeniedFallbackTransport{base: base}
- return &clone
- }
- func (t *grokAccessDeniedFallbackTransport) RoundTrip(req *http.Request) (*http.Response, error) {
- resp, err := t.base.RoundTrip(req)
- if err != nil || !isGrokCLIAccessDeniedFallbackCandidate(req, resp) {
- return resp, err
- }
- body, ok := bufferSmallResponseBody(resp, grokFallbackBodyLimit)
- if !ok || !isGrokCLICompatibilityAccessDenied(body) {
- return resp, nil
- }
- fallbackReq, err := newGrokOfficialAPIFallbackRequest(req)
- if err != nil {
- return resp, nil
- }
- fallbackResp, fallbackErr := t.base.RoundTrip(fallbackReq)
- if fallbackErr != nil {
- slog.Debug("grok_cli_access_denied_api_fallback_failed", "path", req.URL.EscapedPath(), "error", fallbackErr)
- return resp, nil
- }
- if fallbackResp.StatusCode < http.StatusOK || fallbackResp.StatusCode >= http.StatusMultipleChoices {
- if fallbackResp.Body != nil {
- _ = fallbackResp.Body.Close()
- }
- return resp, nil
- }
- if resp.Body != nil {
- _ = resp.Body.Close()
- }
- slog.Warn("grok_cli_access_denied_api_fallback_succeeded", "method", req.Method, "path", req.URL.EscapedPath())
- return fallbackResp, nil
- }
- func isGrokCLICompatibilityAccessDenied(body []byte) bool {
- lower := bytes.ToLower(body)
- if bytes.Contains(lower, []byte("access denied")) {
- return true
- }
- var payload struct {
- Code string `json:"code"`
- Error string `json:"error"`
- }
- if err := json.Unmarshal(body, &payload); err != nil || !strings.EqualFold(strings.TrimSpace(payload.Code), "permission_denied") {
- return false
- }
- 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"
- return strings.HasPrefix(strings.ToLower(strings.TrimSpace(payload.Error)), chatEndpointDeniedPrefix)
- }
- func isGrokCLIAccessDeniedFallbackCandidate(req *http.Request, resp *http.Response) bool {
- return req != nil && req.URL != nil && req.GetBody != nil && resp != nil &&
- resp.StatusCode == http.StatusForbidden &&
- strings.EqualFold(strings.TrimSpace(req.URL.Hostname()), grokCLIProxyHost) &&
- strings.EqualFold(strings.TrimSpace(req.Header.Get("X-XAI-Token-Auth")), "xai-grok-cli") &&
- strings.HasPrefix(strings.ToLower(strings.TrimSpace(req.Header.Get("Authorization"))), "bearer ")
- }
- func newGrokOfficialAPIFallbackRequest(req *http.Request) (*http.Request, error) {
- body, err := req.GetBody()
- if err != nil {
- return nil, err
- }
- fallbackReq := req.Clone(req.Context())
- fallbackReq.Body = body
- fallbackReq.URL = cloneURL(req.URL)
- fallbackReq.URL.Scheme = "https"
- fallbackReq.URL.Host = grokOfficialAPIHost
- fallbackReq.Host = ""
- fallbackReq.RequestURI = ""
- fallbackReq.Header = req.Header.Clone()
- for _, header := range []string{
- "X-XAI-Token-Auth",
- "X-Grok-Client-Version",
- "X-Grok-Client-Surface",
- "X-UserID",
- "X-Email",
- "User-Agent",
- } {
- fallbackReq.Header.Del(header)
- }
- return fallbackReq, nil
- }
- func cloneURL(value *url.URL) *url.URL {
- if value == nil {
- return nil
- }
- clone := *value
- return &clone
- }
- func bufferSmallResponseBody(resp *http.Response, limit int64) ([]byte, bool) {
- if resp == nil || resp.Body == nil || limit <= 0 {
- return nil, false
- }
- original := resp.Body
- body, err := io.ReadAll(io.LimitReader(original, limit+1))
- if err != nil || int64(len(body)) > limit {
- resp.Body = &prefixedReadCloser{
- Reader: io.MultiReader(bytes.NewReader(body), original),
- Closer: original,
- }
- return nil, false
- }
- _ = original.Close()
- resp.Body = io.NopCloser(bytes.NewReader(body))
- resp.ContentLength = int64(len(body))
- return body, true
- }
- type prefixedReadCloser struct {
- io.Reader
- io.Closer
- }
- // applyGrokCLIProxyHeaders applies the official Grok Build client identity at
- // the final shared transport boundary. Keying this behavior to the exact CLI
- // proxy host keeps direct api.x.ai traffic unchanged and automatically covers
- // Responses, Chat Completions, media, quota probes, and account tests.
- //
- // Operator overrides must be >= CLIClientVersion (the preferred pin). Package
- // xai.IsSupportedCLIVersion uses a lower floor (CLIStableVersion) for general
- // validation; transport is stricter so we never silently advertise an older pin
- // than the binary default.
- func applyGrokCLIProxyHeaders(req *http.Request) {
- if req == nil || req.URL == nil || !strings.EqualFold(strings.TrimSpace(req.URL.Hostname()), grokCLIProxyHost) {
- return
- }
- if req.Header == nil {
- req.Header = make(http.Header)
- }
- version := strings.TrimSpace(os.Getenv(grokCLIVersionOverride))
- if !isSupportedGrokCLIVersion(version) {
- version = grokCLIStableVersion
- }
- req.Header.Set("X-XAI-Token-Auth", xai.CLITokenAuth)
- req.Header.Set("x-grok-client-version", version)
- req.Header.Set("x-grok-client-identifier", xai.CLIClientIdentifier)
- req.Header.Set("User-Agent", xai.CLIUserAgent(version))
- }
- func isSupportedGrokCLIVersion(version string) bool {
- canonical := "v" + version
- minimum := "v" + xai.CLIClientVersion
- return semver.IsValid(canonical) &&
- semver.Canonical(canonical) == canonical &&
- semver.Compare(canonical, minimum) >= 0
- }
- // acquireClientWithTLS 获取或创建带 TLS 指纹的客户端
- func (s *httpUpstreamService) acquireClientWithTLS(proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile, upstreamProfile service.HTTPUpstreamProfile) (*upstreamClientEntry, error) {
- return s.getClientEntryWithTLS(proxyURL, accountID, accountConcurrency, profile, upstreamProfile, true, true)
- }
- // getClientEntryWithTLS 获取或创建带 TLS 指纹的客户端条目
- // TLS 指纹客户端使用独立的缓存键,与普通客户端隔离
- func (s *httpUpstreamService) getClientEntryWithTLS(proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile, upstreamProfile service.HTTPUpstreamProfile, markInFlight bool, enforceLimit bool) (*upstreamClientEntry, error) {
- if profile == nil {
- return nil, fmt.Errorf("TLS fingerprint client requires a profile")
- }
- profile = profile.Clone()
- // A fingerprint-enabled account always owns its pool, even when the global
- // default allows clients to share a pool by proxy.
- isolation := config.ConnectionPoolIsolationAccountProxy
- proxyKey, parsedProxy, err := normalizeProxyURL(proxyURL)
- if err != nil {
- return nil, err
- }
- settings := s.resolvePoolSettings(isolation, accountConcurrency)
- settings = s.applyProfilePoolSettings(settings, upstreamProfile)
- proxyDigest := sha256.Sum256([]byte(proxyKey))
- profileKey := profile.CacheKey()
- cacheKey := fmt.Sprintf("tls:account:%d|proxy:%x|profile:%s", accountID, proxyDigest, profileKey)
- poolKey := buildPoolKey(settings, upstreamProtocolModeDefault) + ":tls:" + profileKey
- now := time.Now()
- nowUnix := now.UnixNano()
- // 读锁快速路径
- s.mu.RLock()
- if entry, ok := s.clients[cacheKey]; ok && s.shouldReuseEntry(entry, isolation, proxyKey, poolKey) {
- atomic.StoreInt64(&entry.lastUsed, nowUnix)
- if markInFlight {
- atomic.AddInt64(&entry.inFlight, 1)
- }
- s.mu.RUnlock()
- slog.Debug("tls_fingerprint_reusing_client", "account_id", accountID, "cache_key", cacheKey)
- return entry, nil
- }
- s.mu.RUnlock()
- // 写锁慢路径
- s.mu.Lock()
- if entry, ok := s.clients[cacheKey]; ok {
- if s.shouldReuseEntry(entry, isolation, proxyKey, poolKey) {
- atomic.StoreInt64(&entry.lastUsed, nowUnix)
- if markInFlight {
- atomic.AddInt64(&entry.inFlight, 1)
- }
- s.mu.Unlock()
- slog.Debug("tls_fingerprint_reusing_client", "account_id", accountID, "cache_key", cacheKey)
- return entry, nil
- }
- slog.Debug("tls_fingerprint_evicting_stale_client",
- "account_id", accountID,
- "cache_key", cacheKey,
- "proxy_changed", entry.proxyKey != proxyKey,
- "pool_changed", entry.poolKey != poolKey)
- s.removeClientLocked(cacheKey, entry)
- }
- // 超出缓存上限时尝试淘汰
- if enforceLimit && s.maxUpstreamClients() > 0 {
- s.evictIdleLocked(now)
- if len(s.clients) >= s.maxUpstreamClients() {
- if !s.evictOldestIdleLocked() {
- s.mu.Unlock()
- return nil, errUpstreamClientLimitReached
- }
- }
- }
- // 创建带 TLS 指纹的 Transport
- slog.Debug("tls_fingerprint_creating_new_client", "account_id", accountID, "cache_key", cacheKey)
- transport, err := buildUpstreamTransportWithTLSFingerprint(settings, parsedProxy, profile)
- if err != nil {
- s.mu.Unlock()
- return nil, fmt.Errorf("build TLS fingerprint transport: %w", err)
- }
- client := &http.Client{Transport: transport}
- if s.shouldValidateResolvedIP() {
- client.CheckRedirect = s.redirectChecker
- }
- entry := &upstreamClientEntry{
- client: client,
- proxyKey: proxyKey,
- poolKey: poolKey,
- }
- atomic.StoreInt64(&entry.lastUsed, nowUnix)
- if markInFlight {
- atomic.StoreInt64(&entry.inFlight, 1)
- }
- s.clients[cacheKey] = entry
- s.evictIdleLocked(now)
- s.evictOverLimitLocked()
- s.mu.Unlock()
- return entry, nil
- }
- func (s *httpUpstreamService) shouldValidateResolvedIP() bool {
- if s.cfg == nil {
- return false
- }
- if !s.cfg.Security.URLAllowlist.Enabled {
- return false
- }
- return !s.cfg.Security.URLAllowlist.AllowPrivateHosts
- }
- // validateRequestHost 校验请求主机的解析结果不落在回环、私网、链路本地或未指定地址。
- // 是否全局启用由 security.url_allowlist 决定;带 WithHTTPUpstreamPublicHostsOnly 标记的请求无论配置如何都校验。
- func (s *httpUpstreamService) validateRequestHost(req *http.Request) error {
- publicHostsOnly := req != nil && service.HTTPUpstreamPublicHostsOnly(req.Context())
- if !s.shouldValidateResolvedIP() && !publicHostsOnly {
- return nil
- }
- if req == nil || req.URL == nil {
- return errors.New("request url is nil")
- }
- host := strings.TrimSpace(req.URL.Hostname())
- if host == "" {
- return errors.New("request host is empty")
- }
- if err := urlvalidator.ValidateResolvedIP(host); err != nil {
- return err
- }
- return nil
- }
- func (s *httpUpstreamService) redirectChecker(req *http.Request, via []*http.Request) error {
- if len(via) >= 10 {
- return errors.New("stopped after 10 redirects")
- }
- return s.validateRequestHost(req)
- }
- // acquireClient 获取或创建客户端,并标记为进行中请求
- // 用于请求路径,避免在获取后被淘汰
- func (s *httpUpstreamService) acquireClient(proxyURL string, accountID int64, accountConcurrency int) (*upstreamClientEntry, error) {
- return s.acquireClientWithProfile(proxyURL, accountID, accountConcurrency, service.HTTPUpstreamProfileDefault)
- }
- // acquireClientWithProfile 获取或创建客户端,并按请求 profile 选择协议策略。
- func (s *httpUpstreamService) acquireClientWithProfile(proxyURL string, accountID int64, accountConcurrency int, profile service.HTTPUpstreamProfile) (*upstreamClientEntry, error) {
- return s.getClientEntry(proxyURL, accountID, accountConcurrency, profile, true, true)
- }
- // getOrCreateClient 获取或创建客户端
- // 根据隔离策略和参数决定缓存键,处理代理变更和配置变更
- //
- // 参数:
- // - proxyURL: 代理地址
- // - accountID: 账户 ID
- // - accountConcurrency: 账户并发限制
- //
- // 返回:
- // - *upstreamClientEntry: 客户端缓存条目
- //
- // 隔离策略说明:
- // - proxy: 按代理地址隔离,同一代理共享客户端
- // - account: 按账户隔离,同一账户共享客户端(代理变更时重建)
- // - account_proxy: 按账户+代理组合隔离,最细粒度
- func (s *httpUpstreamService) getOrCreateClient(proxyURL string, accountID int64, accountConcurrency int) (*upstreamClientEntry, error) {
- return s.getClientEntry(proxyURL, accountID, accountConcurrency, service.HTTPUpstreamProfileDefault, false, false)
- }
- // getClientEntry 获取或创建客户端条目
- // markInFlight=true 时会标记进行中请求,用于请求路径防止被淘汰
- // enforceLimit=true 时会限制客户端数量,超限且无法淘汰时返回错误
- func (s *httpUpstreamService) getClientEntry(proxyURL string, accountID int64, accountConcurrency int, profile service.HTTPUpstreamProfile, markInFlight bool, enforceLimit bool) (*upstreamClientEntry, error) {
- // 获取隔离模式
- isolation := s.getIsolationMode()
- // 标准化代理 URL 并解析
- proxyKey, parsedProxy, err := normalizeProxyURL(proxyURL)
- if err != nil {
- return nil, err
- }
- // 根据请求 profile(例如 OpenAI)选择协议模式
- protocolMode := s.resolveProtocolMode(profile, proxyKey, parsedProxy)
- settings := s.resolvePoolSettings(isolation, accountConcurrency)
- settings = s.applyProfilePoolSettings(settings, profile)
- // 构建缓存键(根据隔离策略不同)
- cacheKey := buildCacheKey(isolation, proxyKey, accountID, protocolMode)
- // 构建连接池配置键(用于检测配置变更)
- poolKey := buildPoolKey(settings, protocolMode)
- now := time.Now()
- nowUnix := now.UnixNano()
- // 读锁快速路径:命中缓存直接返回,减少锁竞争
- s.mu.RLock()
- if entry, ok := s.clients[cacheKey]; ok && s.shouldReuseEntry(entry, isolation, proxyKey, poolKey) {
- atomic.StoreInt64(&entry.lastUsed, nowUnix)
- if markInFlight {
- atomic.AddInt64(&entry.inFlight, 1)
- }
- s.mu.RUnlock()
- return entry, nil
- }
- s.mu.RUnlock()
- // 写锁慢路径:创建或重建客户端
- s.mu.Lock()
- if entry, ok := s.clients[cacheKey]; ok {
- if s.shouldReuseEntry(entry, isolation, proxyKey, poolKey) {
- atomic.StoreInt64(&entry.lastUsed, nowUnix)
- if markInFlight {
- atomic.AddInt64(&entry.inFlight, 1)
- }
- s.mu.Unlock()
- return entry, nil
- }
- s.removeClientLocked(cacheKey, entry)
- }
- // 超出缓存上限时尝试淘汰,无法淘汰则拒绝新建
- if enforceLimit && s.maxUpstreamClients() > 0 {
- s.evictIdleLocked(now)
- if len(s.clients) >= s.maxUpstreamClients() {
- if !s.evictOldestIdleLocked() {
- s.mu.Unlock()
- return nil, errUpstreamClientLimitReached
- }
- }
- }
- // 缓存未命中或需要重建,创建新客户端
- transport, err := buildUpstreamTransport(settings, parsedProxy, protocolMode)
- if err != nil {
- s.mu.Unlock()
- return nil, fmt.Errorf("build transport: %w", err)
- }
- client := &http.Client{Transport: transport}
- if s.shouldValidateResolvedIP() {
- client.CheckRedirect = s.redirectChecker
- }
- entry := &upstreamClientEntry{
- client: client,
- proxyKey: proxyKey,
- poolKey: poolKey,
- protocolMode: protocolMode,
- }
- atomic.StoreInt64(&entry.lastUsed, nowUnix)
- if markInFlight {
- atomic.StoreInt64(&entry.inFlight, 1)
- }
- s.clients[cacheKey] = entry
- // 执行淘汰策略:先淘汰空闲超时的,再淘汰超出数量限制的
- s.evictIdleLocked(now)
- s.evictOverLimitLocked()
- s.mu.Unlock()
- return entry, nil
- }
- // shouldReuseEntry 判断缓存条目是否可复用
- // 若代理或连接池配置发生变化,则需要重建客户端
- func (s *httpUpstreamService) shouldReuseEntry(entry *upstreamClientEntry, isolation, proxyKey, poolKey string) bool {
- if entry == nil {
- return false
- }
- if isolation == config.ConnectionPoolIsolationAccount && entry.proxyKey != proxyKey {
- return false
- }
- if entry.poolKey != poolKey {
- return false
- }
- return true
- }
- // removeClientLocked 移除客户端(需持有锁)
- // 从缓存中删除并关闭空闲连接
- //
- // 参数:
- // - key: 缓存键
- // - entry: 客户端条目
- func (s *httpUpstreamService) removeClientLocked(key string, entry *upstreamClientEntry) {
- delete(s.clients, key)
- if entry != nil && entry.client != nil {
- // 关闭空闲连接,释放系统资源
- // 注意:这不会中断活跃连接
- entry.client.CloseIdleConnections()
- }
- }
- // evictIdleLocked 淘汰空闲超时的客户端(需持有锁)
- // 遍历所有客户端,移除超过 TTL 且无活跃请求的条目
- //
- // 参数:
- // - now: 当前时间
- func (s *httpUpstreamService) evictIdleLocked(now time.Time) {
- ttl := s.clientIdleTTL()
- if ttl <= 0 {
- return
- }
- // 计算淘汰截止时间
- cutoff := now.Add(-ttl).UnixNano()
- for key, entry := range s.clients {
- // 跳过有活跃请求的客户端
- if atomic.LoadInt64(&entry.inFlight) != 0 {
- continue
- }
- // 淘汰超时的空闲客户端
- if atomic.LoadInt64(&entry.lastUsed) <= cutoff {
- s.removeClientLocked(key, entry)
- }
- }
- }
- // evictOldestIdleLocked 淘汰最久未使用且无活跃请求的客户端(需持有锁)
- func (s *httpUpstreamService) evictOldestIdleLocked() bool {
- var (
- oldestKey string
- oldestEntry *upstreamClientEntry
- oldestTime int64
- )
- // 查找最久未使用且无活跃请求的客户端
- for key, entry := range s.clients {
- // 跳过有活跃请求的客户端
- if atomic.LoadInt64(&entry.inFlight) != 0 {
- continue
- }
- lastUsed := atomic.LoadInt64(&entry.lastUsed)
- if oldestEntry == nil || lastUsed < oldestTime {
- oldestKey = key
- oldestEntry = entry
- oldestTime = lastUsed
- }
- }
- // 所有客户端都有活跃请求,无法淘汰
- if oldestEntry == nil {
- return false
- }
- s.removeClientLocked(oldestKey, oldestEntry)
- return true
- }
- // evictOverLimitLocked 淘汰超出数量限制的客户端(需持有锁)
- // 使用 LRU 策略,优先淘汰最久未使用且无活跃请求的客户端
- func (s *httpUpstreamService) evictOverLimitLocked() bool {
- maxClients := s.maxUpstreamClients()
- if maxClients <= 0 {
- return false
- }
- evicted := false
- // 循环淘汰直到满足数量限制
- for len(s.clients) > maxClients {
- if !s.evictOldestIdleLocked() {
- return evicted
- }
- evicted = true
- }
- return evicted
- }
- // getIsolationMode 获取连接池隔离模式
- // 从配置中读取,无效值回退到 account_proxy 模式
- //
- // 返回:
- // - string: 隔离模式(proxy/account/account_proxy)
- func (s *httpUpstreamService) getIsolationMode() string {
- if s.cfg == nil {
- return config.ConnectionPoolIsolationAccountProxy
- }
- mode := strings.ToLower(strings.TrimSpace(s.cfg.Gateway.ConnectionPoolIsolation))
- if mode == "" {
- return config.ConnectionPoolIsolationAccountProxy
- }
- switch mode {
- case config.ConnectionPoolIsolationProxy, config.ConnectionPoolIsolationAccount, config.ConnectionPoolIsolationAccountProxy:
- return mode
- default:
- return config.ConnectionPoolIsolationAccountProxy
- }
- }
- // maxUpstreamClients 获取最大客户端缓存数量
- // 从配置中读取,无效值使用默认值
- func (s *httpUpstreamService) maxUpstreamClients() int {
- if s.cfg == nil {
- return defaultMaxUpstreamClients
- }
- if s.cfg.Gateway.MaxUpstreamClients > 0 {
- return s.cfg.Gateway.MaxUpstreamClients
- }
- return defaultMaxUpstreamClients
- }
- // clientIdleTTL 获取客户端空闲回收阈值
- // 从配置中读取,无效值使用默认值
- func (s *httpUpstreamService) clientIdleTTL() time.Duration {
- if s.cfg == nil {
- return time.Duration(defaultClientIdleTTLSeconds) * time.Second
- }
- if s.cfg.Gateway.ClientIdleTTLSeconds > 0 {
- return time.Duration(s.cfg.Gateway.ClientIdleTTLSeconds) * time.Second
- }
- return time.Duration(defaultClientIdleTTLSeconds) * time.Second
- }
- // resolvePoolSettings 解析连接池配置
- // 根据隔离策略和账户并发数动态调整连接池参数
- //
- // 参数:
- // - isolation: 隔离模式
- // - accountConcurrency: 账户并发限制
- //
- // 返回:
- // - poolSettings: 连接池配置
- //
- // 说明:
- // - 账户隔离模式下,连接池大小与账户并发数对应
- // - 这确保了单账户不会占用过多连接资源
- func (s *httpUpstreamService) resolvePoolSettings(isolation string, accountConcurrency int) poolSettings {
- settings := defaultPoolSettings(s.cfg)
- // 账户隔离模式下,根据账户并发数调整连接池大小
- if (isolation == config.ConnectionPoolIsolationAccount || isolation == config.ConnectionPoolIsolationAccountProxy) && accountConcurrency > 0 {
- settings.maxIdleConns = accountConcurrency
- settings.maxIdleConnsPerHost = accountConcurrency
- settings.maxConnsPerHost = accountConcurrency
- }
- return settings
- }
- func (s *httpUpstreamService) applyProfilePoolSettings(settings poolSettings, profile service.HTTPUpstreamProfile) poolSettings {
- switch profile {
- case service.HTTPUpstreamProfileOpenAI:
- settings.responseHeaderTimeout = 0
- if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIResponseHeaderTimeout > 0 {
- settings.responseHeaderTimeout = time.Duration(s.cfg.Gateway.OpenAIResponseHeaderTimeout) * time.Second
- }
- case service.HTTPUpstreamProfileGrok:
- // Grok can stall before its first byte under capacity pressure. Keep the
- // generic 600s gateway timeout from turning one request into a 10-minute
- // resource hold; streaming after headers is unaffected.
- settings.responseHeaderTimeout = 120 * time.Second
- if s != nil && s.cfg != nil {
- settings.responseHeaderTimeout = time.Duration(s.cfg.Gateway.GrokResponseHeaderTimeout) * time.Second
- }
- }
- return settings
- }
- // buildPoolKey 构建连接池配置键,用于检测连接池配置变更。
- func buildPoolKey(settings poolSettings, protocolMode string) string {
- base := fmt.Sprintf(
- "idle:%d|idle_host:%d|max:%d|idle_timeout:%s|header_timeout:%s",
- settings.maxIdleConns,
- settings.maxIdleConnsPerHost,
- settings.maxConnsPerHost,
- settings.idleConnTimeout,
- settings.responseHeaderTimeout,
- )
- if protocolMode == "" || protocolMode == upstreamProtocolModeDefault {
- return base
- }
- return base + "|proto:" + protocolMode
- }
- // buildCacheKey 构建客户端缓存键
- // 根据隔离策略决定缓存键的组成
- //
- // 参数:
- // - isolation: 隔离模式
- // - proxyKey: 代理标识
- // - accountID: 账户 ID
- //
- // 返回:
- // - string: 缓存键
- //
- // 缓存键格式:
- // - proxy 模式: "proxy:{proxyKey}"
- // - account 模式: "account:{accountID}"
- // - account_proxy 模式: "account:{accountID}|proxy:{proxyKey}"
- func buildCacheKey(isolation, proxyKey string, accountID int64, protocolMode string) string {
- var base string
- switch isolation {
- case config.ConnectionPoolIsolationAccount:
- base = fmt.Sprintf("account:%d", accountID)
- case config.ConnectionPoolIsolationAccountProxy:
- base = fmt.Sprintf("account:%d|proxy:%s", accountID, proxyKey)
- default:
- base = fmt.Sprintf("proxy:%s", proxyKey)
- }
- if protocolMode != "" && protocolMode != upstreamProtocolModeDefault {
- base += "|proto:" + protocolMode
- }
- return base
- }
- func (s *httpUpstreamService) resolveOpenAIHTTP2Settings() openAIHTTP2Settings {
- settings := openAIHTTP2Settings{
- enabled: false,
- allowProxyFallbackToHTTP1: true,
- fallbackErrorThreshold: defaultOpenAIHTTP2FallbackErrorThreshold,
- fallbackWindow: defaultOpenAIHTTP2FallbackWindow,
- fallbackTTL: defaultOpenAIHTTP2FallbackTTL,
- }
- if s == nil || s.cfg == nil {
- return settings
- }
- cfg := s.cfg.Gateway.OpenAIHTTP2
- settings.enabled = cfg.Enabled
- settings.allowProxyFallbackToHTTP1 = cfg.AllowProxyFallbackToHTTP1
- if cfg.FallbackErrorThreshold > 0 {
- settings.fallbackErrorThreshold = cfg.FallbackErrorThreshold
- }
- if cfg.FallbackWindowSeconds > 0 {
- settings.fallbackWindow = time.Duration(cfg.FallbackWindowSeconds) * time.Second
- }
- if cfg.FallbackTTLSeconds > 0 {
- settings.fallbackTTL = time.Duration(cfg.FallbackTTLSeconds) * time.Second
- }
- return settings
- }
- func (s *httpUpstreamService) resolveProtocolMode(profile service.HTTPUpstreamProfile, proxyKey string, parsedProxy *url.URL) string {
- if profile == service.HTTPUpstreamProfileLongStream {
- return upstreamProtocolModeLongStreamH2
- }
- if profile == service.HTTPUpstreamProfileGrok {
- return upstreamProtocolModeGrok
- }
- if profile != service.HTTPUpstreamProfileOpenAI {
- return upstreamProtocolModeDefault
- }
- settings := s.resolveOpenAIHTTP2Settings()
- if !settings.enabled {
- return upstreamProtocolModeOpenAIH1
- }
- if parsedProxy == nil {
- return upstreamProtocolModeOpenAIH2
- }
- scheme := strings.ToLower(parsedProxy.Scheme)
- if scheme != "http" && scheme != "https" {
- return upstreamProtocolModeOpenAIH2
- }
- if settings.allowProxyFallbackToHTTP1 && s.isOpenAIHTTP2FallbackActive(proxyKey) {
- return upstreamProtocolModeOpenAIH1Fallback
- }
- return upstreamProtocolModeOpenAIH2
- }
- func (s *httpUpstreamService) isOpenAIHTTP2FallbackActive(proxyKey string) bool {
- raw, ok := s.openAIHTTP2Fallbacks.Load(proxyKey)
- if !ok {
- return false
- }
- state, ok := raw.(*openAIHTTP2FallbackState)
- if !ok || state == nil {
- return false
- }
- return state.isFallbackActive(time.Now())
- }
- func (s *httpUpstreamService) getOrCreateOpenAIHTTP2FallbackState(proxyKey string) *openAIHTTP2FallbackState {
- state := &openAIHTTP2FallbackState{}
- actual, _ := s.openAIHTTP2Fallbacks.LoadOrStore(proxyKey, state)
- cached, ok := actual.(*openAIHTTP2FallbackState)
- if !ok || cached == nil {
- return state
- }
- return cached
- }
- func isHTTPProxyKey(proxyKey string) bool {
- return strings.HasPrefix(proxyKey, "http://") || strings.HasPrefix(proxyKey, "https://")
- }
- func isOpenAIHTTP2CompatibilityError(err error) bool {
- if err == nil {
- return false
- }
- if isUpstreamTimeoutError(err) {
- return false
- }
- msg := strings.ToLower(err.Error())
- if msg == "" {
- return false
- }
- markers := []string{
- "alpn",
- "no application protocol",
- "protocol error",
- "stream error",
- "goaway",
- "refused_stream",
- "frame too large",
- }
- for _, marker := range markers {
- if strings.Contains(msg, marker) {
- return true
- }
- }
- return false
- }
- func isUpstreamTimeoutError(err error) bool {
- if err == nil {
- return false
- }
- if errors.Is(err, context.DeadlineExceeded) {
- return true
- }
- var netErr net.Error
- if errors.As(err, &netErr) && netErr.Timeout() {
- return true
- }
- msg := strings.ToLower(err.Error())
- if msg == "" {
- return false
- }
- timeoutMarkers := []string{
- "timeout awaiting response headers",
- "i/o timeout",
- "context deadline exceeded",
- "client.timeout exceeded while awaiting headers",
- "tls handshake timeout",
- }
- for _, marker := range timeoutMarkers {
- if strings.Contains(msg, marker) {
- return true
- }
- }
- return false
- }
- func (s *httpUpstreamService) recordOpenAIHTTP2Failure(profile service.HTTPUpstreamProfile, protocolMode, proxyKey string, err error) {
- if profile != service.HTTPUpstreamProfileOpenAI || protocolMode != upstreamProtocolModeOpenAIH2 {
- return
- }
- settings := s.resolveOpenAIHTTP2Settings()
- if !settings.enabled || !settings.allowProxyFallbackToHTTP1 {
- return
- }
- if !isHTTPProxyKey(proxyKey) || !isOpenAIHTTP2CompatibilityError(err) {
- return
- }
- state := s.getOrCreateOpenAIHTTP2FallbackState(proxyKey)
- activated, until := state.recordFailure(time.Now(), settings.fallbackErrorThreshold, settings.fallbackWindow, settings.fallbackTTL)
- if activated {
- slog.Warn("openai_http2_proxy_fallback_activated",
- "proxy", proxyKey,
- "fallback_until", until.Format(time.RFC3339))
- }
- }
- func (s *httpUpstreamService) recordOpenAIHTTP2Success(profile service.HTTPUpstreamProfile, protocolMode, proxyKey string) {
- if profile != service.HTTPUpstreamProfileOpenAI || protocolMode != upstreamProtocolModeOpenAIH2 {
- return
- }
- if !isHTTPProxyKey(proxyKey) {
- return
- }
- raw, ok := s.openAIHTTP2Fallbacks.Load(proxyKey)
- if !ok {
- return
- }
- state, ok := raw.(*openAIHTTP2FallbackState)
- if !ok || state == nil {
- return
- }
- state.resetErrorWindow()
- }
- func (s *openAIHTTP2FallbackState) isFallbackActive(now time.Time) bool {
- s.mu.Lock()
- defer s.mu.Unlock()
- if s.fallbackUntil.IsZero() {
- return false
- }
- if now.Before(s.fallbackUntil) {
- return true
- }
- s.fallbackUntil = time.Time{}
- return false
- }
- func (s *openAIHTTP2FallbackState) resetErrorWindow() {
- s.mu.Lock()
- defer s.mu.Unlock()
- s.windowStart = time.Time{}
- s.errorCount = 0
- }
- func (s *openAIHTTP2FallbackState) recordFailure(now time.Time, threshold int, window, ttl time.Duration) (bool, time.Time) {
- if threshold <= 0 {
- threshold = defaultOpenAIHTTP2FallbackErrorThreshold
- }
- if window <= 0 {
- window = defaultOpenAIHTTP2FallbackWindow
- }
- if ttl <= 0 {
- ttl = defaultOpenAIHTTP2FallbackTTL
- }
- s.mu.Lock()
- defer s.mu.Unlock()
- if !s.fallbackUntil.IsZero() && now.Before(s.fallbackUntil) {
- return false, s.fallbackUntil
- }
- if !s.fallbackUntil.IsZero() && !now.Before(s.fallbackUntil) {
- s.fallbackUntil = time.Time{}
- }
- if s.windowStart.IsZero() || now.Sub(s.windowStart) > window {
- s.windowStart = now
- s.errorCount = 0
- }
- s.errorCount++
- if s.errorCount < threshold {
- return false, time.Time{}
- }
- s.fallbackUntil = now.Add(ttl)
- s.windowStart = time.Time{}
- s.errorCount = 0
- return true, s.fallbackUntil
- }
- // normalizeProxyURL 标准化代理 URL
- // 处理空值和解析错误,返回标准化的键和解析后的 URL
- //
- // 参数:
- // - raw: 原始代理 URL 字符串
- //
- // 返回:
- // - string: 标准化的代理键(空返回 "direct")
- // - *url.URL: 解析后的 URL(空返回 nil)
- // - error: 非空代理 URL 解析失败时返回错误(禁止回退到直连)
- func normalizeProxyURL(raw string) (string, *url.URL, error) {
- _, parsed, err := proxyurl.Parse(raw)
- if err != nil {
- var ferr error
- _, parsed, ferr = proxyurl.ParseFlexible(raw, "http")
- if ferr != nil {
- return "", nil, err
- }
- }
- if parsed == nil {
- return directProxyKey, nil, nil
- }
- // 规范化:小写 scheme/host,去除路径和查询参数
- parsed.Scheme = strings.ToLower(parsed.Scheme)
- parsed.Host = strings.ToLower(parsed.Host)
- parsed.Path = ""
- parsed.RawPath = ""
- parsed.RawQuery = ""
- parsed.Fragment = ""
- parsed.ForceQuery = false
- if hostname := parsed.Hostname(); hostname != "" {
- port := parsed.Port()
- if (parsed.Scheme == "http" && port == "80") || (parsed.Scheme == "https" && port == "443") {
- port = ""
- }
- hostname = strings.ToLower(hostname)
- if port != "" {
- parsed.Host = net.JoinHostPort(hostname, port)
- } else {
- parsed.Host = hostname
- }
- }
- return parsed.String(), parsed, nil
- }
- // defaultPoolSettings 获取默认连接池配置
- // 从全局配置中读取,无效值使用常量默认值
- //
- // 参数:
- // - cfg: 全局配置
- //
- // 返回:
- // - poolSettings: 连接池配置
- func defaultPoolSettings(cfg *config.Config) poolSettings {
- maxIdleConns := defaultMaxIdleConns
- maxIdleConnsPerHost := defaultMaxIdleConnsPerHost
- maxConnsPerHost := defaultMaxConnsPerHost
- idleConnTimeout := defaultIdleConnTimeout
- responseHeaderTimeout := defaultResponseHeaderTimeout
- if cfg != nil {
- if cfg.Gateway.MaxIdleConns > 0 {
- maxIdleConns = cfg.Gateway.MaxIdleConns
- }
- if cfg.Gateway.MaxIdleConnsPerHost > 0 {
- maxIdleConnsPerHost = cfg.Gateway.MaxIdleConnsPerHost
- }
- if cfg.Gateway.MaxConnsPerHost >= 0 {
- maxConnsPerHost = cfg.Gateway.MaxConnsPerHost
- }
- if cfg.Gateway.IdleConnTimeoutSeconds > 0 {
- idleConnTimeout = time.Duration(cfg.Gateway.IdleConnTimeoutSeconds) * time.Second
- }
- if cfg.Gateway.ResponseHeaderTimeout >= 0 {
- responseHeaderTimeout = time.Duration(cfg.Gateway.ResponseHeaderTimeout) * time.Second
- }
- }
- return poolSettings{
- maxIdleConns: maxIdleConns,
- maxIdleConnsPerHost: maxIdleConnsPerHost,
- maxConnsPerHost: maxConnsPerHost,
- idleConnTimeout: idleConnTimeout,
- responseHeaderTimeout: responseHeaderTimeout,
- }
- }
- // newUpstreamDialer 构建上游 Transport 的 TCP dialer。
- //
- // 必须显式提供:http.Transport 的 DialContext 为 nil 时使用零值 net.Dialer,
- // 建连没有任何超时上限,只能等内核 TCP 重传耗尽(Linux 约 130 秒)。
- func newUpstreamDialer() *net.Dialer {
- return &net.Dialer{
- Timeout: defaultUpstreamDialTimeout,
- KeepAlive: defaultUpstreamDialKeepAlive,
- }
- }
- // buildUpstreamTransport 构建上游请求的 Transport
- // 使用配置文件中的连接池参数,支持生产环境调优
- //
- // 参数:
- // - settings: 连接池配置
- // - proxyURL: 代理 URL(nil 表示直连)
- //
- // 返回:
- // - *http.Transport: 配置好的 Transport 实例
- // - error: 代理配置错误
- //
- // Transport 参数说明:
- // - DialContext: DNS 解析 + TCP 建连超时(不设置则无上限,退化为内核默认重传)
- // - TLSHandshakeTimeout: TLS 握手超时
- // - MaxIdleConns: 所有主机的最大空闲连接总数
- // - MaxIdleConnsPerHost: 每主机最大空闲连接数(影响连接复用率)
- // - MaxConnsPerHost: 每主机最大连接数(达到后新请求等待)
- // - IdleConnTimeout: 空闲连接超时(超时后关闭)
- // - ResponseHeaderTimeout: 等待响应头超时(不影响流式传输)
- func buildUpstreamTransport(settings poolSettings, proxyURL *url.URL, protocolMode string) (*http.Transport, error) {
- transport := &http.Transport{
- DialContext: newUpstreamDialer().DialContext,
- TLSHandshakeTimeout: defaultUpstreamTLSHandshakeTimeout,
- MaxIdleConns: settings.maxIdleConns,
- MaxIdleConnsPerHost: settings.maxIdleConnsPerHost,
- MaxConnsPerHost: settings.maxConnsPerHost,
- IdleConnTimeout: settings.idleConnTimeout,
- ResponseHeaderTimeout: settings.responseHeaderTimeout,
- }
- switch protocolMode {
- case upstreamProtocolModeLongStreamH2, upstreamProtocolModeOpenAIH2:
- transport.ForceAttemptHTTP2 = true
- // 显式配置 http2 并启用 PING 健康探测,剔除代理/NAT 静默掐断的死连接,
- // 避免请求挂在死连接上直到 TCP 重传超时(分钟级)。
- if _, err := enableHTTP2KeepAlive(transport); err != nil {
- return nil, err
- }
- case upstreamProtocolModeOpenAIH1:
- transport.ForceAttemptHTTP2 = false
- transport.TLSNextProto = make(map[string]func(string, *tls.Conn) http.RoundTripper)
- case upstreamProtocolModeOpenAIH1Fallback:
- // 显式禁用 HTTP/2,确保代理不兼容场景回退到 HTTP/1.1。
- transport.ForceAttemptHTTP2 = false
- transport.TLSNextProto = make(map[string]func(string, *tls.Conn) http.RoundTripper)
- }
- if err := proxyutil.ConfigureTransportProxy(transport, proxyURL); err != nil {
- return nil, err
- }
- return transport, nil
- }
- // enableHTTP2KeepAlive 在 http.Transport 上显式配置 HTTP/2 并启用连接健康探测。
- // Go 默认惰性配置 http2 且 ReadIdleTimeout=0(不发健康 PING),无法检测被代理/NAT
- // 静默掐断的死连接。此处主动设置 ReadIdleTimeout/PingTimeout,让死连接被提前 PING
- // 出并关闭,请求得以重建连接而非挂到 TCP 重传超时。返回底层 *http2.Transport 便于测试。
- func enableHTTP2KeepAlive(transport *http.Transport) (*http2.Transport, error) {
- h2, err := http2.ConfigureTransports(transport)
- if err != nil {
- return nil, err
- }
- if h2 != nil {
- h2.ReadIdleTimeout = longStreamHTTP2ReadIdleTimeout
- h2.PingTimeout = longStreamHTTP2PingTimeout
- }
- return h2, nil
- }
- // buildUpstreamTransportWithTLSFingerprint 构建带 TLS 指纹伪装的 Transport
- // 使用 utls 库模拟 Claude CLI 的 TLS 指纹
- //
- // 参数:
- // - settings: 连接池配置
- // - proxyURL: 代理 URL(nil 表示直连)
- // - profile: TLS 指纹配置
- //
- // 返回:
- // - *http.Transport: 配置好的 Transport 实例
- // - error: 配置错误
- //
- // 代理类型处理:
- // - nil/空: 直连,使用 TLSFingerprintDialer
- // - http/https: HTTP 代理,使用 HTTPProxyDialer(CONNECT 隧道 + utls 握手)
- // - socks5: SOCKS5 代理,使用 SOCKS5ProxyDialer(SOCKS5 隧道 + utls 握手)
- func buildUpstreamTransportWithTLSFingerprint(settings poolSettings, proxyURL *url.URL, profile *tlsfingerprint.Profile) (*http.Transport, error) {
- transport := &http.Transport{
- MaxIdleConns: settings.maxIdleConns,
- MaxIdleConnsPerHost: settings.maxIdleConnsPerHost,
- MaxConnsPerHost: settings.maxConnsPerHost,
- IdleConnTimeout: settings.idleConnTimeout,
- ResponseHeaderTimeout: settings.responseHeaderTimeout,
- // 禁用默认的 TLS,我们使用自定义的 DialTLSContext
- ForceAttemptHTTP2: false,
- }
- if err := tlsfingerprint.ConfigureTransport(transport, profile, proxyURL, tlsfingerprint.TransportOptions{
- DialContext: newUpstreamDialer().DialContext, HandshakeTimeout: defaultUpstreamTLSHandshakeTimeout,
- }); err != nil {
- return nil, err
- }
- return transport, nil
- }
- // trackedBody 带跟踪功能的响应体包装器
- // 在 Close 时执行回调,用于更新请求计数
- type trackedBody struct {
- io.ReadCloser // 原始响应体
- once sync.Once
- onClose func() // 关闭时的回调函数
- }
- // Close 关闭响应体并执行回调
- // 使用 sync.Once 确保回调只执行一次
- func (b *trackedBody) Close() error {
- err := b.ReadCloser.Close()
- if b.onClose != nil {
- b.once.Do(b.onClose)
- }
- return err
- }
- // wrapTrackedBody 包装响应体以跟踪关闭事件
- // 用于在响应体关闭时更新 inFlight 计数
- //
- // 参数:
- // - body: 原始响应体
- // - onClose: 关闭时的回调函数
- //
- // 返回:
- // - io.ReadCloser: 包装后的响应体
- func wrapTrackedBody(body io.ReadCloser, onClose func()) io.ReadCloser {
- if body == nil {
- return body
- }
- return &trackedBody{ReadCloser: body, onClose: onClose}
- }
- // decompressResponseBody 根据 Content-Encoding 解压响应体。
- // 当请求显式设置了 accept-encoding 时,Go 的 Transport 不会自动解压,需要手动处理。
- // 解压成功后会删除 Content-Encoding 和 Content-Length header(长度已不准确)。
- func decompressResponseBody(resp *http.Response) {
- if resp == nil || resp.Body == nil {
- return
- }
- ce := strings.ToLower(strings.TrimSpace(resp.Header.Get("Content-Encoding")))
- if ce == "" {
- return
- }
- originalBody := resp.Body
- var reader io.Reader
- switch ce {
- case "gzip":
- gr, err := gzip.NewReader(resp.Body)
- if err != nil {
- return // 解压失败,保持原样
- }
- reader = gr
- case "br":
- reader = brotli.NewReader(resp.Body)
- case "deflate":
- reader = flate.NewReader(resp.Body)
- case "zstd":
- bufferedBody := bufio.NewReader(resp.Body)
- resp.Body = &decompressedBody{reader: bufferedBody, closer: originalBody}
- headerBytes, _ := bufferedBody.Peek(zstd.HeaderMaxSize)
- var header zstd.Header
- if err := header.Decode(headerBytes); err != nil {
- slog.Warn("zstd_decompress_failed", "error", err)
- return
- }
- zr, err := zstd.NewReader(bufferedBody)
- if err != nil {
- slog.Warn("zstd_decompress_failed", "error", err)
- return
- }
- reader = &zstdResponseReader{ReadCloser: zr.IOReadCloser()}
- default:
- return
- }
- resp.Body = &decompressedBody{reader: reader, closer: originalBody}
- resp.Header.Del("Content-Encoding")
- resp.Header.Del("Content-Length") // 解压后长度不确定
- resp.ContentLength = -1
- }
- type zstdResponseReader struct {
- io.ReadCloser
- warnOnce sync.Once
- }
- func (r *zstdResponseReader) Read(p []byte) (int, error) {
- n, err := r.ReadCloser.Read(p)
- if err != nil && !errors.Is(err, io.EOF) {
- r.warnOnce.Do(func() {
- slog.Warn("zstd_decompress_failed", "error", err)
- })
- }
- return n, err
- }
- // decompressedBody 组合解压 reader 和原始 body 的 close。
- type decompressedBody struct {
- reader io.Reader
- closer io.Closer
- }
- func (d *decompressedBody) Read(p []byte) (int, error) {
- return d.reader.Read(p)
- }
- func (d *decompressedBody) Close() error {
- // 如果 reader 本身也是 Closer(如 gzip.Reader),先关闭它
- if rc, ok := d.reader.(io.Closer); ok {
- _ = rc.Close()
- }
- return d.closer.Close()
- }
|