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 }