mirror of
https://github.com/netbirdio/netbird.git
synced 2026-08-22 23:51:29 +02:00
[proxy] Forward an oversized model listing whole
The discovery filter read one byte past its 1 MiB cap to detect a body too large to rewrite, then closed the upstream body and forwarded the buffer it had — the response reached the client truncated at exactly the cap, with Content-Length rewritten to match so nothing looked wrong until the client tried to parse it. Splice the bytes already read back in front of the unread remainder and forward the response as the upstream sent it, headers untouched.
This commit is contained in:
@@ -51,18 +51,25 @@ func filterModelListing(resp *http.Response, permitted map[string]struct{}) erro
|
||||
return nil
|
||||
}
|
||||
|
||||
// One byte past the cap, so an oversized body is detectable without
|
||||
// buffering all of it.
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, maxDiscoveryBodyBytes+1))
|
||||
closeErr := resp.Body.Close()
|
||||
if err != nil {
|
||||
_ = resp.Body.Close()
|
||||
return err
|
||||
}
|
||||
if closeErr != nil {
|
||||
return closeErr
|
||||
}
|
||||
if len(body) > maxDiscoveryBodyBytes {
|
||||
restoreBody(resp, body)
|
||||
// Too large to filter. Put the bytes already read back in front of the
|
||||
// unread remainder and forward the response exactly as the upstream
|
||||
// sent it, headers included. Buffering what was read and closing here
|
||||
// would truncate the body at the cap and hand the client a short,
|
||||
// invalid listing — worse than not filtering at all.
|
||||
resp.Body = spliceBody(body, resp.Body)
|
||||
return nil
|
||||
}
|
||||
if err := resp.Body.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
filtered, ok := filterListingBody(body, permitted)
|
||||
if !ok {
|
||||
@@ -164,6 +171,19 @@ func modelIDForms(id string) []string {
|
||||
|
||||
// restoreBody puts body back on the response and fixes the length headers
|
||||
// so the client reads exactly what is there.
|
||||
// spliceBody returns a ReadCloser that yields prefix followed by whatever is
|
||||
// left in rest, closing rest when closed. It lets the filter put back bytes it
|
||||
// consumed while deciding, without owning the rest of the stream.
|
||||
func spliceBody(prefix []byte, rest io.ReadCloser) io.ReadCloser {
|
||||
return struct {
|
||||
io.Reader
|
||||
io.Closer
|
||||
}{
|
||||
Reader: io.MultiReader(bytes.NewReader(prefix), rest),
|
||||
Closer: rest,
|
||||
}
|
||||
}
|
||||
|
||||
func restoreBody(resp *http.Response, body []byte) {
|
||||
resp.Body = io.NopCloser(bytes.NewReader(body))
|
||||
resp.ContentLength = int64(len(body))
|
||||
|
||||
@@ -170,3 +170,48 @@ func TestModelDiscoveryFilter_KeepsSlashBearingIDs(t *testing.T) {
|
||||
assert.Equal(t, []string{"Qwen/Qwen2.5-0.5B-Instruct"}, ids,
|
||||
"a slash inside the model id is part of the id, not a provider prefix")
|
||||
}
|
||||
|
||||
// TestModelDiscoveryFilter_ForwardsOversizedBodyIntact covers a listing past
|
||||
// the buffering cap. The filter reads one byte beyond the cap to detect the
|
||||
// size; forwarding only what it read would hand the client a body truncated
|
||||
// at exactly 1 MiB — valid-looking, short, and unparseable as JSON. The bytes
|
||||
// already read must be spliced back in front of the unread remainder so the
|
||||
// response reaches the client exactly as the upstream sent it.
|
||||
func TestModelDiscoveryFilter_ForwardsOversizedBodyIntact(t *testing.T) {
|
||||
// A well-formed listing whose single entry pads the body past the cap.
|
||||
padding := strings.Repeat("x", maxDiscoveryBodyBytes)
|
||||
body := `{"object":"list","data":[{"id":"gpt-4o","note":"` + padding + `"}]}`
|
||||
require.Greater(t, len(body), maxDiscoveryBodyBytes+1,
|
||||
"the fixture must exceed the cap by more than the one-byte probe")
|
||||
|
||||
resp := jsonListingResponse(body) //nolint:bodyclose // in-memory body
|
||||
require.NoError(t, modelDiscoveryFilter(nil, nil)(resp)) //nolint:bodyclose // in-memory body
|
||||
|
||||
got, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, len(body), len(got),
|
||||
"an oversized listing must reach the client whole, not truncated at the cap")
|
||||
assert.Equal(t, body, string(got), "the forwarded bytes must be the upstream's own")
|
||||
|
||||
var doc map[string]json.RawMessage
|
||||
assert.NoError(t, json.Unmarshal(got, &doc),
|
||||
"the forwarded body must still parse as JSON")
|
||||
}
|
||||
|
||||
// TestModelDiscoveryFilter_OversizedBodyKeepsUpstreamHeaders pins that the
|
||||
// oversized path leaves the response metadata alone. Rewriting Content-Length
|
||||
// to the truncated prefix is what made the corruption invisible to the client
|
||||
// until it tried to parse.
|
||||
func TestModelDiscoveryFilter_OversizedBodyKeepsUpstreamHeaders(t *testing.T) {
|
||||
padding := strings.Repeat("x", maxDiscoveryBodyBytes)
|
||||
body := `{"object":"list","data":[{"id":"gpt-4o","note":"` + padding + `"}]}`
|
||||
|
||||
resp := jsonListingResponse(body) //nolint:bodyclose // in-memory body
|
||||
resp.Header.Set("Content-Length", strconv.Itoa(len(body)))
|
||||
require.NoError(t, modelDiscoveryFilter(nil, nil)(resp)) //nolint:bodyclose // in-memory body
|
||||
|
||||
assert.Equal(t, int64(len(body)), resp.ContentLength,
|
||||
"ContentLength must keep describing the body the client receives")
|
||||
assert.Equal(t, strconv.Itoa(len(body)), resp.Header.Get("Content-Length"),
|
||||
"the Content-Length header must not be rewritten to the truncated prefix")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user