Fix: bound model response reads (#19211)

### Summary As title, from #15614, close #15613.
2026-09-08 02:08:01 +08:00 · 2026-09-04 17:46:59 +08:00
parent ff69dba542
commit fd1e515ecf
5 changed files with 237 additions and 30 deletions
--- a/internal/entity/models/base_model.go
+++ b/internal/entity/models/base_model.go
@@ -213,13 +213,17 @@ func (b *BaseModel) doRequest(ctx context.Context, url string, apiConfig *APICon
 	}
 	defer resp.Body.Close()

-	body, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read response: %w", err)
+	if resp.StatusCode != http.StatusOK {
+		body, err := readModelErrorBody(resp.Body)
+		if err != nil {
+			return nil, fmt.Errorf("API request failed with status %d; failed to read error response: %w", resp.StatusCode, err)
+		}
+		return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
 	}

-	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
+	body, err := readModelResponseBody(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response: %w", err)
 	}

 	return body, nil
@@ -244,13 +248,17 @@ func (b *BaseModel) doGetRequest(ctx context.Context, url string, apiConfig *API
 	}
 	defer resp.Body.Close()

-	body, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read response: %w", err)
+	if resp.StatusCode != http.StatusOK {
+		body, err := readModelErrorBody(resp.Body)
+		if err != nil {
+			return nil, fmt.Errorf("API request failed with status %d; failed to read error response: %w", resp.StatusCode, err)
+		}
+		return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
 	}

-	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
+	body, err := readModelResponseBody(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response: %w", err)
 	}

 	return body, nil
@@ -274,7 +282,10 @@ func (b *BaseModel) doStreamRequest(ctx context.Context, url string, apiConfig *
 	defer resp.Body.Close()

 	if resp.StatusCode != http.StatusOK {
-		body, _ := io.ReadAll(resp.Body)
+		body, err := readModelErrorBody(resp.Body)
+		if err != nil {
+			return fmt.Errorf("API request failed with status %d; failed to read error response: %w", resp.StatusCode, err)
+		}
 		return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
 	}

--- a/internal/entity/models/openai.go
+++ b/internal/entity/models/openai.go
@@ -204,15 +204,19 @@ func (o *OpenAIModel) Embed(ctx context.Context, modelName *string, request Embe
 	}
 	defer resp.Body.Close()

+	if resp.StatusCode != http.StatusOK {
+		body, err := readModelErrorBody(resp.Body)
+		if err != nil {
+			return nil, fmt.Errorf("OpenAI embeddings API error: %s, failed to read error response body: %w", resp.Status, err)
+		}
+		return nil, fmt.Errorf("OpenAI embeddings API error: %s, body: %s", resp.Status, string(body))
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response: %w", err)
 	}

-	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("OpenAI embeddings API error: %s, body: %s", resp.Status, string(body))
-	}
-
 	var parsed openaiEmbeddingResponse
 	if err = json.Unmarshal(body, &parsed); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -259,13 +263,17 @@ func (o *OpenAIModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([]L
 	}
 	defer resp.Body.Close()

-	body, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read response: %w", err)
+	if resp.StatusCode != http.StatusOK {
+		body, err := readModelErrorBody(resp.Body)
+		if err != nil {
+			return nil, fmt.Errorf("API request failed with status %d; failed to read error response: %w", resp.StatusCode, err)
+		}
+		return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
 	}

-	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
+	body, err := readModelResponseBody(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response: %w", err)
 	}

 	// Parse response
@@ -318,15 +326,19 @@ func (o *OpenAIModel) TranscribeAudio(ctx context.Context, modelName *string, fi
 	}
 	defer resp.Body.Close()

+	if resp.StatusCode != http.StatusOK {
+		respBody, err := readModelErrorBody(resp.Body)
+		if err != nil {
+			return nil, fmt.Errorf("OpenAI ASR API error: %s, failed to read error response body: %w", resp.Status, err)
+		}
+		return nil, fmt.Errorf("OpenAI ASR API error: %s, body: %s", resp.Status, string(respBody))
+	}
+
 	respBody, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}

-	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("OpenAI ASR API error: %s, body: %s", resp.Status, string(respBody))
-	}
-
 	return decodeOpenAIASRResponse(respBody, responseFormat)
 }

