diff --git a/internal/llm/client.go b/internal/llm/client.go index bc93b36..7b6bba0 100644 --- a/internal/llm/client.go +++ b/internal/llm/client.go @@ -9,6 +9,7 @@ import ( "context" "encoding/json" "fmt" + "io" "log" "net/http" "sync" @@ -187,6 +188,12 @@ type resp struct { } `json:"choices"` } +// MaxResponseBytes bounds one llama-server completion response. Even a full +// 32k-token answer is far smaller than 1 MiB; anything larger is a malformed +// or hostile peer, not useful model output. The same client reaches the LAN +// workstation, so loopback trust is not a sufficient bound. +const MaxResponseBytes int64 = 1 << 20 + func (c *Client) Complete(ctx context.Context, r Req) (string, error) { // Priority first: a background request can sit here for a while, and it // must not be counted against the swap drain while it waits. @@ -230,8 +237,15 @@ func (c *Client) Complete(ctx context.Context, r Req) (string, error) { if httpResp.StatusCode != 200 { return "", fmt.Errorf("llm: status %d", httpResp.StatusCode) } + raw, err := io.ReadAll(io.LimitReader(httpResp.Body, MaxResponseBytes+1)) + if err != nil { + return "", fmt.Errorf("llm: read response: %w", err) + } + if int64(len(raw)) > MaxResponseBytes { + return "", fmt.Errorf("llm: response exceeds %d bytes", MaxResponseBytes) + } var out resp - if err := json.NewDecoder(httpResp.Body).Decode(&out); err != nil { + if err := json.Unmarshal(raw, &out); err != nil { return "", err } if len(out.Choices) == 0 { diff --git a/internal/llm/client_test.go b/internal/llm/client_test.go index 9526c72..237dfad 100644 --- a/internal/llm/client_test.go +++ b/internal/llm/client_test.go @@ -91,3 +91,19 @@ func TestSetBaseURL(t *testing.T) { t.Errorf("request after the swap went to %q; want B", hit) } } + +func TestCompleteRejectsOversizeResponse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"`)) + _, _ = w.Write([]byte(strings.Repeat("x", int(MaxResponseBytes)))) + _, _ = w.Write([]byte(`"}}]}`)) + })) + defer srv.Close() + + c := New(srv.URL, 5*time.Second) + _, err := c.Complete(context.Background(), Req{User: "x"}) + if err == nil || !strings.Contains(err.Error(), "response exceeds") { + t.Fatalf("Complete oversize error = %v, want bounded-response error", err) + } +}