git.gammaspectra.live/P2Pool/consensus/v3@v3.8.0/utils/limitreader.go (about) 1 package utils 2 3 import "io" 4 5 type ReaderAndByteReader interface { 6 io.Reader 7 io.ByteReader 8 } 9 10 // LimitByteReader returns a Reader that reads from r 11 // but stops with EOF after n bytes. 12 // The underlying implementation is a *LimitedReader. 13 func LimitByteReader(r ReaderAndByteReader, n int64) *LimitedByteReader { 14 return &LimitedByteReader{r, n} 15 } 16 17 // A LimitedByteReader reads from R but limits the amount of 18 // data returned to just N bytes. Each call to Read 19 // updates N to reflect the new amount remaining. 20 // Read returns EOF when N <= 0 or when the underlying R returns EOF. 21 type LimitedByteReader struct { 22 R ReaderAndByteReader // underlying reader 23 N int64 // max bytes remaining 24 } 25 26 func (l *LimitedByteReader) Read(p []byte) (n int, err error) { 27 if l.N <= 0 { 28 return 0, io.EOF 29 } 30 if int64(len(p)) > l.N { 31 p = p[0:l.N] 32 } 33 n, err = l.R.Read(p) 34 l.N -= int64(n) 35 return 36 } 37 38 func (l *LimitedByteReader) ReadByte() (v byte, err error) { 39 if l.N <= 0 { 40 return 0, io.EOF 41 } 42 v, err = l.R.ReadByte() 43 l.N-- 44 return 45 } 46 47 func (l *LimitedByteReader) Left() int64 { 48 return l.N 49 }