github.com/go-playground/pkg/v5@v5.29.1/io/limit_reader.go (about)

     1  package ioext
     2  
     3  import (
     4  	"errors"
     5  	"io"
     6  )
     7  
     8  var (
     9  	// ErrLimitedReaderEOF is an error returned by LimitedReader to give feedback to the fact that we did not hit an
    10  	// EOF of the Reader but hit the limit imposed by the LimitedReader.
    11  	ErrLimitedReaderEOF = errors.New("LimitedReader EOF: limit reached")
    12  )
    13  
    14  // LimitReader returns a LimitedReader that reads from r
    15  // but stops with ErrLimitedReaderEOF after n bytes.
    16  func LimitReader(r io.Reader, n int64) *LimitedReader {
    17  	return &LimitedReader{R: r, N: n}
    18  }
    19  
    20  // A LimitedReader reads from R but limits the amount of
    21  // data returned to just N bytes. Each call to Read
    22  // updates N to reflect the new amount remaining.
    23  // Read returns ErrLimitedReaderEOF when N <= 0 or when the underlying R returns EOF.
    24  // Unlike the std io.LimitedReader this provides feedback
    25  // that the limit was reached through the returned error.
    26  type LimitedReader struct {
    27  	R io.Reader
    28  	N int64 // bytes allotted
    29  }
    30  
    31  func (l *LimitedReader) Read(p []byte) (n int, err error) {
    32  	if int64(len(p)) > l.N {
    33  		p = p[0 : l.N+1]
    34  	}
    35  	n, err = l.R.Read(p)
    36  	l.N -= int64(n)
    37  	if err != nil {
    38  		return
    39  	}
    40  	if l.N < 0 {
    41  		return n, ErrLimitedReaderEOF
    42  	}
    43  	return
    44  }