Bound and strictly decode provider responses

internal/llm/openai_compatible_client.go

@@ -25,6 +25,8 @@ var (
 	ErrMalformedResponse = errors.New("malformed llm response")
 )
 
+const maxOpenAIChatResponseBytes int64 = 16 << 20
+
 type requestFailedError struct {
 	cause error
 }
@@ -156,10 +158,13 @@ func (c *OpenAICompatibleClient) Generate(ctx context.Context, req domain.Genera
 		_, _ = io.Copy(io.Discard, io.LimitReader(httpResp.Body, 4096))
 		return nil, fmt.Errorf("%w: status=%d", ErrUnexpectedStatus, httpResp.StatusCode)
 	}
+	if httpResp.ContentLength > maxOpenAIChatResponseBytes {
+		return nil, openAIChatResponseTooLargeError()
+	}
 
-	var wireResp openAIChatResponse
-	if err := json.NewDecoder(httpResp.Body).Decode(&wireResp); err != nil {
-		return nil, fmt.Errorf("%w: failed to decode response: %v", ErrMalformedResponse, err)
+	wireResp, err := decodeOpenAIChatResponse(httpResp.Body)
+	if err != nil {
+		return nil, err
 	}
 
 	if len(wireResp.Choices) == 0 {
@@ -182,6 +187,46 @@ func (c *OpenAICompatibleClient) Generate(ctx context.Context, req domain.Genera
 	}, nil
 }
 
+func decodeOpenAIChatResponse(body io.Reader) (openAIChatResponse, error) {
+	limited := &io.LimitedReader{
+		R: body,
+		N: maxOpenAIChatResponseBytes + 1,
+	}
+	decoder := json.NewDecoder(limited)
+
+	var response openAIChatResponse
+	if err := decoder.Decode(&response); err != nil {
+		if limited.N == 0 {
+			return openAIChatResponse{}, openAIChatResponseTooLargeError()
+		}
+		return openAIChatResponse{}, fmt.Errorf("%w: failed to decode response", ErrMalformedResponse)
+	}
+	if limited.N == 0 {
+		return openAIChatResponse{}, openAIChatResponseTooLargeError()
+	}
+
+	var trailing any
+	if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
+		if limited.N == 0 {
+			return openAIChatResponse{}, openAIChatResponseTooLargeError()
+		}
+		return openAIChatResponse{}, fmt.Errorf("%w: response contains trailing data", ErrMalformedResponse)
+	}
+	if limited.N == 0 {
+		return openAIChatResponse{}, openAIChatResponseTooLargeError()
+	}
+
+	return response, nil
+}
+
+func openAIChatResponseTooLargeError() error {
+	return fmt.Errorf(
+		"%w: response exceeds %d-byte limit",
+		ErrMalformedResponse,
+		maxOpenAIChatResponseBytes,
+	)
+}
+
 func openAIChatRequestFromGenerateRequest(req domain.GenerateRequest, defaultModel string) (openAIChatRequest, error) {
 	model := strings.TrimSpace(req.Target.Model)
 	if model == "" {

internal/llm/openai_compatible_client_test.go

@@ -17,7 +17,10 @@ import (
 	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
 )
 
-var errTransportStopped = errors.New("transport stopped after request inspection")
+var (
+	errTransportStopped      = errors.New("transport stopped after request inspection")
+	errResponseReadPastLimit = errors.New("response reader was read past the allowed boundary")
+)
 
 type deadlineCapturingTransport struct {
 	deadline    time.Time
@@ -52,6 +55,72 @@ func (waitingContextTransport) RoundTrip(req *http.Request) (*http.Response, err
 	return nil, req.Context().Err()
 }
 
+type countingReadCloser struct {
+	reader    io.Reader
+	bytesRead int64
+	closed    bool
+}
+
+func (r *countingReadCloser) Read(p []byte) (int, error) {
+	n, err := r.reader.Read(p)
+	r.bytesRead += int64(n)
+	return n, err
+}
+
+func (r *countingReadCloser) Close() error {
+	r.closed = true
+	return nil
+}
+
+type repeatingByteReader byte
+
+func (r repeatingByteReader) Read(p []byte) (int, error) {
+	for i := range p {
+		p[i] = byte(r)
+	}
+	return len(p), nil
+}
+
+type guardedRepeatingReader struct {
+	value         byte
+	remaining     int64
+	readPastLimit bool
+}
+
+func (r *guardedRepeatingReader) Read(p []byte) (int, error) {
+	if r.remaining == 0 {
+		r.readPastLimit = true
+		return 0, errResponseReadPastLimit
+	}
+	if int64(len(p)) > r.remaining {
+		p = p[:r.remaining]
+	}
+	for i := range p {
+		p[i] = r.value
+	}
+	r.remaining -= int64(len(p))
+	return len(p), nil
+}
+
+const (
+	successResponsePrefix = `{"choices":[{"message":{"content":"`
+	successResponseSuffix = `"}}]}`
+	responseContentMarker = "provider-secret-fragment"
+)
+
+func sizedSuccessResponseBody(size int64) *countingReadCloser {
+	contentBytes := size - int64(len(successResponsePrefix)+len(successResponseSuffix))
+	if contentBytes < int64(len(responseContentMarker)) {
+		panic("successful response size is too small")
+	}
+	return &countingReadCloser{reader: io.MultiReader(
+		strings.NewReader(successResponsePrefix),
+		strings.NewReader(responseContentMarker),
+		io.LimitReader(repeatingByteReader('x'), contentBytes-int64(len(responseContentMarker))),
+		strings.NewReader(successResponseSuffix),
+	)}
+}
+
 func assertDeadlineNear(t *testing.T, deadline, before, after time.Time, duration time.Duration) {
 	t.Helper()
 
@@ -1055,6 +1124,205 @@ func TestOpenAICompatibleClientMalformedResponseMissingChoices(t *testing.T) {
 	}
 }
 
+func TestOpenAICompatibleClientBoundsSuccessfulResponseBodies(t *testing.T) {
+	tests := []struct {
+		name          string
+		size          int64
+		contentLength int64
+		wantErr       bool
+	}{
+		{name: "just below limit with content length", size: maxOpenAIChatResponseBytes - 1, contentLength: maxOpenAIChatResponseBytes - 1},
+		{name: "exact limit with content length", size: maxOpenAIChatResponseBytes, contentLength: maxOpenAIChatResponseBytes},
+		{name: "one byte over with content length", size: maxOpenAIChatResponseBytes + 1, contentLength: maxOpenAIChatResponseBytes + 1, wantErr: true},
+		{name: "just below limit without content length", size: maxOpenAIChatResponseBytes - 1, contentLength: -1},
+		{name: "exact limit without content length", size: maxOpenAIChatResponseBytes, contentLength: -1},
+		{name: "one byte over without content length", size: maxOpenAIChatResponseBytes + 1, contentLength: -1, wantErr: true},
+		{name: "one byte over with underreported content length", size: maxOpenAIChatResponseBytes + 1, contentLength: maxOpenAIChatResponseBytes - 1, wantErr: true},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			body := sizedSuccessResponseBody(tc.size)
+			client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
+				BaseURL: "https://provider.example/v1",
+				Model:   "m",
+				HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+					return &http.Response{
+						StatusCode:    http.StatusOK,
+						Header:        make(http.Header),
+						Body:          body,
+						ContentLength: tc.contentLength,
+						Request:       req,
+					}, nil
+				})},
+			})
+			if err != nil {
+				t.Fatalf("construct client: %v", err)
+			}
+
+			response, err := client.Generate(context.Background(), domain.GenerateRequest{
+				Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
+			})
+			if tc.wantErr {
+				if response != nil {
+					t.Fatalf("expected no partial response, got %+v", response)
+				}
+				if !errors.Is(err, ErrMalformedResponse) {
+					t.Fatalf("expected ErrMalformedResponse, got %v", err)
+				}
+				if strings.Contains(err.Error(), responseContentMarker) {
+					t.Fatalf("error exposed provider content: %v", err)
+				}
+			} else {
+				if err != nil {
+					t.Fatalf("generate: %v", err)
+				}
+				wantContentBytes := tc.size - int64(len(successResponsePrefix)+len(successResponseSuffix))
+				if int64(len(response.Content)) != wantContentBytes || !strings.HasPrefix(response.Content, responseContentMarker) {
+					t.Fatalf("unexpected response content length or prefix")
+				}
+			}
+			if body.bytesRead > maxOpenAIChatResponseBytes+1 {
+				t.Fatalf("read %d bytes, limit is %d", body.bytesRead, maxOpenAIChatResponseBytes+1)
+			}
+			if tc.wantErr && tc.contentLength > maxOpenAIChatResponseBytes && body.bytesRead != 0 {
+				t.Fatalf("read %d bytes despite oversized Content-Length", body.bytesRead)
+			}
+			if tc.wantErr && tc.contentLength <= maxOpenAIChatResponseBytes && body.bytesRead != maxOpenAIChatResponseBytes+1 {
+				t.Fatalf("read %d bytes, want one byte beyond the limit", body.bytesRead)
+			}
+			if !tc.wantErr && body.bytesRead != tc.size {
+				t.Fatalf("read %d bytes, want complete %d-byte response", body.bytesRead, tc.size)
+			}
+			if !body.closed {
+				t.Fatal("response body was not closed")
+			}
+		})
+	}
+}
+
+func TestOpenAICompatibleClientRejectsContinuingOversizedResponse(t *testing.T) {
+	prefix := successResponsePrefix + responseContentMarker
+	continuation := &guardedRepeatingReader{
+		value:     'x',
+		remaining: maxOpenAIChatResponseBytes + 1 - int64(len(prefix)),
+	}
+	body := &countingReadCloser{reader: io.MultiReader(
+		strings.NewReader(prefix),
+		continuation,
+	)}
+	client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
+		BaseURL: "https://provider.example/v1",
+		Model:   "m",
+		HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+			return &http.Response{
+				StatusCode:    http.StatusOK,
+				Header:        make(http.Header),
+				Body:          body,
+				ContentLength: -1,
+				Request:       req,
+			}, nil
+		})},
+	})
+	if err != nil {
+		t.Fatalf("construct client: %v", err)
+	}
+
+	type outcome struct {
+		response *domain.GenerateResponse
+		err      error
+	}
+	done := make(chan outcome, 1)
+	go func() {
+		response, generateErr := client.Generate(context.Background(), domain.GenerateRequest{
+			Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
+		})
+		done <- outcome{response: response, err: generateErr}
+	}()
+
+	select {
+	case result := <-done:
+		if result.response != nil {
+			t.Fatalf("expected no partial response, got %+v", result.response)
+		}
+		if !errors.Is(result.err, ErrMalformedResponse) {
+			t.Fatalf("expected ErrMalformedResponse, got %v", result.err)
+		}
+		if strings.Contains(result.err.Error(), responseContentMarker) {
+			t.Fatalf("error exposed provider content: %v", result.err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out rejecting continuing oversized response")
+	}
+	if body.bytesRead != maxOpenAIChatResponseBytes+1 {
+		t.Fatalf("read %d bytes, want %d", body.bytesRead, maxOpenAIChatResponseBytes+1)
+	}
+	if continuation.readPastLimit {
+		t.Fatal("response reader was read past one byte beyond the limit")
+	}
+	if !body.closed {
+		t.Fatal("response body was not closed")
+	}
+}
+
+func TestOpenAICompatibleClientRequiresSingleResponseDocument(t *testing.T) {
+	validResponse := `{"choices":[{"message":{"content":"ok"}}]}`
+	tests := []struct {
+		name    string
+		body    string
+		wantErr bool
+	}{
+		{name: "one document", body: validResponse},
+		{name: "trailing whitespace", body: validResponse + " \n\t\r  "},
+		{name: "trailing garbage", body: validResponse + " " + responseContentMarker, wantErr: true},
+		{name: "second JSON value", body: validResponse + ` {"detail":"` + responseContentMarker + `"}`, wantErr: true},
+		{name: "truncated document", body: validResponse[:len(validResponse)-2], wantErr: true},
+		{name: "missing choices", body: `{"choices":[]}`, wantErr: true},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			body := &countingReadCloser{reader: strings.NewReader(tc.body)}
+			client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
+				BaseURL: "https://provider.example/v1",
+				Model:   "m",
+				HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+					return &http.Response{
+						StatusCode:    http.StatusOK,
+						Header:        make(http.Header),
+						Body:          body,
+						ContentLength: int64(len(tc.body)),
+						Request:       req,
+					}, nil
+				})},
+			})
+			if err != nil {
+				t.Fatalf("construct client: %v", err)
+			}
+
+			response, err := client.Generate(context.Background(), domain.GenerateRequest{
+				Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
+			})
+			if tc.wantErr {
+				if response != nil {
+					t.Fatalf("expected no partial response, got %+v", response)
+				}
+				if !errors.Is(err, ErrMalformedResponse) {
+					t.Fatalf("expected ErrMalformedResponse, got %v", err)
+				}
+				if strings.Contains(err.Error(), responseContentMarker) {
+					t.Fatalf("error exposed provider content: %v", err)
+				}
+			} else if err != nil || response == nil || response.Content != "ok" {
+				t.Fatalf("response = %+v, error = %v", response, err)
+			}
+			if !body.closed {
+				t.Fatal("response body was not closed")
+			}
+		})
+	}
+}
+
 func TestOpenAICompatibleClientGenerationTimeoutSetsEarlierDeadline(t *testing.T) {
 	transport := &deadlineCapturingTransport{err: context.DeadlineExceeded}
 	generationTimeout := 2 * time.Second