pool.go 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218
  1. // Package httpclient 提供共享 HTTP 客户端池
  2. //
  3. // 性能优化说明:
  4. // 原实现在多个服务中重复创建 http.Client:
  5. // 1. proxy_probe_service.go: 每次探测创建新客户端
  6. // 2. pricing_service.go: 每次请求创建新客户端
  7. // 3. turnstile_service.go: 每次验证创建新客户端
  8. // 4. github_release_service.go: 每次请求创建新客户端
  9. // 5. claude_usage_service.go: 每次请求创建新客户端
  10. //
  11. // 新实现使用统一的客户端池:
  12. // 1. 相同配置复用同一 http.Client 实例
  13. // 2. 复用 Transport 连接池,减少 TCP/TLS 握手开销
  14. // 3. 支持 HTTP/HTTPS/SOCKS5/SOCKS5H 代理
  15. // 4. 代理配置失败时直接返回错误,不会回退到直连(避免 IP 关联风险)
  16. package httpclient
  17. import (
  18. "fmt"
  19. "net"
  20. "net/http"
  21. "strings"
  22. "sync"
  23. "time"
  24. "github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
  25. "github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
  26. "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
  27. "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
  28. )
  29. // Transport 连接池默认配置
  30. const (
  31. defaultMaxIdleConns = 100 // 最大空闲连接数
  32. defaultMaxIdleConnsPerHost = 10 // 每个主机最大空闲连接数
  33. defaultIdleConnTimeout = 90 * time.Second // 空闲连接超时时间(建议小于上游 LB 超时)
  34. defaultDialTimeout = 5 * time.Second // TCP 连接超时(含代理握手),代理不通时快速失败
  35. defaultTLSHandshakeTimeout = 5 * time.Second // TLS 握手超时
  36. validatedHostTTL = 30 * time.Second // DNS Rebinding 校验缓存 TTL
  37. )
  38. // Options 定义共享 HTTP 客户端的构建参数
  39. type Options struct {
  40. ProxyURL string // 代理 URL(支持 http/https/socks5/socks5h)
  41. Timeout time.Duration // 请求总超时时间
  42. ResponseHeaderTimeout time.Duration // 等待响应头超时时间
  43. InsecureSkipVerify bool // 是否跳过 TLS 证书验证(已禁用,不允许设置为 true)
  44. ValidateResolvedIP bool // 是否校验解析后的 IP(防止 DNS Rebinding)
  45. AllowPrivateHosts bool // 允许私有地址解析(与 ValidateResolvedIP 一起使用)
  46. // 可选的连接池参数(不设置则使用默认值)
  47. MaxIdleConns int // 最大空闲连接总数(默认 100)
  48. MaxIdleConnsPerHost int // 每主机最大空闲连接(默认 10)
  49. MaxConnsPerHost int // 每主机最大连接数(默认 0 无限制)
  50. }
  51. // sharedClients 存储按配置参数缓存的 http.Client 实例
  52. var sharedClients sync.Map
  53. // 允许测试替换校验函数,生产默认指向真实实现。
  54. var validateResolvedIP = urlvalidator.ValidateResolvedIP
  55. // GetClient 返回共享的 HTTP 客户端实例
  56. // 性能优化:相同配置复用同一客户端,避免重复创建 Transport
  57. // 安全说明:代理配置失败时直接返回错误,不会回退到直连,避免 IP 关联风险
  58. func GetClient(opts Options) (*http.Client, error) {
  59. key := buildClientKey(opts)
  60. if cached, ok := sharedClients.Load(key); ok {
  61. if client, ok := cached.(*http.Client); ok {
  62. return client, nil
  63. }
  64. }
  65. client, err := buildClient(opts)
  66. if err != nil {
  67. return nil, err
  68. }
  69. actual, _ := sharedClients.LoadOrStore(key, client)
  70. if c, ok := actual.(*http.Client); ok {
  71. return c, nil
  72. }
  73. return client, nil
  74. }
  75. func buildClient(opts Options) (*http.Client, error) {
  76. transport, err := buildTransport(opts)
  77. if err != nil {
  78. return nil, err
  79. }
  80. var rt http.RoundTripper = transport
  81. if opts.ValidateResolvedIP && !opts.AllowPrivateHosts {
  82. rt = newValidatedTransport(transport)
  83. }
  84. rt = servertiming.WrapRoundTripper(rt)
  85. return &http.Client{
  86. Transport: rt,
  87. Timeout: opts.Timeout,
  88. }, nil
  89. }
  90. func buildTransport(opts Options) (*http.Transport, error) {
  91. // 使用自定义值或默认值
  92. maxIdleConns := opts.MaxIdleConns
  93. if maxIdleConns <= 0 {
  94. maxIdleConns = defaultMaxIdleConns
  95. }
  96. maxIdleConnsPerHost := opts.MaxIdleConnsPerHost
  97. if maxIdleConnsPerHost <= 0 {
  98. maxIdleConnsPerHost = defaultMaxIdleConnsPerHost
  99. }
  100. transport := &http.Transport{
  101. DialContext: (&net.Dialer{
  102. Timeout: defaultDialTimeout,
  103. }).DialContext,
  104. TLSHandshakeTimeout: defaultTLSHandshakeTimeout,
  105. MaxIdleConns: maxIdleConns,
  106. MaxIdleConnsPerHost: maxIdleConnsPerHost,
  107. MaxConnsPerHost: opts.MaxConnsPerHost, // 0 表示无限制
  108. IdleConnTimeout: defaultIdleConnTimeout,
  109. ResponseHeaderTimeout: opts.ResponseHeaderTimeout,
  110. }
  111. if opts.InsecureSkipVerify {
  112. // 安全要求:禁止跳过证书验证,避免中间人攻击。
  113. return nil, fmt.Errorf("insecure_skip_verify is not allowed; install a trusted certificate instead")
  114. }
  115. _, parsed, err := proxyurl.Parse(opts.ProxyURL)
  116. if err != nil {
  117. // 尝试多格式兼容解析(host:port:user:pass 等),默认 http
  118. var ferr error
  119. _, parsed, ferr = proxyurl.ParseFlexible(opts.ProxyURL, "http")
  120. if ferr != nil {
  121. return nil, err
  122. }
  123. }
  124. if parsed == nil {
  125. return transport, nil
  126. }
  127. if err := proxyutil.ConfigureTransportProxy(transport, parsed); err != nil {
  128. return nil, err
  129. }
  130. return transport, nil
  131. }
  132. func buildClientKey(opts Options) string {
  133. return fmt.Sprintf("%s|%s|%s|%t|%t|%t|%d|%d|%d",
  134. strings.TrimSpace(opts.ProxyURL),
  135. opts.Timeout.String(),
  136. opts.ResponseHeaderTimeout.String(),
  137. opts.InsecureSkipVerify,
  138. opts.ValidateResolvedIP,
  139. opts.AllowPrivateHosts,
  140. opts.MaxIdleConns,
  141. opts.MaxIdleConnsPerHost,
  142. opts.MaxConnsPerHost,
  143. )
  144. }
  145. type validatedTransport struct {
  146. base http.RoundTripper
  147. validatedHosts sync.Map // map[string]time.Time, value 为过期时间
  148. now func() time.Time
  149. }
  150. func newValidatedTransport(base http.RoundTripper) *validatedTransport {
  151. return &validatedTransport{
  152. base: base,
  153. now: time.Now,
  154. }
  155. }
  156. func (t *validatedTransport) isValidatedHost(host string, now time.Time) bool {
  157. if t == nil {
  158. return false
  159. }
  160. raw, ok := t.validatedHosts.Load(host)
  161. if !ok {
  162. return false
  163. }
  164. expireAt, ok := raw.(time.Time)
  165. if !ok {
  166. t.validatedHosts.Delete(host)
  167. return false
  168. }
  169. if now.Before(expireAt) {
  170. return true
  171. }
  172. t.validatedHosts.Delete(host)
  173. return false
  174. }
  175. func (t *validatedTransport) RoundTrip(req *http.Request) (*http.Response, error) {
  176. if req != nil && req.URL != nil {
  177. host := strings.ToLower(strings.TrimSpace(req.URL.Hostname()))
  178. if host != "" {
  179. now := time.Now()
  180. if t != nil && t.now != nil {
  181. now = t.now()
  182. }
  183. if !t.isValidatedHost(host, now) {
  184. if err := validateResolvedIP(host); err != nil {
  185. return nil, err
  186. }
  187. t.validatedHosts.Store(host, now.Add(validatedHostTTL))
  188. }
  189. }
  190. }
  191. if t == nil || t.base == nil {
  192. return nil, fmt.Errorf("validated transport base is nil")
  193. }
  194. return t.base.RoundTrip(req)
  195. }