86 lines
1.5 KiB
Go
86 lines
1.5 KiB
Go
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")
|
|
}
|