mirror of
https://github.com/golang/go
synced 2024-11-19 02:14:43 -07:00
net/http: deflake TestRequestBodyLimit
This test can return with a Transport still processing an in-flight request, resulting in a test failure due to the leaked Transport. Avoid this by waiting for the Transport to close the request body before returning. Fixes #60264 Change-Id: I8d8b54f633c2e28da2b1bf1bc01ce09dd77769de Reviewed-on: https://go-review.googlesource.com/c/go/+/522695 Reviewed-by: Bryan Mills <bcmills@google.com> Auto-Submit: Damien Neil <dneil@google.com> TryBot-Result: Gopher Robot <gobot@golang.org> Run-TryBot: Damien Neil <dneil@google.com>
This commit is contained in:
parent
611c3af30d
commit
92bbecc518
@ -2967,15 +2967,36 @@ func (b neverEnding) Read(p []byte) (n int, err error) {
|
|||||||
return len(p), nil
|
return len(p), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
type countReader struct {
|
type bodyLimitReader struct {
|
||||||
r io.Reader
|
mu sync.Mutex
|
||||||
n *int64
|
count int
|
||||||
|
limit int
|
||||||
|
closed chan struct{}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (cr countReader) Read(p []byte) (n int, err error) {
|
func (r *bodyLimitReader) Read(p []byte) (int, error) {
|
||||||
n, err = cr.r.Read(p)
|
r.mu.Lock()
|
||||||
atomic.AddInt64(cr.n, int64(n))
|
defer r.mu.Unlock()
|
||||||
return
|
select {
|
||||||
|
case <-r.closed:
|
||||||
|
return 0, errors.New("closed")
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
if r.count > r.limit {
|
||||||
|
return 0, errors.New("at limit")
|
||||||
|
}
|
||||||
|
r.count += len(p)
|
||||||
|
for i := range p {
|
||||||
|
p[i] = 'a'
|
||||||
|
}
|
||||||
|
return len(p), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *bodyLimitReader) Close() error {
|
||||||
|
r.mu.Lock()
|
||||||
|
defer r.mu.Unlock()
|
||||||
|
close(r.closed)
|
||||||
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestRequestBodyLimit(t *testing.T) { run(t, testRequestBodyLimit) }
|
func TestRequestBodyLimit(t *testing.T) { run(t, testRequestBodyLimit) }
|
||||||
@ -2999,8 +3020,11 @@ func testRequestBodyLimit(t *testing.T, mode testMode) {
|
|||||||
}
|
}
|
||||||
}))
|
}))
|
||||||
|
|
||||||
nWritten := new(int64)
|
body := &bodyLimitReader{
|
||||||
req, _ := NewRequest("POST", cst.ts.URL, io.LimitReader(countReader{neverEnding('a'), nWritten}, limit*200))
|
closed: make(chan struct{}),
|
||||||
|
limit: limit * 200,
|
||||||
|
}
|
||||||
|
req, _ := NewRequest("POST", cst.ts.URL, body)
|
||||||
|
|
||||||
// Send the POST, but don't care it succeeds or not. The
|
// Send the POST, but don't care it succeeds or not. The
|
||||||
// remote side is going to reply and then close the TCP
|
// remote side is going to reply and then close the TCP
|
||||||
@ -3015,10 +3039,13 @@ func testRequestBodyLimit(t *testing.T, mode testMode) {
|
|||||||
if err == nil {
|
if err == nil {
|
||||||
resp.Body.Close()
|
resp.Body.Close()
|
||||||
}
|
}
|
||||||
|
// Wait for the Transport to finish writing the request body.
|
||||||
|
// It will close the body when done.
|
||||||
|
<-body.closed
|
||||||
|
|
||||||
if atomic.LoadInt64(nWritten) > limit*100 {
|
if body.count > limit*100 {
|
||||||
t.Errorf("handler restricted the request body to %d bytes, but client managed to write %d",
|
t.Errorf("handler restricted the request body to %d bytes, but client managed to write %d",
|
||||||
limit, nWritten)
|
limit, body.count)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user