From 5839d84f6ff288f3ebcbd5cd68fff830bacdaf76 Mon Sep 17 00:00:00 2001 From: shubo <814183583@qq.com> Date: Tue, 30 Jun 2026 16:52:39 +0800 Subject: [PATCH] fix: close response bodies on non-ok responses --- controller/channel-billing.go | 13 ++++---- controller/channel.go | 2 +- service/task_polling.go | 2 +- service/task_polling_test.go | 62 +++++++++++++++++++++++++++++++---- 4 files changed, 65 insertions(+), 14 deletions(-) diff --git a/controller/channel-billing.go b/controller/channel-billing.go index 751ee3600ac9..10d4a2e36656 100644 --- a/controller/channel-billing.go +++ b/controller/channel-billing.go @@ -136,7 +136,7 @@ func GetClaudeAuthHeader(token string) http.Header { return h } -func GetResponseBody(method, url string, channel *model.Channel, headers http.Header) ([]byte, error) { +func GetResponseBody(method, url string, channel *model.Channel, headers http.Header) (body []byte, err error) { req, err := http.NewRequest(method, url, nil) if err != nil { return nil, err @@ -152,14 +152,15 @@ func GetResponseBody(method, url string, channel *model.Channel, headers http.He if err != nil { return nil, err } + defer func() { + if closeErr := res.Body.Close(); err == nil && closeErr != nil { + err = closeErr + } + }() if res.StatusCode != http.StatusOK { return nil, fmt.Errorf("status code: %d", res.StatusCode) } - body, err := io.ReadAll(res.Body) - if err != nil { - return nil, err - } - err = res.Body.Close() + body, err = io.ReadAll(res.Body) if err != nil { return nil, err } diff --git a/controller/channel.go b/controller/channel.go index a2b5687ab319..1fe39af19001 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -1241,6 +1241,7 @@ func FetchModels(c *gin.Context) { }) return } + defer response.Body.Close() //check status code if response.StatusCode != http.StatusOK { c.JSON(http.StatusInternalServerError, gin.H{ @@ -1249,7 +1250,6 @@ func FetchModels(c *gin.Context) { }) return } - defer response.Body.Close() var result struct { Data []struct { diff --git a/service/task_polling.go b/service/task_polling.go index 179fa7342082..36b3f5c196d3 100644 --- a/service/task_polling.go +++ b/service/task_polling.go @@ -243,11 +243,11 @@ func updateSunoTasks(ctx context.Context, channelId int, taskIds []string, taskM common.SysLog(fmt.Sprintf("Get Task Do req error: %v", err)) return err } + defer resp.Body.Close() if resp.StatusCode != http.StatusOK { logger.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode)) return fmt.Errorf("Get Task status code: %d", resp.StatusCode) } - defer resp.Body.Close() responseBody, err := io.ReadAll(resp.Body) if err != nil { common.SysLog(fmt.Sprintf("Get Suno Task parse body error: %v", err)) diff --git a/service/task_polling_test.go b/service/task_polling_test.go index 3164d0a9e29c..522ac09242b0 100644 --- a/service/task_polling_test.go +++ b/service/task_polling_test.go @@ -27,6 +27,8 @@ type taskPollingFetchAdaptor struct { blockStarted chan struct{} releaseBlock chan struct{} blockOnce sync.Once + statusCode int + body io.ReadCloser } func (a *taskPollingFetchAdaptor) Init(_ *relaycommon.RelayInfo) {} @@ -51,6 +53,16 @@ func (a *taskPollingFetchAdaptor) FetchTask(_ string, _ string, body map[string] default: } } + statusCode := a.statusCode + if statusCode == 0 { + statusCode = http.StatusOK + } + if a.body != nil { + return &http.Response{ + StatusCode: statusCode, + Body: a.body, + }, nil + } response := dto.TaskResponse[model.Task]{ Code: dto.TaskSuccessCode, @@ -65,7 +77,7 @@ func (a *taskPollingFetchAdaptor) FetchTask(_ string, _ string, body map[string] return nil, err } return &http.Response{ - StatusCode: http.StatusOK, + StatusCode: statusCode, Body: io.NopCloser(bytes.NewReader(responseBody)), }, nil } @@ -90,14 +102,29 @@ func (a *taskPollingFetchAdaptor) fetchedTaskIDs() []string { return append([]string(nil), a.taskIDs...) } +type closeTrackingReadCloser struct { + closed bool +} + +func (b *closeTrackingReadCloser) Read(_ []byte) (int, error) { + return 0, io.EOF +} + +func (b *closeTrackingReadCloser) Close() error { + b.closed = true + return nil +} + func seedTaskPollingChannel(t *testing.T, id int, disableSleep bool) { t.Helper() + baseURL := "http://example.test" ch := &model.Channel{ - Id: id, - Type: constant.ChannelTypeKling, - Name: "polling_channel", - Key: "sk-test", - Status: common.ChannelStatusEnabled, + Id: id, + Type: constant.ChannelTypeKling, + Name: "polling_channel", + Key: "sk-test", + Status: common.ChannelStatusEnabled, + BaseURL: &baseURL, } if disableSleep { ch.SetOtherSettings(dto.ChannelOtherSettings{DisableTaskPollingSleep: true}) @@ -125,6 +152,29 @@ func seedPollingTask(t *testing.T, channelID int, publicID string, upstreamID st return task } +func TestUpdateSunoTasksClosesNonOKResponseBody(t *testing.T) { + truncate(t) + + const channelID = 1001 + seedTaskPollingChannel(t, channelID, true) + task := seedPollingTask(t, channelID, "suno_public_non_ok", "suno_upstream_non_ok") + body := &closeTrackingReadCloser{} + adaptor := &taskPollingFetchAdaptor{ + statusCode: http.StatusInternalServerError, + body: body, + } + previousFactory := GetTaskAdaptorFunc + GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor } + t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory }) + + err := updateSunoTasks(context.Background(), channelID, []string{task.GetUpstreamTaskID()}, map[string]*model.Task{ + task.GetUpstreamTaskID(): task, + }) + + require.Error(t, err) + assert.True(t, body.closed) +} + func TestUpdateVideoTasksDefaultSleepWaitsBetweenTasks(t *testing.T) { truncate(t)