diff --git a/pkg/rewindablereader/reader.go b/pkg/rewindablereader/reader.go index c3563e90..52e08c3d 100644 --- a/pkg/rewindablereader/reader.go +++ b/pkg/rewindablereader/reader.go @@ -6,25 +6,30 @@ import ( "io" ) -const ( - maxRecordedSize = 1 * 1024 * 1024 -) - // Reader is a reader that can be (and must be) rewinded once. type Reader struct { - R io.Reader + R io.Reader + MaxSize int - entries [][]byte - size int - rewinded bool + initialized bool + entries [][]byte + size int + rewinded bool } // Read implements io.Reader. func (r *Reader) Read(p []byte) (int, error) { + if !r.initialized { + r.initialized = true + if r.MaxSize == 0 { + r.MaxSize = 1 * 1024 * 1024 + } + } + if !r.rewinded { n, err := r.R.Read(p) - if (r.size + n) > maxRecordedSize { + if (r.size + n) > r.MaxSize { return 0, errors.New("max recorded size exceeded") }