mirror of https://github.com/coder/coder.git
44 lines
749 B
Go
44 lines
749 B
Go
package xio
|
|
|
|
import (
|
|
"io"
|
|
|
|
"golang.org/x/xerrors"
|
|
)
|
|
|
|
var ErrLimitReached = xerrors.Errorf("i/o limit reached")
|
|
|
|
// LimitWriter will only write bytes to the underlying writer until the limit is reached.
|
|
type LimitWriter struct {
|
|
Limit int64
|
|
N int64
|
|
W io.Writer
|
|
}
|
|
|
|
func NewLimitWriter(w io.Writer, n int64) *LimitWriter {
|
|
// If anyone tries this, just make a 0 writer.
|
|
if n < 0 {
|
|
n = 0
|
|
}
|
|
return &LimitWriter{
|
|
Limit: n,
|
|
N: 0,
|
|
W: w,
|
|
}
|
|
}
|
|
|
|
func (l *LimitWriter) Write(p []byte) (int, error) {
|
|
if l.N >= l.Limit {
|
|
return 0, ErrLimitReached
|
|
}
|
|
|
|
// Write 0 bytes if the limit is to be exceeded.
|
|
if int64(len(p)) > l.Limit-l.N {
|
|
return 0, ErrLimitReached
|
|
}
|
|
|
|
n, err := l.W.Write(p)
|
|
l.N += int64(n)
|
|
return n, err
|
|
}
|