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

package channel

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"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)
	}
	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, "SSE ping error: "+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, fmt.Sprintf(
			"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++
	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)
		if info.DurableAttemptID != 0 {
			guardUsage = service.NewGuardUsageCollector(c, info)
			guard.OnEvent = guardUsage.Observe
			guard.OnClose = func() { guardUsage.CheckpointFinal(guard.Reason) }
		}
		req = req.WithContext(upstream)
	}
	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, "do request failed: "+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")
	}
	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 common2.DebugEnabled {
		policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting)
		logger.LogDebug(c, fmt.Sprintf(
			"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)
}