Files
new-api-custom/relay/channel/api_request.go
T

692 lines
22 KiB
Go

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:<name>}: 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>" / "regex:<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)
}