feat(backend): 提交后端审计修复与配套测试改动

This commit is contained in:
yangjianbo
2026-02-14 11:23:10 +08:00
parent 862199143e
commit d04b47b3ca
22 changed files with 653 additions and 55 deletions

View File

@@ -3332,7 +3332,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
// 不需要重试(成功或不可重试的错误),跳出循环
// DEBUG: 输出响应 headers用于检测 rate limit 信息)
if account.Platform == PlatformGemini && resp.StatusCode < 400 {
if account.Platform == PlatformGemini && resp.StatusCode < 400 && s.cfg != nil && s.cfg.Gateway.GeminiDebugResponseHeaders {
logger.LegacyPrintf("service.gateway", "[DEBUG] Gemini API Response Headers for account %d:", account.ID)
for k, v := range resp.Header {
logger.LegacyPrintf("service.gateway", "[DEBUG] %s: %v", k, v)
@@ -4467,8 +4467,19 @@ func (s *GatewayService) handleNonStreamingResponse(ctx context.Context, resp *h
// 更新5h窗口状态
s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header)
body, err := io.ReadAll(resp.Body)
maxBytes := resolveUpstreamResponseReadLimit(s.cfg)
body, err := readUpstreamResponseBodyLimited(resp.Body, maxBytes)
if err != nil {
if errors.Is(err, ErrUpstreamResponseBodyTooLarge) {
setOpsUpstreamError(c, http.StatusBadGateway, "upstream response too large", "")
c.JSON(http.StatusBadGateway, gin.H{
"type": "error",
"error": gin.H{
"type": "upstream_error",
"message": "Upstream response too large",
},
})
}
return nil, err
}
@@ -4990,9 +5001,15 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
}
// 读取响应体
respBody, err := io.ReadAll(resp.Body)
maxReadBytes := resolveUpstreamResponseReadLimit(s.cfg)
respBody, err := readUpstreamResponseBodyLimited(resp.Body, maxReadBytes)
_ = resp.Body.Close()
if err != nil {
if errors.Is(err, ErrUpstreamResponseBodyTooLarge) {
setOpsUpstreamError(c, http.StatusBadGateway, "upstream response too large", "")
s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream response too large")
return err
}
s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response")
return err
}
@@ -5007,9 +5024,14 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
retryResp, retryErr := s.httpUpstream.DoWithTLS(retryReq, proxyURL, account.ID, account.Concurrency, account.IsTLSFingerprintEnabled())
if retryErr == nil {
resp = retryResp
respBody, err = io.ReadAll(resp.Body)
respBody, err = readUpstreamResponseBodyLimited(resp.Body, maxReadBytes)
_ = resp.Body.Close()
if err != nil {
if errors.Is(err, ErrUpstreamResponseBodyTooLarge) {
setOpsUpstreamError(c, http.StatusBadGateway, "upstream response too large", "")
s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream response too large")
return err
}
s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response")
return err
}