You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
54 lines
1.1 KiB
Go
54 lines
1.1 KiB
Go
package noxy
|
|
|
|
import (
|
|
"bufio"
|
|
"errors"
|
|
"io"
|
|
)
|
|
|
|
var ErrSizeLimitExceeded = errors.New("size limit exceeded")
|
|
|
|
// HardLimitReader reads from r up to max bytes.
|
|
// it returns ErrSizeLimitExceeded if the number of read bytes is equal or exceeds max.
|
|
func HardLimitReader(r io.Reader, max int64) io.Reader {
|
|
return &hardLimitReader{r, max}
|
|
}
|
|
|
|
type hardLimitReader struct {
|
|
r io.Reader
|
|
n int64 // remaining bytes
|
|
}
|
|
|
|
func (l *hardLimitReader) Read(p []byte) (n int, err error) {
|
|
if l.n <= 0 {
|
|
return 0, ErrSizeLimitExceeded
|
|
}
|
|
if int64(len(p)) > l.n {
|
|
p = p[0:l.n]
|
|
}
|
|
n, err = l.r.Read(p)
|
|
l.n -= int64(n)
|
|
return
|
|
}
|
|
|
|
// SinkTeeReader stops writing to w at the first encountered write error
|
|
// but continues to propagate reads. the underlying writer is always buffered.
|
|
func SinkTeeReader(r io.Reader, w io.Writer) io.Reader {
|
|
return &sinkTeeReader{r: r, w: bufio.NewWriter(w)}
|
|
}
|
|
|
|
type sinkTeeReader struct {
|
|
r io.Reader
|
|
w *bufio.Writer
|
|
werr error
|
|
}
|
|
|
|
func (st *sinkTeeReader) Read(p []byte) (n int, err error) {
|
|
n, err = st.r.Read(p)
|
|
if n > 0 && st.werr == nil {
|
|
_, st.werr = st.w.Write(p[:n])
|
|
st.werr = st.w.Flush()
|
|
}
|
|
return
|
|
}
|