32 lines
612 B
Go
32 lines
612 B
Go
|
package limio
|
||
|
|
|
||
|
|
import "io"
|
||
|
|
|
||
|
|
type limitedWriter struct {
|
||
|
|
W io.Writer
|
||
|
|
N int64
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l *limitedWriter) Write(p []byte) (n int, err error) {
|
||
|
|
if l.N <= 0 {
|
||
|
|
return 0, io.ErrShortWrite
|
||
|
|
}
|
||
|
|
if int64(len(p)) > l.N {
|
||
|
|
p = p[0:l.N]
|
||
|
|
err = io.ErrShortWrite
|
||
|
|
}
|
||
|
|
n, er := l.W.Write(p)
|
||
|
|
if er != nil {
|
||
|
|
err = er
|
||
|
|
}
|
||
|
|
l.N -= int64(n)
|
||
|
|
return n, err
|
||
|
|
}
|
||
|
|
|
||
|
|
// LimitWriter works like io.LimitReader. It writes at most n bytes
|
||
|
|
// to the underlying Writer. It returns io.ErrShortWrite if more than n
|
||
|
|
// bytes are attempted to be written.
|
||
|
|
func LimitWriter(w io.Writer, n int64) io.Writer {
|
||
|
|
return &limitedWriter{W: w, N: n}
|
||
|
|
}
|