more comments about buffering and add new tests

This commit is contained in:
WeidiDeng 2025-01-09 10:49:06 +08:00
parent 0e570e0cc7
commit 2f2c77167a
No known key found for this signature in database
GPG key ID: 25F87CE1741EC7CD
2 changed files with 88 additions and 0 deletions

View file

@ -0,0 +1,84 @@
package reverseproxy
import (
"io"
"testing"
)
type zeroReader struct{}
func (zeroReader) Read(p []byte) (int, error) {
for i := range p {
p[i] = 0
}
return len(p), nil
}
func TestBuffering(t *testing.T) {
var (
h Handler
zr zeroReader
)
type args struct {
body io.ReadCloser
limit int64
}
tests := []struct {
name string
args args
resultCheck func(io.ReadCloser, int64, args) bool
}{
{
name: "0 limit, body is returned as is",
args: args{
body: io.NopCloser(&zr),
limit: 0,
},
resultCheck: func(res io.ReadCloser, read int64, args args) bool {
return res == args.body && read == args.limit && read == 0
},
},
{
name: "negative limit, body is read completely",
args: args{
body: io.NopCloser(io.LimitReader(&zr, 100)),
limit: -1,
},
resultCheck: func(res io.ReadCloser, read int64, args args) bool {
brc, ok := res.(bodyReadCloser)
return ok && brc.body == nil && brc.buf.Len() == 100 && read == 100
},
},
{
name: "positive limit, body is read partially",
args: args{
body: io.NopCloser(io.LimitReader(&zr, 100)),
limit: 50,
},
resultCheck: func(res io.ReadCloser, read int64, args args) bool {
brc, ok := res.(bodyReadCloser)
return ok && brc.body != nil && brc.buf.Len() == 50 && read == 50
},
},
{
name: "positive limit, body is read completely",
args: args{
body: io.NopCloser(io.LimitReader(&zr, 100)),
limit: 101,
},
resultCheck: func(res io.ReadCloser, read int64, args args) bool {
brc, ok := res.(bodyReadCloser)
return ok && brc.body == nil && brc.buf.Len() == 100 && read == 100
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
res, read := h.bufferedBody(tt.args.body, tt.args.limit)
if !tt.resultCheck(res, read, tt.args) {
t.Error("Handler.bufferedBody() test failed")
return
}
})
}
}

View file

@ -1221,6 +1221,10 @@ func (h Handler) provisionUpstream(upstream *Upstream) {
// then returns a reader for the buffer along with how many bytes were buffered. Always close
// the return value when done with it, just like if it was the original body! If limit is 0
// (which it shouldn't be), this function returns its input; i.e. is a no-op, for safety.
// Otherwise, it returns bodyReadCloser, the original body will be closed and body will be nil
// if it's explicitly configured to buffer all or EOF is reached when reading.
// TODO: the error during reading is discarded if the limit is negative, should the error be propagated
// to upstream/downstream?
func (h Handler) bufferedBody(originalBody io.ReadCloser, limit int64) (io.ReadCloser, int64) {
if limit == 0 {
return originalBody, 0