From 4025f223c4753a0840a95b0dd91e6edd1f0521ee Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Wed, 3 Mar 2021 17:18:16 +0300
Subject: [PATCH] Correct rewinding for faketls

---
 wrappers/stream/rewind.go | 46 ++++++++++-----------------------------
 1 file changed, 12 insertions(+), 34 deletions(-)

diff --git a/wrappers/stream/rewind.go b/wrappers/stream/rewind.go
index b0e6afb..d8d7425 100644
--- a/wrappers/stream/rewind.go
+++ b/wrappers/stream/rewind.go
@@ -2,7 +2,6 @@ package stream
 
 import (
 	"bytes"
-	"errors"
 	"io"
 	"net"
 	"sync"
@@ -18,10 +17,10 @@ type ReadWriteCloseRewinder interface {
 }
 
 type wrapperRewind struct {
-	parent   conntypes.StreamReadWriteCloser
-	buf      bytes.Buffer
-	mutex    sync.Mutex
-	rewinded bool
+	parent       conntypes.StreamReadWriteCloser
+	activeReader io.Reader
+	buf          bytes.Buffer
+	mutex        sync.Mutex
 }
 
 func (w *wrapperRewind) Write(p []byte) (int, error) {
@@ -36,38 +35,14 @@ func (w *wrapperRewind) Read(p []byte) (int, error) {
 	w.mutex.Lock()
 	defer w.mutex.Unlock()
 
-	if w.rewinded {
-		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
+	return w.activeReader.Read(p)
 }
 
-func (w *wrapperRewind) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+func (w *wrapperRewind) ReadTimeout(p []byte, _ time.Duration) (int, error) {
 	w.mutex.Lock()
 	defer w.mutex.Unlock()
 
-	if w.rewinded {
-		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
+	return w.activeReader.Read(p)
 }
 
 func (w *wrapperRewind) Conn() net.Conn {
@@ -94,12 +69,15 @@ func (w *wrapperRewind) Close() error {
 
 func (w *wrapperRewind) Rewind() {
 	w.mutex.Lock()
-	w.rewinded = true
+	w.activeReader = io.MultiReader(&w.buf, w.parent)
 	w.mutex.Unlock()
 }
 
 func NewRewind(parent conntypes.StreamReadWriteCloser) ReadWriteCloseRewinder {
-	return &wrapperRewind{
+	rv := &wrapperRewind{
 		parent: parent,
 	}
+	rv.activeReader = io.TeeReader(parent, &rv.buf)
+
+	return rv
 }