You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

552 lines
16 KiB

  1. package channel
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "bytes"
  7. "io"
  8. "net/http"
  9. "regexp"
  10. "strings"
  11. "sync"
  12. "time"
  13. common2 "github.com/QuantumNous/new-api/common"
  14. "github.com/QuantumNous/new-api/logger"
  15. "github.com/QuantumNous/new-api/relay/common"
  16. "github.com/QuantumNous/new-api/relay/constant"
  17. "github.com/QuantumNous/new-api/relay/helper"
  18. "github.com/QuantumNous/new-api/service"
  19. "github.com/QuantumNous/new-api/setting/operation_setting"
  20. "github.com/QuantumNous/new-api/types"
  21. "github.com/bytedance/gopkg/util/gopool"
  22. "github.com/gin-gonic/gin"
  23. "github.com/gorilla/websocket"
  24. )
  25. func SetupApiRequestHeader(info *common.RelayInfo, c *gin.Context, req *http.Header) {
  26. if info.RelayMode == constant.RelayModeAudioTranscription || info.RelayMode == constant.RelayModeAudioTranslation {
  27. // multipart/form-data
  28. } else if info.RelayMode == constant.RelayModeRealtime {
  29. // websocket
  30. } else {
  31. req.Set("Content-Type", c.Request.Header.Get("Content-Type"))
  32. req.Set("Accept", c.Request.Header.Get("Accept"))
  33. if info.IsStream && c.Request.Header.Get("Accept") == "" {
  34. req.Set("Accept", "text/event-stream")
  35. }
  36. }
  37. }
  38. const clientHeaderPlaceholderPrefix = "{client_header:"
  39. const (
  40. headerPassthroughAllKey = "*"
  41. headerPassthroughRegexPrefix = "re:"
  42. headerPassthroughRegexPrefixV2 = "regex:"
  43. )
  44. var passthroughSkipHeaderNamesLower = map[string]struct{}{
  45. // RFC 7230 hop-by-hop headers.
  46. "connection": {},
  47. "keep-alive": {},
  48. "proxy-authenticate": {},
  49. "proxy-authorization": {},
  50. "te": {},
  51. "trailer": {},
  52. "transfer-encoding": {},
  53. "upgrade": {},
  54. "cookie": {},
  55. // Additional headers that should not be forwarded by name-matching passthrough rules.
  56. "host": {},
  57. "content-length": {},
  58. "accept-encoding": {},
  59. // Do not passthrough credentials by wildcard/regex.
  60. "authorization": {},
  61. "x-api-key": {},
  62. "x-goog-api-key": {},
  63. // WebSocket handshake headers are generated by the client/dialer.
  64. "sec-websocket-key": {},
  65. "sec-websocket-version": {},
  66. "sec-websocket-extensions": {},
  67. }
  68. var headerPassthroughRegexCache sync.Map // map[string]*regexp.Regexp
  69. func getHeaderPassthroughRegex(pattern string) (*regexp.Regexp, error) {
  70. pattern = strings.TrimSpace(pattern)
  71. if pattern == "" {
  72. return nil, errors.New("empty regex pattern")
  73. }
  74. if v, ok := headerPassthroughRegexCache.Load(pattern); ok {
  75. if re, ok := v.(*regexp.Regexp); ok {
  76. return re, nil
  77. }
  78. headerPassthroughRegexCache.Delete(pattern)
  79. }
  80. compiled, err := regexp.Compile(pattern)
  81. if err != nil {
  82. return nil, err
  83. }
  84. actual, _ := headerPassthroughRegexCache.LoadOrStore(pattern, compiled)
  85. if re, ok := actual.(*regexp.Regexp); ok {
  86. return re, nil
  87. }
  88. return compiled, nil
  89. }
  90. func isHeaderPassthroughRuleKey(key string) bool {
  91. key = strings.TrimSpace(key)
  92. if key == "" {
  93. return false
  94. }
  95. if key == headerPassthroughAllKey {
  96. return true
  97. }
  98. lower := strings.ToLower(key)
  99. return strings.HasPrefix(lower, headerPassthroughRegexPrefix) || strings.HasPrefix(lower, headerPassthroughRegexPrefixV2)
  100. }
  101. func shouldSkipPassthroughHeader(name string) bool {
  102. name = strings.TrimSpace(name)
  103. if name == "" {
  104. return true
  105. }
  106. lower := strings.ToLower(name)
  107. if _, ok := passthroughSkipHeaderNamesLower[lower]; ok {
  108. return true
  109. }
  110. return false
  111. }
  112. func applyHeaderOverridePlaceholders(template string, c *gin.Context, apiKey string) (string, bool, error) {
  113. trimmed := strings.TrimSpace(template)
  114. if strings.HasPrefix(trimmed, clientHeaderPlaceholderPrefix) {
  115. afterPrefix := trimmed[len(clientHeaderPlaceholderPrefix):]
  116. end := strings.Index(afterPrefix, "}")
  117. if end < 0 || end != len(afterPrefix)-1 {
  118. return "", false, fmt.Errorf("client_header placeholder must be the full value: %q", template)
  119. }
  120. name := strings.TrimSpace(afterPrefix[:end])
  121. if name == "" {
  122. return "", false, fmt.Errorf("client_header placeholder name is empty: %q", template)
  123. }
  124. if c == nil || c.Request == nil {
  125. return "", false, fmt.Errorf("missing request context for client_header placeholder")
  126. }
  127. clientHeaderValue := c.Request.Header.Get(name)
  128. if strings.TrimSpace(clientHeaderValue) == "" {
  129. return "", false, nil
  130. }
  131. // Do not interpolate {api_key} inside client-supplied content.
  132. return clientHeaderValue, true, nil
  133. }
  134. if strings.Contains(template, "{api_key}") {
  135. template = strings.ReplaceAll(template, "{api_key}", apiKey)
  136. }
  137. if strings.TrimSpace(template) == "" {
  138. return "", false, nil
  139. }
  140. return template, true, nil
  141. }
  142. // processHeaderOverride applies channel header overrides, with placeholder substitution.
  143. // Supported placeholders:
  144. // - {api_key}: resolved to the channel API key
  145. // - {client_header:<name>}: resolved to the incoming request header value
  146. //
  147. // Header passthrough rules (keys only; values are ignored):
  148. // - "*": passthrough all incoming headers by name (excluding unsafe headers)
  149. // - "re:<regex>" / "regex:<regex>": passthrough headers whose names match the regex (Go regexp)
  150. //
  151. // Passthrough rules are applied first, then normal overrides are applied, so explicit overrides win.
  152. func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]string, error) {
  153. headerOverride := make(map[string]string)
  154. passAll := false
  155. var passthroughRegex []*regexp.Regexp
  156. if !info.IsChannelTest {
  157. for k := range info.HeadersOverride {
  158. key := strings.TrimSpace(k)
  159. if key == "" {
  160. continue
  161. }
  162. if key == headerPassthroughAllKey {
  163. passAll = true
  164. continue
  165. }
  166. lower := strings.ToLower(key)
  167. var pattern string
  168. switch {
  169. case strings.HasPrefix(lower, headerPassthroughRegexPrefix):
  170. pattern = strings.TrimSpace(key[len(headerPassthroughRegexPrefix):])
  171. case strings.HasPrefix(lower, headerPassthroughRegexPrefixV2):
  172. pattern = strings.TrimSpace(key[len(headerPassthroughRegexPrefixV2):])
  173. default:
  174. continue
  175. }
  176. if pattern == "" {
  177. return nil, types.NewError(fmt.Errorf("header passthrough regex pattern is empty: %q", k), types.ErrorCodeChannelHeaderOverrideInvalid)
  178. }
  179. compiled, err := getHeaderPassthroughRegex(pattern)
  180. if err != nil {
  181. return nil, types.NewError(err, types.ErrorCodeChannelHeaderOverrideInvalid)
  182. }
  183. passthroughRegex = append(passthroughRegex, compiled)
  184. }
  185. }
  186. if passAll || len(passthroughRegex) > 0 {
  187. if c == nil || c.Request == nil {
  188. return nil, types.NewError(fmt.Errorf("missing request context for header passthrough"), types.ErrorCodeChannelHeaderOverrideInvalid)
  189. }
  190. for name := range c.Request.Header {
  191. if shouldSkipPassthroughHeader(name) {
  192. continue
  193. }
  194. if !passAll {
  195. matched := false
  196. for _, re := range passthroughRegex {
  197. if re.MatchString(name) {
  198. matched = true
  199. break
  200. }
  201. }
  202. if !matched {
  203. continue
  204. }
  205. }
  206. value := strings.TrimSpace(c.Request.Header.Get(name))
  207. if value == "" {
  208. continue
  209. }
  210. headerOverride[name] = value
  211. }
  212. }
  213. for k, v := range info.HeadersOverride {
  214. if isHeaderPassthroughRuleKey(k) {
  215. continue
  216. }
  217. key := strings.TrimSpace(k)
  218. if key == "" {
  219. continue
  220. }
  221. str, ok := v.(string)
  222. if !ok {
  223. return nil, types.NewError(nil, types.ErrorCodeChannelHeaderOverrideInvalid)
  224. }
  225. if info.IsChannelTest && strings.HasPrefix(strings.TrimSpace(str), clientHeaderPlaceholderPrefix) {
  226. continue
  227. }
  228. value, include, err := applyHeaderOverridePlaceholders(str, c, info.ApiKey)
  229. if err != nil {
  230. return nil, types.NewError(err, types.ErrorCodeChannelHeaderOverrideInvalid)
  231. }
  232. if !include {
  233. continue
  234. }
  235. headerOverride[key] = value
  236. }
  237. return headerOverride, nil
  238. }
  239. func applyHeaderOverrideToRequest(req *http.Request, headerOverride map[string]string) {
  240. if req == nil {
  241. return
  242. }
  243. for key, value := range headerOverride {
  244. req.Header.Set(key, value)
  245. // set Host in req
  246. if strings.EqualFold(key, "Host") {
  247. req.Host = value
  248. }
  249. }
  250. }
  251. func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) {
  252. fullRequestURL, err := a.GetRequestURL(info)
  253. if err != nil {
  254. return nil, fmt.Errorf("get request url failed: %w", err)
  255. }
  256. if common2.DebugEnabled {
  257. common2.SysLog("fullRequestURL: " + fullRequestURL)
  258. }
  259. req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
  260. if err != nil {
  261. return nil, fmt.Errorf("new request failed: %w", err)
  262. }
  263. headers := req.Header
  264. err = a.SetupRequestHeader(c, &headers, info)
  265. if err != nil {
  266. return nil, fmt.Errorf("setup request header failed: %w", err)
  267. }
  268. headerOverride, err := processHeaderOverride(info, c)
  269. if err != nil {
  270. return nil, err
  271. }
  272. applyHeaderOverrideToRequest(req, headerOverride)
  273. dumpUpstreamRequest(req)
  274. resp, err := doRequest(c, req, info)
  275. if err != nil {
  276. return nil, fmt.Errorf("do request failed: %w", err)
  277. }
  278. return resp, nil
  279. }
  280. func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) {
  281. fullRequestURL, err := a.GetRequestURL(info)
  282. if err != nil {
  283. return nil, fmt.Errorf("get request url failed: %w", err)
  284. }
  285. if common2.DebugEnabled {
  286. common2.SysLog("fullRequestURL: " + fullRequestURL)
  287. }
  288. req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
  289. if err != nil {
  290. return nil, fmt.Errorf("new request failed: %w", err)
  291. }
  292. req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
  293. headers := req.Header
  294. err = a.SetupRequestHeader(c, &headers, info)
  295. if err != nil {
  296. return nil, fmt.Errorf("setup request header failed: %w", err)
  297. }
  298. headerOverride, err := processHeaderOverride(info, c)
  299. if err != nil {
  300. return nil, err
  301. }
  302. applyHeaderOverrideToRequest(req, headerOverride)
  303. dumpUpstreamRequest(req)
  304. resp, err := doRequest(c, req, info)
  305. if err != nil {
  306. return nil, fmt.Errorf("do request failed: %w", err)
  307. }
  308. return resp, nil
  309. }
  310. func dumpUpstreamRequest(req *http.Request) {
  311. if !common2.DebugEnabled || req == nil {
  312. return
  313. }
  314. bodyBytes, _ := io.ReadAll(req.Body)
  315. req.Body = io.NopCloser(bytes.NewReader(bodyBytes))
  316. common2.SysLog(fmt.Sprintf("[UpstreamRequest] URL: %s\nHeaders: %v\nBody: %s", req.URL.String(), req.Header, string(bodyBytes)))
  317. }
  318. func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*websocket.Conn, error) {
  319. fullRequestURL, err := a.GetRequestURL(info)
  320. if err != nil {
  321. return nil, fmt.Errorf("get request url failed: %w", err)
  322. }
  323. targetHeader := http.Header{}
  324. err = a.SetupRequestHeader(c, &targetHeader, info)
  325. headerOverride, err := processHeaderOverride(info, c)
  326. if err != nil {
  327. return nil, err
  328. }
  329. for key, value := range headerOverride {
  330. targetHeader.Set(key, value)
  331. }
  332. targetHeader.Set("Content-Type", c.Request.Header.Get("Content-Type"))
  333. targetConn, resp, err := websocket.DefaultDialer.Dial(fullRequestURL, targetHeader)
  334. if err != nil {
  335. return nil, fmt.Errorf("dial failed to %s: %w", fullRequestURL, err)
  336. }
  337. if resp != nil {
  338. common.CaptureUpstreamIDFromHTTPResponse(info, resp)
  339. common.SetRelayUpstreamID(c, info.UpstreamID)
  340. }
  341. // send request body
  342. //all, err := io.ReadAll(requestBody)
  343. //err = service.WssString(c, targetConn, string(all))
  344. return targetConn, nil
  345. }
  346. func startPingKeepAlive(c *gin.Context, pingInterval time.Duration) context.CancelFunc {
  347. pingerCtx, stopPinger := context.WithCancel(context.Background())
  348. gopool.Go(func() {
  349. defer func() {
  350. // 增加panic恢复处理
  351. if r := recover(); r != nil {
  352. if common2.DebugEnabled {
  353. println("SSE ping goroutine panic recovered:", fmt.Sprintf("%v", r))
  354. }
  355. }
  356. if common2.DebugEnabled {
  357. println("SSE ping goroutine stopped.")
  358. }
  359. }()
  360. if pingInterval <= 0 {
  361. pingInterval = helper.DefaultPingInterval
  362. }
  363. ticker := time.NewTicker(pingInterval)
  364. // 确保在任何情况下都清理ticker
  365. defer func() {
  366. ticker.Stop()
  367. if common2.DebugEnabled {
  368. println("SSE ping ticker stopped")
  369. }
  370. }()
  371. var pingMutex sync.Mutex
  372. if common2.DebugEnabled {
  373. println("SSE ping goroutine started")
  374. }
  375. // 增加超时控制,防止goroutine长时间运行
  376. maxPingDuration := 120 * time.Minute // 最大ping持续时间
  377. pingTimeout := time.NewTimer(maxPingDuration)
  378. defer pingTimeout.Stop()
  379. for {
  380. select {
  381. // 发送 ping 数据
  382. case <-ticker.C:
  383. if err := sendPingData(c, &pingMutex); err != nil {
  384. if common2.DebugEnabled {
  385. println("SSE ping error, stopping goroutine:", err.Error())
  386. }
  387. return
  388. }
  389. // 收到退出信号
  390. case <-pingerCtx.Done():
  391. return
  392. // request 结束
  393. case <-c.Request.Context().Done():
  394. return
  395. // 超时保护,防止goroutine无限运行
  396. case <-pingTimeout.C:
  397. if common2.DebugEnabled {
  398. println("SSE ping goroutine timeout, stopping")
  399. }
  400. return
  401. }
  402. }
  403. })
  404. return stopPinger
  405. }
  406. func sendPingData(c *gin.Context, mutex *sync.Mutex) error {
  407. // 增加超时控制,防止锁死等待
  408. done := make(chan error, 1)
  409. go func() {
  410. mutex.Lock()
  411. defer mutex.Unlock()
  412. err := helper.PingData(c)
  413. if err != nil {
  414. logger.LogError(c, "SSE ping error: "+err.Error())
  415. done <- err
  416. return
  417. }
  418. if common2.DebugEnabled {
  419. println("SSE ping data sent.")
  420. }
  421. done <- nil
  422. }()
  423. // 设置发送ping数据的超时时间
  424. select {
  425. case err := <-done:
  426. return err
  427. case <-time.After(10 * time.Second):
  428. return errors.New("SSE ping data send timeout")
  429. case <-c.Request.Context().Done():
  430. return errors.New("request context cancelled during ping")
  431. }
  432. }
  433. func DoRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
  434. return doRequest(c, req, info)
  435. }
  436. func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
  437. var client *http.Client
  438. var err error
  439. if info.ChannelSetting.Proxy != "" {
  440. client, err = service.NewProxyHttpClient(info.ChannelSetting.Proxy)
  441. if err != nil {
  442. return nil, fmt.Errorf("new proxy http client failed: %w", err)
  443. }
  444. } else {
  445. client = service.GetHttpClient()
  446. }
  447. var stopPinger context.CancelFunc
  448. if info.IsStream {
  449. helper.SetEventStreamHeaders(c)
  450. // 处理流式请求的 ping 保活
  451. generalSettings := operation_setting.GetGeneralSetting()
  452. if generalSettings.PingIntervalEnabled && !info.DisablePing {
  453. pingInterval := time.Duration(generalSettings.PingIntervalSeconds) * time.Second
  454. stopPinger = startPingKeepAlive(c, pingInterval)
  455. // 使用defer确保在任何情况下都能停止ping goroutine
  456. defer func() {
  457. if stopPinger != nil {
  458. stopPinger()
  459. if common2.DebugEnabled {
  460. println("SSE ping goroutine stopped by defer")
  461. }
  462. }
  463. }()
  464. }
  465. }
  466. resp, err := client.Do(req)
  467. if err != nil {
  468. logger.LogError(c, "do request failed: "+err.Error())
  469. return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithHideErrMsg("upstream error: do request failed"))
  470. }
  471. if resp == nil {
  472. return nil, errors.New("resp is nil")
  473. }
  474. common.CaptureUpstreamIDFromHTTPResponse(info, resp)
  475. common.SetRelayUpstreamID(c, info.UpstreamID)
  476. _ = req.Body.Close()
  477. _ = c.Request.Body.Close()
  478. return resp, nil
  479. }
  480. func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) {
  481. fullRequestURL, err := a.BuildRequestURL(info)
  482. if err != nil {
  483. return nil, err
  484. }
  485. req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
  486. if err != nil {
  487. return nil, fmt.Errorf("new request failed: %w", err)
  488. }
  489. req.GetBody = func() (io.ReadCloser, error) {
  490. return io.NopCloser(requestBody), nil
  491. }
  492. err = a.BuildRequestHeader(c, req, info)
  493. if err != nil {
  494. return nil, fmt.Errorf("setup request header failed: %w", err)
  495. }
  496. resp, err := doRequest(c, req, info)
  497. if err != nil {
  498. return nil, fmt.Errorf("do request failed: %w", err)
  499. }
  500. return resp, nil
  501. }