Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

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