Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 43 additions & 0 deletions pkg/utils/apiutil/apiutil.go
Original file line number Diff line number Diff line change
Expand Up @@ -484,7 +484,50 @@ func NewCustomReverseProxies(dialClient *http.Client, urls []url.URL) http.Handl
return p
}

// EnsureRewindableBody makes r's body safe to be retried by net/http
// Transport on connection loss. It is a no-op when the body is nil,
// http.NoBody, or r.GetBody is already set. Otherwise it drains the body
// into memory and wires up GetBody so the transport can rewind.
//
// Callers should ensure the body fits in memory; this helper buffers the
// entire payload. It guards against the
// "net/http: cannot rewind body after connection loss" error that can
// occur when a server-side request (GetBody == nil) is forwarded via
// http.Client.Do and the underlying keep-alive connection goes stale.
func EnsureRewindableBody(r *http.Request) error {
if r.Body == nil || r.Body == http.NoBody || r.GetBody != nil {
return nil
}
buf, err := io.ReadAll(r.Body)
_ = r.Body.Close()
if err != nil {
return err
}
// Restore NoBody semantics for empty payloads so that Transport's
// outgoingLength returns 0 and no body probing happens.
if len(buf) == 0 {
r.Body = http.NoBody
r.GetBody = func() (io.ReadCloser, error) { return http.NoBody, nil }
r.ContentLength = 0
return nil
}
r.Body = io.NopCloser(bytes.NewReader(buf))
r.GetBody = func() (io.ReadCloser, error) {
return io.NopCloser(bytes.NewReader(buf)), nil
}
// We now know the exact length; set it so the transport can pick
// Content-Length framing over chunked encoding when forwarding.
r.ContentLength = int64(len(buf))
return nil
Comment on lines +501 to +521

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Verify current helper has unbounded body buffering and where it is used.
rg -n 'func EnsureRewindableBody|io.ReadAll\(r.Body\)|LimitedReader|maxRewindableBodyBytes' pkg/utils/apiutil/apiutil.go
rg -n 'EnsureRewindableBody\(' pkg/utils/apiutil/apiutil.go pkg/utils/requestutil/request_info.go

Repository: tikv/pd

Length of output: 425


🏁 Script executed:

# Examine the ServeHTTP handler context where EnsureRewindableBody is called
sed -n '520,535p' pkg/utils/apiutil/apiutil.go

# Check for any request size limit configuration in the server setup
rg -n 'MaxHeaderBytes|ContentLength|MaxBytesReader|LimitedReader|maxBody|size.*limit' pkg/utils/apiutil/ pkg/utils/requestutil/ -A 2 -B 2 | head -50

# Check the request_info.go usage context
sed -n '65,85p' pkg/utils/requestutil/request_info.go

Repository: tikv/pd

Length of output: 2553


🏁 Script executed:

# Check for server configuration and request size limits
rg -n 'http\.Server|MaxHeaderBytes|MaxBytesReader|ListenAndServe|NewServer' pkg/ cmd/ --type go -B 2 -A 3 | head -60

# Check if there's any middleware that validates request size
rg -n 'ContentLength.*<|request.*size|body.*limit|ReadLimitCloser' pkg/ cmd/ --type go | head -20

# Check what the error message would look like
sed -n '497,530p' pkg/utils/apiutil/apiutil.go

Repository: tikv/pd

Length of output: 4467


Add a max body-size guard before buffering.

Line 501 reads the full request body into memory with no upper bound. When called in the ServeHTTP reverse proxy handler (line 525), this creates a memory-exhaustion vector—an attacker can send large payloads to exhaust memory and trigger OOM without any server-side size limit configured.

Additionally, line 528 exposes the raw error message to the client via http.Error(w, err.Error(), ...), which should not leak internal details.

🔧 Suggested fix
+const maxRewindableBodyBytes = 8 << 20 // 8 MiB, tune as needed
+
 func EnsureRewindableBody(r *http.Request) error {
 	if r.Body == nil || r.Body == http.NoBody || r.GetBody != nil {
 		return nil
 	}
-	buf, err := io.ReadAll(r.Body)
+	if r.ContentLength > maxRewindableBodyBytes {
+		_ = r.Body.Close()
+		return errors.Errorf("request body too large: %d", r.ContentLength)
+	}
+	lr := &io.LimitedReader{R: r.Body, N: maxRewindableBodyBytes + 1}
+	buf, err := io.ReadAll(lr)
 	_ = r.Body.Close()
 	if err != nil {
 		return err
 	}
+	if int64(len(buf)) > maxRewindableBodyBytes {
+		return errors.Errorf("request body too large")
+	}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@pkg/utils/apiutil/apiutil.go` around lines 501 - 521, The code that calls
io.ReadAll(r.Body) must enforce a max-body-size to prevent OOM: change the read
to use a bounded reader (e.g., buf, err := io.ReadAll(io.LimitReader(r.Body,
maxBodySize+1))) and check if len(buf) > maxBodySize then return a specific
sentinel error (e.g., ErrRequestBodyTooLarge) instead of buffering unbounded
data; restore r.Body/GetBody from the buffered bytes only when within the limit.
In the ServeHTTP reverse proxy handler, do not propagate err.Error() to
clients—handle ErrRequestBodyTooLarge by responding with a generic 413 Payload
Too Large message (or use http.MaxBytesReader at the handler boundary as an
alternative) and for other errors return a sanitized message. Ensure the symbols
referenced are the io.ReadAll call and r.Body / r.GetBody assignments and the
ServeHTTP reverse proxy handler.

}

func (p *customReverseProxies) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if err := EnsureRewindableBody(r); err != nil {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to move it before Do?

log.Error("failed to read request body", zap.Error(err))
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}

for _, url := range p.urls {
r.RequestURI = ""
r.URL.Host = url.Host
Expand Down
103 changes: 103 additions & 0 deletions pkg/utils/apiutil/apiutil_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ package apiutil

import (
"bytes"
"errors"
"io"
"net/http"
"net/http/httptest"
Expand Down Expand Up @@ -210,6 +211,108 @@ func TestGetIPPortFromHTTPRequest(t *testing.T) {
}
}

type errReader struct{ err error }

func (e *errReader) Read(_ []byte) (int, error) { return 0, e.err }
func (*errReader) Close() error { return nil }

type closeTracker struct {
io.Reader
closed bool
}

func (c *closeTracker) Close() error {
c.closed = true
return nil
}

func TestEnsureRewindableBody(t *testing.T) {
t.Run("nil body is a no-op", func(t *testing.T) {
re := require.New(t)
r := &http.Request{}
re.NoError(EnsureRewindableBody(r))
re.Nil(r.Body)
re.Nil(r.GetBody)
})

t.Run("http.NoBody is a no-op", func(t *testing.T) {
re := require.New(t)
r := &http.Request{Body: http.NoBody}
re.NoError(EnsureRewindableBody(r))
re.Equal(http.NoBody, r.Body)
re.Nil(r.GetBody)
})

t.Run("existing GetBody is preserved", func(t *testing.T) {
re := require.New(t)
orig := io.NopCloser(bytes.NewBufferString("payload"))
called := false
getBody := func() (io.ReadCloser, error) {
called = true
return io.NopCloser(bytes.NewBufferString("payload")), nil
}
r := &http.Request{Body: orig, GetBody: getBody}
re.NoError(EnsureRewindableBody(r))
// Body untouched, GetBody untouched (we only check it wasn't replaced
// by invoking it and confirming our own sentinel).
re.Equal(orig, r.Body)
_, err := r.GetBody()
re.NoError(err)
re.True(called)
})

t.Run("empty body is restored to NoBody", func(t *testing.T) {
re := require.New(t)
tracker := &closeTracker{Reader: bytes.NewReader(nil)}
r := &http.Request{Body: tracker}
re.NoError(EnsureRewindableBody(r))
re.True(tracker.closed, "original body should be closed")
re.Equal(http.NoBody, r.Body)
re.EqualValues(0, r.ContentLength)
re.NotNil(r.GetBody)
rc, err := r.GetBody()
re.NoError(err)
re.Equal(http.NoBody, rc)
})

t.Run("non-empty body becomes rewindable", func(t *testing.T) {
re := require.New(t)
payload := []byte(`{"hello":"world"}`)
tracker := &closeTracker{Reader: bytes.NewReader(payload)}
r := &http.Request{Body: tracker, ContentLength: -1}
re.NoError(EnsureRewindableBody(r))
re.True(tracker.closed, "original body should be closed")
re.EqualValues(len(payload), r.ContentLength)
re.NotNil(r.GetBody)

// Draining r.Body once should yield the payload.
got, err := io.ReadAll(r.Body)
re.NoError(err)
re.Equal(payload, got)
re.NoError(r.Body.Close())

// GetBody should be invokable multiple times and each returned
// ReadCloser should independently yield the same payload -- this is
// the actual rewindability guarantee we care about.
for range 3 {
rc, err := r.GetBody()
re.NoError(err)
got, err := io.ReadAll(rc)
re.NoError(err)
re.Equal(payload, got)
re.NoError(rc.Close())
}
})

t.Run("read error is propagated", func(t *testing.T) {
re := require.New(t)
wantErr := errors.New("boom")
r := &http.Request{Body: &errReader{err: wantErr}}
err := EnsureRewindableBody(r)
re.ErrorIs(err, wantErr)
})
}

func TestParseHexKeys(t *testing.T) {
re := require.New(t)
// Test for hex format
Expand Down
18 changes: 10 additions & 8 deletions pkg/utils/requestutil/request_info.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
package requestutil

import (
"bytes"
"encoding/json"
"fmt"
"io"
Expand Down Expand Up @@ -69,13 +68,16 @@ func getURLParam(r *http.Request) string {
}

func getBodyParam(r *http.Request) string {
if r.Body == nil {
// Make the body rewindable so downstream forwarding can retry on
// connection loss, then read it back via GetBody for audit logging.
if err := apiutil.EnsureRewindableBody(r); err != nil || r.GetBody == nil {
return ""
}
// http request body is a io.Reader between bytes.Reader and strings.Reader, it only has EOF error
buf, _ := io.ReadAll(r.Body)
r.Body.Close()
bodyParam := string(buf)
r.Body = io.NopCloser(bytes.NewBuffer(buf))
return bodyParam
rc, err := r.GetBody()
if err != nil {
return ""
}
defer rc.Close()
buf, _ := io.ReadAll(rc)
return string(buf)
}
Loading