Merge remote-tracking branch 'origin/master' into stable

This commit is contained in:
9seconds
2021-03-03 17:25:34 +03:00
+9 -31
View File
@@ -2,7 +2,6 @@ package stream
import ( import (
"bytes" "bytes"
"errors"
"io" "io"
"net" "net"
"sync" "sync"
@@ -19,9 +18,9 @@ type ReadWriteCloseRewinder interface {
type wrapperRewind struct { type wrapperRewind struct {
parent conntypes.StreamReadWriteCloser parent conntypes.StreamReadWriteCloser
activeReader io.Reader
buf bytes.Buffer buf bytes.Buffer
mutex sync.Mutex mutex sync.Mutex
rewinded bool
} }
func (w *wrapperRewind) Write(p []byte) (int, error) { func (w *wrapperRewind) Write(p []byte) (int, error) {
@@ -36,38 +35,14 @@ func (w *wrapperRewind) Read(p []byte) (int, error) {
w.mutex.Lock() w.mutex.Lock()
defer w.mutex.Unlock() defer w.mutex.Unlock()
if w.rewinded { return w.activeReader.Read(p)
if n, err := w.buf.Read(p); errors.Is(err, io.EOF) {
return n, err // nolint: wrapcheck
}
}
n, err := w.parent.Read(p)
if !w.rewinded {
w.buf.Write(p[:n])
}
return n, err // nolint: wrapcheck
} }
func (w *wrapperRewind) ReadTimeout(p []byte, timeout time.Duration) (int, error) { func (w *wrapperRewind) ReadTimeout(p []byte, _ time.Duration) (int, error) {
w.mutex.Lock() w.mutex.Lock()
defer w.mutex.Unlock() defer w.mutex.Unlock()
if w.rewinded { return w.activeReader.Read(p)
if n, err := w.buf.Read(p); errors.Is(err, io.EOF) {
return n, err // nolint: wrapcheck
}
}
n, err := w.parent.ReadTimeout(p, timeout)
if !w.rewinded {
w.buf.Write(p[:n])
}
return n, err // nolint: wrapcheck
} }
func (w *wrapperRewind) Conn() net.Conn { func (w *wrapperRewind) Conn() net.Conn {
@@ -94,12 +69,15 @@ func (w *wrapperRewind) Close() error {
func (w *wrapperRewind) Rewind() { func (w *wrapperRewind) Rewind() {
w.mutex.Lock() w.mutex.Lock()
w.rewinded = true w.activeReader = io.MultiReader(&w.buf, w.parent)
w.mutex.Unlock() w.mutex.Unlock()
} }
func NewRewind(parent conntypes.StreamReadWriteCloser) ReadWriteCloseRewinder { func NewRewind(parent conntypes.StreamReadWriteCloser) ReadWriteCloseRewinder {
return &wrapperRewind{ rv := &wrapperRewind{
parent: parent, parent: parent,
} }
rv.activeReader = io.TeeReader(parent, &rv.buf)
return rv
} }