noxy/io.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
}