mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-09-08 10:14:35 +08:00
Fix: bound model response reads (#19211)
### Summary As title, from #15614, close #15613.
This commit is contained in:
@@ -213,13 +213,17 @@ func (b *BaseModel) doRequest(ctx context.Context, url string, apiConfig *APICon
|
|||||||
}
|
}
|
||||||
defer resp.Body.Close()
|
defer resp.Body.Close()
|
||||||
|
|
||||||
body, err := io.ReadAll(resp.Body)
|
if resp.StatusCode != http.StatusOK {
|
||||||
if err != nil {
|
body, err := readModelErrorBody(resp.Body)
|
||||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
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 {
|
body, err := readModelResponseBody(resp.Body)
|
||||||
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
return body, nil
|
return body, nil
|
||||||
@@ -244,13 +248,17 @@ func (b *BaseModel) doGetRequest(ctx context.Context, url string, apiConfig *API
|
|||||||
}
|
}
|
||||||
defer resp.Body.Close()
|
defer resp.Body.Close()
|
||||||
|
|
||||||
body, err := io.ReadAll(resp.Body)
|
if resp.StatusCode != http.StatusOK {
|
||||||
if err != nil {
|
body, err := readModelErrorBody(resp.Body)
|
||||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
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 {
|
body, err := readModelResponseBody(resp.Body)
|
||||||
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
return body, nil
|
return body, nil
|
||||||
@@ -274,7 +282,10 @@ func (b *BaseModel) doStreamRequest(ctx context.Context, url string, apiConfig *
|
|||||||
defer resp.Body.Close()
|
defer resp.Body.Close()
|
||||||
|
|
||||||
if resp.StatusCode != http.StatusOK {
|
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))
|
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -204,15 +204,19 @@ func (o *OpenAIModel) Embed(ctx context.Context, modelName *string, request Embe
|
|||||||
}
|
}
|
||||||
defer resp.Body.Close()
|
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)
|
body, err := io.ReadAll(resp.Body)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
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
|
var parsed openaiEmbeddingResponse
|
||||||
if err = json.Unmarshal(body, &parsed); err != nil {
|
if err = json.Unmarshal(body, &parsed); err != nil {
|
||||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
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()
|
defer resp.Body.Close()
|
||||||
|
|
||||||
body, err := io.ReadAll(resp.Body)
|
if resp.StatusCode != http.StatusOK {
|
||||||
if err != nil {
|
body, err := readModelErrorBody(resp.Body)
|
||||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
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 {
|
body, err := readModelResponseBody(resp.Body)
|
||||||
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Parse response
|
// Parse response
|
||||||
@@ -318,15 +326,19 @@ func (o *OpenAIModel) TranscribeAudio(ctx context.Context, modelName *string, fi
|
|||||||
}
|
}
|
||||||
defer resp.Body.Close()
|
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)
|
respBody, err := io.ReadAll(resp.Body)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("failed to read response body: %w", err)
|
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)
|
return decodeOpenAIASRResponse(respBody, responseFormat)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -348,7 +360,10 @@ func (o *OpenAIModel) TranscribeAudioWithSender(ctx context.Context, modelName *
|
|||||||
defer resp.Body.Close()
|
defer resp.Body.Close()
|
||||||
|
|
||||||
if resp.StatusCode != http.StatusOK {
|
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))
|
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()
|
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)
|
body, err := io.ReadAll(resp.Body)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("failed to read response body: %w", err)
|
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
|
return &TTSResponse{Audio: body}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -469,7 +488,10 @@ func (o *OpenAIModel) AudioSpeechWithSender(ctx context.Context, modelName *stri
|
|||||||
defer resp.Body.Close()
|
defer resp.Body.Close()
|
||||||
|
|
||||||
if resp.StatusCode != http.StatusOK {
|
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))
|
return fmt.Errorf("OpenAI TTS stream API error: %s, body: %s", resp.Status, string(body))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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) {
|
func TestOpenAIAudioSpeechWithSenderStreamsRawAudio(t *testing.T) {
|
||||||
withSSRFBypass(t)
|
withSSRFBypass(t)
|
||||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user