package routes import ( "bytes" "errors" "fmt" "io" "sour.is/x/toolbox/log" ) type PreviewReader struct { r io.Reader buf bytes.Buffer } func NewPreviewReader(r io.Reader) *PreviewReader { return &PreviewReader{r: r} } var ErrReaderDrained = errors.New("io.Reader has been drained") func (pr *PreviewReader) Read(p []byte) (n int, err error) { if pr.r == nil { return 0, ErrReaderDrained } i, err := pr.r.Read(p) log.Debugf("PreviewReader: buffer %d bytes", i) _, berr := pr.buf.Write(p[:i]) if berr != nil { return i, berr } return i, err } func (pr *PreviewReader) Drain() io.Reader { dr := &drainReader{r: pr.r, buf: &pr.buf} pr.r = nil return dr } type drainReader struct { r io.Reader buf *bytes.Buffer } var _ io.Seeker = (*drainReader)(nil) func (dr *drainReader) Read(p []byte) (n int, err error) { i := 0 if dr.buf.Len() > 0 { i, err = dr.buf.Read(p) if err != nil && err != io.EOF { return i, err } if err != nil && err == io.EOF { err = nil } if err != nil { return i, err } } ri, err := dr.r.Read(p[i:]) // log.Debugs("drainReader:", "drain", i, "read", ri, "cap", len(p), "err", err) return ri + i, err } // Seek attempt if the underlying reader supports it. func (dr *drainReader) Seek(offset int64, whence int) (int64, error) { if dr.buf.Len() > 0 { return 0, fmt.Errorf("unable to seek") } if r, ok := dr.r.(io.Seeker); ok { return r.Seek(offset, whence) } return 0, fmt.Errorf("unable to seek") }