package channel import ( "context" "errors" "fmt" "io" "net/http" "net/http/httptrace" "regexp" "strings" "sync" "time" common2 "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relay/helper" "github.com/QuantumNous/new-api/relaykit/types" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/bytedance/gopkg/util/gopool" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" ) // ApplyUpstreamBodyMetadata restores metadata that net/http cannot infer from // a ReplayableBody. Callers must pass the original body because NewRequest // hides its dynamic type behind req.Body's io.ReadCloser wrapper. func ApplyUpstreamBodyMetadata(req *http.Request, body io.Reader) { replayable, ok := body.(common2.ReplayableBody) if !ok { return } // BodyStorage structurally satisfies ReplayableBody, but it also exposes // io.Closer. If a caller passes the storage directly instead of using // NewReplayableBodyReader, hide Close before the transport takes ownership // of req.Body so the shared replay source remains available to GetBody. if _, rawStorage := body.(common2.BodyStorage); rawStorage { req.Body = io.NopCloser(body) } req.ContentLength = replayable.Size() if req.GetBody == nil { req.GetBody = replayable.NewReader } } func SetupApiRequestHeader(info *common.RelayInfo, c *gin.Context, req *http.Header) { if info.RelayMode == constant.RelayModeAudioTranscription || info.RelayMode == constant.RelayModeAudioTranslation { // multipart/form-data } else if info.RelayMode == constant.RelayModeRealtime { // websocket } else { req.Set("Content-Type", c.Request.Header.Get("Content-Type")) req.Set("Accept", c.Request.Header.Get("Accept")) if info.IsStream && c.Request.Header.Get("Accept") == "" { req.Set("Accept", "text/event-stream") } } } const clientHeaderPlaceholderPrefix = "{client_header:" const ( headerPassthroughAllKey = "*" headerPassthroughRegexPrefix = "re:" headerPassthroughRegexPrefixV2 = "regex:" ) var passthroughSkipHeaderNamesLower = map[string]struct{}{ // RFC 7230 hop-by-hop headers. "connection": {}, "keep-alive": {}, "proxy-authenticate": {}, "proxy-authorization": {}, "te": {}, "trailer": {}, "transfer-encoding": {}, "upgrade": {}, "cookie": {}, // Additional headers that should not be forwarded by name-matching passthrough rules. "host": {}, "content-length": {}, "accept-encoding": {}, // Do not passthrough credentials by wildcard/regex. "authorization": {}, "x-api-key": {}, "x-goog-api-key": {}, // WebSocket handshake headers are generated by the client/dialer. "sec-websocket-key": {}, "sec-websocket-version": {}, "sec-websocket-extensions": {}, } var headerPassthroughRegexCache sync.Map // map[string]*regexp.Regexp func getHeaderPassthroughRegex(pattern string) (*regexp.Regexp, error) { pattern = strings.TrimSpace(pattern) if pattern == "" { return nil, errors.New("empty regex pattern") } if v, ok := headerPassthroughRegexCache.Load(pattern); ok { if re, ok := v.(*regexp.Regexp); ok { return re, nil } headerPassthroughRegexCache.Delete(pattern) } compiled, err := regexp.Compile(pattern) if err != nil { return nil, err } actual, _ := headerPassthroughRegexCache.LoadOrStore(pattern, compiled) if re, ok := actual.(*regexp.Regexp); ok { return re, nil } return compiled, nil } func IsHeaderPassthroughRuleKey(key string) bool { return isHeaderPassthroughRuleKey(key) } func isHeaderPassthroughRuleKey(key string) bool { key = strings.TrimSpace(key) if key == "" { return false } if key == headerPassthroughAllKey { return true } lower := strings.ToLower(key) return strings.HasPrefix(lower, headerPassthroughRegexPrefix) || strings.HasPrefix(lower, headerPassthroughRegexPrefixV2) } func shouldSkipPassthroughHeader(name string) bool { name = strings.TrimSpace(name) if name == "" { return true } lower := strings.ToLower(name) if _, ok := passthroughSkipHeaderNamesLower[lower]; ok { return true } return false } func applyHeaderOverridePlaceholders(template string, c *gin.Context, apiKey string) (string, bool, error) { trimmed := strings.TrimSpace(template) if strings.HasPrefix(trimmed, clientHeaderPlaceholderPrefix) { afterPrefix := trimmed[len(clientHeaderPlaceholderPrefix):] end := strings.Index(afterPrefix, "}") if end < 0 || end != len(afterPrefix)-1 { return "", false, fmt.Errorf("client_header placeholder must be the full value: %q", template) } name := strings.TrimSpace(afterPrefix[:end]) if name == "" { return "", false, fmt.Errorf("client_header placeholder name is empty: %q", template) } if c == nil || c.Request == nil { return "", false, fmt.Errorf("missing request context for client_header placeholder") } clientHeaderValue := c.Request.Header.Get(name) if strings.TrimSpace(clientHeaderValue) == "" { return "", false, nil } // Do not interpolate {api_key} inside client-supplied content. return clientHeaderValue, true, nil } if strings.Contains(template, "{api_key}") { template = strings.ReplaceAll(template, "{api_key}", apiKey) } if strings.TrimSpace(template) == "" { return "", false, nil } return template, true, nil } // processHeaderOverride applies channel header overrides, with placeholder substitution. // Supported placeholders: // - {api_key}: resolved to the channel API key // - {client_header:}: resolved to the incoming request header value // // Header passthrough rules (keys only; values are ignored): // - "*": passthrough all incoming headers by name (excluding unsafe headers) // - "re:" / "regex:": passthrough headers whose names match the regex (Go regexp) // // Passthrough rules are applied first, then normal overrides are applied, so explicit overrides win. func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]string, error) { headerOverride := make(map[string]string) if info == nil { return headerOverride, nil } headerOverrideSource := common.GetEffectiveHeaderOverride(info) passAll := false var passthroughRegex []*regexp.Regexp if !info.IsChannelTest { for k := range headerOverrideSource { key := strings.TrimSpace(strings.ToLower(k)) if key == "" { continue } if key == headerPassthroughAllKey { passAll = true continue } var pattern string switch { case strings.HasPrefix(key, headerPassthroughRegexPrefix): pattern = strings.TrimSpace(key[len(headerPassthroughRegexPrefix):]) case strings.HasPrefix(key, headerPassthroughRegexPrefixV2): pattern = strings.TrimSpace(key[len(headerPassthroughRegexPrefixV2):]) default: continue } if pattern == "" { return nil, types.NewError(fmt.Errorf("header passthrough regex pattern is empty: %q", k), types.ErrorCodeChannelHeaderOverrideInvalid) } compiled, err := getHeaderPassthroughRegex(pattern) if err != nil { return nil, types.NewError(err, types.ErrorCodeChannelHeaderOverrideInvalid) } passthroughRegex = append(passthroughRegex, compiled) } } if passAll || len(passthroughRegex) > 0 { if c == nil || c.Request == nil { return nil, types.NewError(fmt.Errorf("missing request context for header passthrough"), types.ErrorCodeChannelHeaderOverrideInvalid) } for name := range c.Request.Header { if shouldSkipPassthroughHeader(name) { continue } if !passAll { matched := false for _, re := range passthroughRegex { if re.MatchString(name) { matched = true break } } if !matched { continue } } value := strings.TrimSpace(c.Request.Header.Get(name)) if value == "" { continue } headerOverride[strings.ToLower(strings.TrimSpace(name))] = value } } for k, v := range headerOverrideSource { if isHeaderPassthroughRuleKey(k) { continue } key := strings.TrimSpace(strings.ToLower(k)) if key == "" { continue } str, ok := v.(string) if !ok { return nil, types.NewError(nil, types.ErrorCodeChannelHeaderOverrideInvalid) } if info.IsChannelTest && strings.HasPrefix(strings.TrimSpace(str), clientHeaderPlaceholderPrefix) { continue } value, include, err := applyHeaderOverridePlaceholders(str, c, info.ApiKey) if err != nil { return nil, types.NewError(err, types.ErrorCodeChannelHeaderOverrideInvalid) } if !include { continue } headerOverride[key] = value } return headerOverride, nil } func ResolveHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]string, error) { return processHeaderOverride(info, c) } func applyHeaderOverrideToRequest(req *http.Request, headerOverride map[string]string) { if req == nil { return } for key, value := range headerOverride { req.Header.Set(key, value) // set Host in req if strings.EqualFold(key, "Host") { req.Host = value } } } func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) { fullRequestURL, err := a.GetRequestURL(info) if err != nil { return nil, fmt.Errorf("get request url failed: %w", err) } logger.LogDebug(c, "fullRequestURL: %s", common.SanitizeURLForLog(fullRequestURL)) req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody) if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } ApplyUpstreamBodyMetadata(req, requestBody) headers := req.Header err = a.SetupRequestHeader(c, &headers, info) if err != nil { return nil, fmt.Errorf("setup request header failed: %w", err) } // 在 SetupRequestHeader 之后应用 Header Override,确保用户设置优先级最高 // 这样可以覆盖默认的 Authorization header 设置 headerOverride, err := processHeaderOverride(info, c) if err != nil { return nil, err } applyHeaderOverrideToRequest(req, headerOverride) resp, err := doRequest(c, req, info) if err != nil { return nil, fmt.Errorf("do request failed: %w", err) } return resp, nil } func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) { fullRequestURL, err := a.GetRequestURL(info) if err != nil { return nil, fmt.Errorf("get request url failed: %w", err) } logger.LogDebug(c, "fullRequestURL: %s", common.SanitizeURLForLog(fullRequestURL)) req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody) if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } ApplyUpstreamBodyMetadata(req, requestBody) // set form data req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) headers := req.Header err = a.SetupRequestHeader(c, &headers, info) if err != nil { return nil, fmt.Errorf("setup request header failed: %w", err) } // 在 SetupRequestHeader 之后应用 Header Override,确保用户设置优先级最高 // 这样可以覆盖默认的 Authorization header 设置 headerOverride, err := processHeaderOverride(info, c) if err != nil { return nil, err } applyHeaderOverrideToRequest(req, headerOverride) resp, err := doRequest(c, req, info) if err != nil { return nil, fmt.Errorf("do request failed: %w", err) } return resp, nil } // DoWssRequest dials the adaptor's upstream over WebSocket for the realtime // and Responses WebSocket relays. It honors the channel proxy, is bound to the // request context, and reports a rejected handshake as a *types.NewAPIError // carrying the upstream status code (types.NewError preserves it). func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*websocket.Conn, error) { fullRequestURL, err := a.GetRequestURL(info) if err != nil { return nil, fmt.Errorf("get request url failed: %w", err) } fullRequestURL = toWebSocketURL(fullRequestURL) targetHeader := http.Header{} err = a.SetupRequestHeader(c, &targetHeader, info) if err != nil { return nil, fmt.Errorf("setup request header failed: %w", err) } // 在 SetupRequestHeader 之后应用 Header Override,确保用户设置优先级最高 // 这样可以覆盖默认的 Authorization header 设置 headerOverride, err := processHeaderOverride(info, c) if err != nil { return nil, err } for key, value := range headerOverride { targetHeader.Set(key, value) } targetHeader.Set("Content-Type", c.Request.Header.Get("Content-Type")) dialer := *websocket.DefaultDialer if info.ChannelSetting.Proxy != "" { proxyURL, _, proxyErr := common2.ParseProxyURLRuntime(info.ChannelSetting.Proxy) if proxyErr != nil { return nil, proxyErr } dialer.Proxy = http.ProxyURL(proxyURL) } if info.ChannelSetting.TLSInsecureSkipVerify { dialer.TLSClientConfig = common2.InsecureTLSConfig } targetConn, resp, err := dialer.DialContext(c.Request.Context(), fullRequestURL, targetHeader) if err != nil { statusCode := http.StatusInternalServerError if resp != nil { statusCode = resp.StatusCode if resp.Body != nil { _ = resp.Body.Close() } } return nil, types.NewErrorWithStatusCode(fmt.Errorf("dial failed to %s: %w", common.SanitizeURLForLog(fullRequestURL), err), types.ErrorCodeDoRequestFailed, statusCode) } return targetConn, nil } // toWebSocketURL maps an http(s) endpoint to ws(s). Realtime adaptors already // return ws(s) URLs, which pass through unchanged. func toWebSocketURL(raw string) string { switch { case strings.HasPrefix(raw, "https://"): return "wss://" + strings.TrimPrefix(raw, "https://") case strings.HasPrefix(raw, "http://"): return "ws://" + strings.TrimPrefix(raw, "http://") default: return raw } } func startPingKeepAlive(c *gin.Context, pingInterval time.Duration) (context.CancelFunc, <-chan struct{}) { pingerCtx, stopPinger := context.WithCancel(context.Background()) done := make(chan struct{}) gopool.Go(func() { defer close(done) defer func() { // 增加panic恢复处理 if r := recover(); r != nil { logger.LogDebug(c, "SSE ping goroutine panic recovered: %v", r) } logger.LogDebug(c, "SSE ping goroutine stopped") }() if pingInterval <= 0 { pingInterval = helper.DefaultPingInterval } ticker := time.NewTicker(pingInterval) // 确保在任何情况下都清理ticker defer func() { ticker.Stop() logger.LogDebug(c, "SSE ping ticker stopped") }() var pingMutex sync.Mutex logger.LogDebug(c, "SSE ping goroutine started") // 增加超时控制,防止goroutine长时间运行 maxPingDuration := 120 * time.Minute // 最大ping持续时间 pingTimeout := time.NewTimer(maxPingDuration) defer pingTimeout.Stop() for { select { // 发送 ping 数据 case <-ticker.C: if err := sendPingData(c, &pingMutex); err != nil { logger.LogDebug(c, "SSE ping error, stopping goroutine: %s", err.Error()) return } // 收到退出信号 case <-pingerCtx.Done(): return // request 结束 case <-c.Request.Context().Done(): return // 超时保护,防止goroutine无限运行 case <-pingTimeout.C: logger.LogDebug(c, "SSE ping goroutine timeout, stopping") return } } }) return stopPinger, done } func sendPingData(c *gin.Context, mutex *sync.Mutex) error { mutex.Lock() defer mutex.Unlock() // Bound the write so a slow client cannot block this goroutine forever; // doRequest's defer waits for the pinger to exit before returning. helper.ExtendWriteDeadline(c) err := helper.PingData(c) if err != nil { logger.LogError(c, common2.LogText("SSE ping error: %s", err.Error())) return err } logger.LogDebug(c, "SSE ping data sent") return nil } func DoRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) { return doRequest(c, req, info) } // keepUpstreamRedirectResponse stops net/http from following redirects while // returning the upstream 3xx response to the relay without an extra error. func keepUpstreamRedirectResponse(_ *http.Request, _ []*http.Request) error { return http.ErrUseLastResponse } func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) { guardEnabled := info.ResponseGuard != nil && info.AstraAutoSwitch if guardEnabled && info.UpstreamRequests >= common.ResponseGuardMaxAttempts { return nil, types.NewErrorWithStatusCode(errors.New("response model attempt budget exhausted"), common.ResponseGuardErrorCode, http.StatusBadGateway, types.ErrOptionWithSkipRetry()) } if guardEnabled && c.Request.Context().Err() != nil { return nil, types.NewErrorWithStatusCode(c.Request.Context().Err(), common.ResponseGuardErrorCode, http.StatusBadGateway, types.ErrOptionWithSkipRetry()) } client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting) if err != nil { return nil, fmt.Errorf("new proxy http client failed: %w", err) } // Clients are cached and shared across channels, so override redirect // behavior on a shallow copy instead of mutating the cached client. This // still reuses its transport and connection pools, including HTTP/2's // transparent stream retries. relayClient := *client relayClient.CheckRedirect = keepUpstreamRedirectResponse if common2.DebugEnabled && req != nil && req.URL != nil { policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting) logger.LogDebug(c, "http transport select: host=%s protocol=%s shards=%d policy=%s", req.URL.Host, policy.Protocol, policy.Shards, policy.String()) } var stopPinger context.CancelFunc var pingerDone <-chan struct{} if info.IsStream && !guardEnabled { helper.SetEventStreamHeaders(c) // 处理流式请求的 ping 保活 generalSettings := operation_setting.GetGeneralSetting() if generalSettings.PingIntervalEnabled && !info.DisablePing { pingInterval := time.Duration(generalSettings.PingIntervalSeconds) * time.Second stopPinger, pingerDone = startPingKeepAlive(c, pingInterval) // 使用defer确保在任何情况下都能停止ping goroutine defer func() { if stopPinger != nil { stopPinger() <-pingerDone logger.LogDebug(c, "SSE ping goroutine stopped by defer") } }() } } info.UpstreamRequests++ info.UpstreamTiming = nil if info.UsingGroup == common2.GetEnvOrDefaultString("UPSTREAM_TIMING_GROUP", "") && c.Request.URL.Path == "/v1/responses" { start := c.GetTime("request_received_at") if start.IsZero() { start = info.StartTime } info.UpstreamTiming = common.NewUpstreamTiming(start, info.UpstreamRequests, info.ChannelId, req.ContentLength) } c.Set(common.UpstreamTimingContextKey, info.UpstreamTiming) var guard *common.ResponseGuardAttempt var guardUsage *service.GuardUsageCollector if guardEnabled { var upstream context.Context upstream, guard = info.ResponseGuard.Begin(c.Request.Context(), info.UpstreamRequests, info.ChannelId) guard.Timing = info.UpstreamTiming if info.DurableAttemptID != 0 { guardUsage = service.NewGuardUsageCollector(c, info) guard.OnEvent = guardUsage.Observe guard.OnClose = func() { guardUsage.CheckpointFinal(guard.Reason) } } req = req.WithContext(upstream) } if info.UpstreamTiming != nil { req = req.WithContext(httptrace.WithClientTrace(req.Context(), info.UpstreamTiming.ClientTrace())) } resp, err := relayClient.Do(req) if err != nil { if guard != nil { guard.Close() guard.Finish(false) if c.Request.Context().Err() != nil { guard.Reject("client_cancelled") return nil, guard.APIError() } if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { guard.Reject("confirmation_timeout") return nil, guard.APIError() } guard.Reason = "transport_error" } logger.LogError(c, common2.LogText("do request failed: %s", err.Error())) return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithHideErrMsg("upstream error: do request failed")) } if resp == nil { if guard != nil { guard.Close() } return nil, errors.New("resp is nil") } info.UpstreamTiming.Response(resp.Proto, resp.StatusCode) if guard != nil { if err := guard.Confirm(resp, c.Request.Context()); err != nil { if guardUsage != nil && guardUsage.FinishRejected(info, guard, resp) { guard.Finish(false) return nil, guard.APIError() } guard.Close() _ = resp.Body.Close() guard.Finish(false) return nil, guard.APIError() } if guard.Reason == "accepted" { info.UpstreamTiming.Mark("model_confirmed") } } if common2.DebugEnabled { policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting) logger.LogDebug(c, "http transport negotiated: host=%s protocol=%s shards=%d policy=%s negotiated=%s", req.URL.Host, policy.Protocol, policy.Shards, policy.String(), resp.Proto) } if upID := resp.Header.Get(common2.RequestIdKey); upID != "" { c.Set(common2.UpstreamRequestIdKey, upID) } _ = req.Body.Close() _ = c.Request.Body.Close() return resp, nil } func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) { fullRequestURL, err := a.BuildRequestURL(info) if err != nil { return nil, err } req, err := newTaskAPIRequest(c, fullRequestURL, requestBody) if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } ApplyUpstreamBodyMetadata(req, requestBody) // Do NOT wrap requestBody in a GetBody closure here: returning the same // (already consumed) reader would make any transport-level retry silently // replay an empty body. http.NewRequest already derives a correct, // snapshot-based GetBody for *bytes.Reader/Buffer/strings.Reader bodies // (which most task adaptors pass in); ApplyUpstreamBodyMetadata wires the // same contract for bodies that explicitly implement ReplayableBody. // Otherwise GetBody stays nil so the transport fails the retry instead of // sending a corrupted request. err = a.BuildRequestHeader(c, req, info) if err != nil { return nil, fmt.Errorf("setup request header failed: %w", err) } resp, err := doRequest(c, req, info) if err != nil { return nil, fmt.Errorf("do request failed: %w", err) } return resp, nil } func newTaskAPIRequest(c *gin.Context, fullRequestURL string, requestBody io.Reader) (*http.Request, error) { if c == nil || c.Request == nil { return nil, errors.New("task client request is missing") } return http.NewRequestWithContext(c.Request.Context(), c.Request.Method, fullRequestURL, requestBody) }