@@ -348,7 +360,10 @@ func (o *OpenAIModel) TranscribeAudioWithSender(ctx context.Context, modelName *
 	defer resp.Body.Close()

 	if resp.StatusCode != http.StatusOK {
-		respBody, _ := io.ReadAll(resp.Body)
+		respBody, err := readModelErrorBody(resp.Body)
+		if err != nil {
+			return fmt.Errorf("OpenAI ASR stream API error: %s, failed to read error response body: %w", resp.Status, err)
+		}
 		return fmt.Errorf("OpenAI ASR stream API error: %s, body: %s", resp.Status, string(respBody))
 	}

@@ -437,15 +452,19 @@ func (o *OpenAIModel) AudioSpeech(ctx context.Context, modelName *string, audioC
 	}
 	defer resp.Body.Close()

+	if resp.StatusCode != http.StatusOK {
+		body, err := readModelErrorBody(resp.Body)
+		if err != nil {
+			return nil, fmt.Errorf("OpenAI TTS API error: %s, failed to read error response body: %w", resp.Status, err)
+		}
+		return nil, fmt.Errorf("OpenAI TTS API error: %s, body: %s", resp.Status, string(body))
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}

-	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("OpenAI TTS API error: %s, body: %s", resp.Status, string(body))
-	}
-
 	return &TTSResponse{Audio: body}, nil
 }

@@ -469,7 +488,10 @@ func (o *OpenAIModel) AudioSpeechWithSender(ctx context.Context, modelName *stri
 	defer resp.Body.Close()

 	if resp.StatusCode != http.StatusOK {
-		body, _ := io.ReadAll(resp.Body)
+		body, err := readModelErrorBody(resp.Body)
+		if err != nil {
+			return fmt.Errorf("OpenAI TTS stream API error: %s, failed to read error response body: %w", resp.Status, err)
+		}
 		return fmt.Errorf("OpenAI TTS stream API error: %s, body: %s", resp.Status, string(body))
 	}

--- a/internal/entity/models/openai_test.go
+++ b/internal/entity/models/openai_test.go
@@ -323,6 +323,95 @@ func TestOpenAIAudioSpeechRejectsNonStringVoice(t *testing.T) {
 	}
 }

