fix: fix getAndValidateTextRequest failed: unexpected end of JSON input (close #1043)

This commit is contained in:
JustSong 2024-02-26 22:52:16 +08:00
parent 6b27d6659a
commit b747cdbc6f
2 changed files with 26 additions and 6 deletions

View File

@ -8,12 +8,24 @@ import (
"strings" "strings"
) )
func UnmarshalBodyReusable(c *gin.Context, v any) error { const KeyRequestBody = "key_request_body"
func GetRequestBody(c *gin.Context) ([]byte, error) {
requestBody, _ := c.Get(KeyRequestBody)
if requestBody != nil {
return requestBody.([]byte), nil
}
requestBody, err := io.ReadAll(c.Request.Body) requestBody, err := io.ReadAll(c.Request.Body)
if err != nil { if err != nil {
return err return nil, err
} }
err = c.Request.Body.Close() _ = c.Request.Body.Close()
c.Set(KeyRequestBody, requestBody)
return requestBody.([]byte), nil
}
func UnmarshalBodyReusable(c *gin.Context, v any) error {
requestBody, err := GetRequestBody(c)
if err != nil { if err != nil {
return err return err
} }

View File

@ -1,9 +1,11 @@
package controller package controller
import ( import (
"bytes"
"context" "context"
"fmt" "fmt"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
@ -13,6 +15,7 @@ import (
"github.com/songquanpeng/one-api/relay/controller" "github.com/songquanpeng/one-api/relay/controller"
"github.com/songquanpeng/one-api/relay/model" "github.com/songquanpeng/one-api/relay/model"
"github.com/songquanpeng/one-api/relay/util" "github.com/songquanpeng/one-api/relay/util"
"io"
"net/http" "net/http"
) )
@ -50,8 +53,8 @@ func Relay(c *gin.Context) {
go processChannelRelayError(ctx, channelId, channelName, bizErr) go processChannelRelayError(ctx, channelId, channelName, bizErr)
requestId := c.GetString(logger.RequestIdKey) requestId := c.GetString(logger.RequestIdKey)
retryTimes := config.RetryTimes retryTimes := config.RetryTimes
if !shouldRetry(bizErr.StatusCode) { if !shouldRetry(c, bizErr.StatusCode) {
logger.Errorf(ctx, "relay error happen, but status code is %d, won't retry in this case", bizErr.StatusCode) logger.Errorf(ctx, "relay error happen, status code is %d, won't retry in this case", bizErr.StatusCode)
retryTimes = 0 retryTimes = 0
} }
for i := retryTimes; i > 0; i-- { for i := retryTimes; i > 0; i-- {
@ -65,6 +68,8 @@ func Relay(c *gin.Context) {
continue continue
} }
middleware.SetupContextForSelectedChannel(c, channel, originalModel) middleware.SetupContextForSelectedChannel(c, channel, originalModel)
requestBody, err := common.GetRequestBody(c)
c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
bizErr = relay(c, relayMode) bizErr = relay(c, relayMode)
if bizErr == nil { if bizErr == nil {
return return
@ -85,7 +90,10 @@ func Relay(c *gin.Context) {
} }
} }
func shouldRetry(statusCode int) bool { func shouldRetry(c *gin.Context, statusCode int) bool {
if _, ok := c.Get("specific_channel_id"); ok {
return false
}
if statusCode == http.StatusTooManyRequests { if statusCode == http.StatusTooManyRequests {
return true return true
} }