github.com/insolar/vanilla@v0.0.0-20201023172447-248fdf805322/iokit/limited.go (about)

     1  // Copyright 2020 Insolar Network Ltd.
     2  // All rights reserved.
     3  // This material is licensed under the Insolar License version 1.0,
     4  // available at https://github.com/insolar/assured-ledger/blob/master/LICENSE.md.
     5  
     6  package iokit
     7  
     8  import "io"
     9  
    10  type LenReader interface {
    11  	io.Reader
    12  	Len() int
    13  }
    14  
    15  func LimitWriter(w io.Writer, n int64) *LimitedWriter {
    16  	return &LimitedWriter{TeeWriter{W: w}, n}
    17  }
    18  
    19  func LimitReader(r io.Reader, n int64) *LimitedReader {
    20  	return &LimitedReader{TeeReader{R: r}, n}
    21  }
    22  
    23  func NewLimitedTeeWriter(main, copy io.Writer, limit int64) *LimitedWriter {
    24  	return &LimitedWriter{TeeWriter{main, teeTemplate{copy, 0}}, limit}
    25  }
    26  
    27  func NewLimitedTeeWriterWithSkip(main, copy io.Writer, skipLeadingBytes int, limit int64) *LimitedWriter {
    28  	return &LimitedWriter{TeeWriter{main, teeTemplate{copy, skipLeadingBytes}}, limit}
    29  }
    30  
    31  func NewLimitedTeeReader(main io.Reader, copy io.Writer, limit int64) *LimitedReader {
    32  	return &LimitedReader{TeeReader{main, teeTemplate{copy, 0}}, limit}
    33  }
    34  
    35  func NewLimitedTeeReaderWithSkip(main io.Reader, copy io.Writer, skipLeadingBytes int, limit int64) *LimitedReader {
    36  	return &LimitedReader{TeeReader{main, teeTemplate{copy, skipLeadingBytes}}, limit}
    37  }
    38  
    39  /***********************************************/
    40  
    41  var _ io.Reader = &LimitedReader{}
    42  
    43  type LimitedReader struct {
    44  	r TeeReader
    45  	n int64 // max bytes remaining
    46  }
    47  
    48  func (l *LimitedReader) Read(p []byte) (n int, err error) {
    49  	if l.n <= 0 {
    50  		return 0, io.ErrUnexpectedEOF
    51  	}
    52  	if int64(len(p)) > l.n {
    53  		p = p[0:l.n]
    54  	}
    55  	n, err = l.r.Read(p)
    56  	l.n -= int64(n)
    57  	return
    58  }
    59  
    60  func (l *LimitedReader) RemainingBytes() int64 {
    61  	return l.n
    62  }
    63  
    64  func (l *LimitedReader) Len() int {
    65  	return int(l.n)
    66  }
    67  
    68  /***********************************************/
    69  
    70  var _ io.Writer = &LimitedWriter{}
    71  
    72  type LimitedWriter struct {
    73  	w TeeWriter
    74  	n int64 // max bytes remaining
    75  }
    76  
    77  func (l *LimitedWriter) Write(p []byte) (n int, err error) {
    78  	if l.n <= 0 {
    79  		return 0, io.ErrShortWrite
    80  	}
    81  	if int64(len(p)) > l.n {
    82  		p = p[0:l.n]
    83  	}
    84  	n, err = l.w.Write(p)
    85  	l.n -= int64(n)
    86  	return
    87  }
    88  
    89  func (l *LimitedWriter) RemainingBytes() int64 {
    90  	return l.n
    91  }