+func TestOpenAIChatWithMessagesBoundsErrorResponseBody(t *testing.T) {
+	withSSRFBypass(t)
+	ctx := t.Context()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusBadGateway)
+		_, _ = io.WriteString(w, strings.Repeat("x", int(maxModelErrorBodyBytes)+1))
+	}))
+	defer srv.Close()
+
+	apiKey := "test-key"
+	_, err := newOpenAIForTest(srv.URL).ChatWithMessages(
+		ctx,
+		"gpt-4o-mini",
+		[]Message{{Role: "user", Content: "hello"}},
+		&APIConfig{ApiKey: &apiKey},
+		nil,
+		nil,
+	)
+	if err == nil {
+		t.Fatal("expected error")
+	}
+	if !strings.Contains(err.Error(), "API request failed with status 502") {
+		t.Fatalf("err=%v, want status 502", err)
+	}
+	if !strings.Contains(err.Error(), "failed to read error response") {
+		t.Fatalf("err=%v, want failed to read error response", err)
+	}
+	if !strings.Contains(err.Error(), "response body exceeds") {
+		t.Fatalf("err=%v, want response body exceeds", err)
+	}
+}
+
+func TestOpenAIListModelsBoundsErrorResponseBody(t *testing.T) {
+	withSSRFBypass(t)
+	ctx := t.Context()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusBadGateway)
+		_, _ = io.WriteString(w, strings.Repeat("x", int(maxModelErrorBodyBytes)+1))
+	}))
+	defer srv.Close()
+
+	apiKey := "test-key"
+	_, err := newOpenAIForTest(srv.URL).ListModels(ctx, &APIConfig{ApiKey: &apiKey})
+	if err == nil {
+		t.Fatal("expected error")
+	}
+	if !strings.Contains(err.Error(), "API request failed with status 502") {
+		t.Fatalf("err=%v, want status 502", err)
+	}
+	if !strings.Contains(err.Error(), "failed to read error response") {
+		t.Fatalf("err=%v, want failed to read error response", err)
+	}
+	if !strings.Contains(err.Error(), "response body exceeds") {
+		t.Fatalf("err=%v, want response body exceeds", err)
+	}
+}
+
+func TestOpenAIAudioSpeechBoundsErrorResponseBody(t *testing.T) {
+	withSSRFBypass(t)
+	ctx := t.Context()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+		_, _ = io.WriteString(w, strings.Repeat("x", int(maxModelErrorBodyBytes)+1))
+	}))
+	defer srv.Close()
+
+	apiKey := "test-key"
+	model := "tts-1"
+	input := "hello"
+
+	_, err := newOpenAIForTest(srv.URL).AudioSpeech(
+		ctx,
+		&model,
+		&input,
+		&APIConfig{ApiKey: &apiKey},
+		&TTSConfig{Params: map[string]interface{}{"voice": "alloy"}},
+		nil,
+	)
+	if err == nil {
+		t.Fatal("expected error")
+	}
+	if !strings.Contains(err.Error(), "failed to read error response body") {
+		t.Fatalf("err=%v, want failed to read error response body", err)
+	}
+	if !strings.Contains(err.Error(), "response body exceeds") {
+		t.Fatalf("err=%v, want response body exceeds", err)
+	}
+}
+
 func TestOpenAIAudioSpeechWithSenderStreamsRawAudio(t *testing.T) {
 	withSSRFBypass(t)
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
--- a/internal/entity/models/response_body.go
+++ b/internal/entity/models/response_body.go
@@ -0,0 +1,55 @@
+//
+//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+package models
+
+import (
+	"fmt"
+	"io"
+)
+
+const (
+	// maxModelResponseBodyBytes caps JSON/text provider responses that should stay modest.
+	maxModelResponseBodyBytes int64 = 16 << 20
+	// maxModelErrorBodyBytes caps provider error pages before they are included in errors.
+	maxModelErrorBodyBytes int64 = 1 << 20
+)
+
+func readModelResponseBodyLimited(body io.Reader, maxBytes int64) ([]byte, error) {
+	if body == nil {
+		return nil, fmt.Errorf("response body is nil")
+	}
+	if maxBytes <= 0 {
+		return nil, fmt.Errorf("response body limit must be positive")
+	}
+
+	data, err := io.ReadAll(io.LimitReader(body, maxBytes+1))
+	if err != nil {
+		return nil, err
+	}
+	if int64(len(data)) > maxBytes {
+		return nil, fmt.Errorf("response body exceeds %d bytes", maxBytes)
+	}
+	return data, nil
+}
+
+func readModelResponseBody(body io.Reader) ([]byte, error) {
+	return readModelResponseBodyLimited(body, maxModelResponseBodyBytes)
+}
+
+func readModelErrorBody(body io.Reader) ([]byte, error) {
+	return readModelResponseBodyLimited(body, maxModelErrorBodyBytes)
+}
--- a/internal/entity/models/response_body_test.go
+++ b/internal/entity/models/response_body_test.go
@@ -0,0 +1,30 @@
+package models
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestReadModelResponseBodyLimited(t *testing.T) {
+	body, err := readModelResponseBodyLimited(strings.NewReader("abcd"), 4)
+	if err != nil {
+		t.Fatalf("readModelResponseBodyLimited: %v", err)
+	}
+	if string(body) != "abcd" {
+		t.Fatalf("body=%q, want abcd", string(body))
+	}
+}
+
+func TestReadModelResponseBodyLimitedRejectsOversizedBody(t *testing.T) {
+	_, err := readModelResponseBodyLimited(strings.NewReader("abcde"), 4)
+	if err == nil || !strings.Contains(err.Error(), "response body exceeds 4 bytes") {
+		t.Fatalf("err=%v, want response body exceeds 4 bytes", err)
+	}
+}
+
+func TestReadModelResponseBodyLimitedRejectsInvalidLimit(t *testing.T) {
+	_, err := readModelResponseBodyLimited(strings.NewReader("abcd"), 0)
+	if err == nil || !strings.Contains(err.Error(), "response body limit must be positive") {
+		t.Fatalf("err=%v, want positive limit error", err)
+	}
+}