github.com/cozy/cozy-stack@v0.0.0-20240603063001-31110fa4cae1/pkg/utils/limit_writer.go (about)

     1  package utils
     2  
     3  import "io"
     4  
     5  type limitedWriter struct {
     6  	w io.Writer
     7  	n int64
     8  
     9  	discard bool
    10  }
    11  
    12  func (l *limitedWriter) Write(p []byte) (n int, err error) {
    13  	if l.n <= 0 {
    14  		if l.discard {
    15  			return len(p), nil
    16  		}
    17  		return 0, io.ErrShortWrite
    18  	}
    19  	var discarded bool
    20  	if int64(len(p)) > l.n {
    21  		p = p[0:l.n]
    22  		if l.discard {
    23  			discarded = true
    24  		} else {
    25  			err = io.ErrShortWrite
    26  		}
    27  	}
    28  	n, errw := l.w.Write(p)
    29  	if errw != nil {
    30  		err = errw
    31  	}
    32  	l.n -= int64(n)
    33  	if discarded {
    34  		n = len(p)
    35  	}
    36  	return n, err
    37  }
    38  
    39  // LimitWriter works like io.LimitReader. It writes at most n bytes to the
    40  // underlying Writer. It returns io.ErrShortWrite if more than n bytes are
    41  // attempted to be written.
    42  func LimitWriter(w io.Writer, n int64) io.Writer {
    43  	return &limitedWriter{w, n, false}
    44  }
    45  
    46  // LimitWriterDiscard works like io.LimitReader. It writes at most n bytes to
    47  // the underlying Writer. It does not return any error if more than n bytes are
    48  // attempted to be written, the data is discarded.
    49  func LimitWriterDiscard(w io.Writer, n int64) io.Writer {
    50  	return &limitedWriter{w, n, true}
    51  }