Rewrite cloaking

faketls/client_protocol.go

@@ -8,7 +8,6 @@ import (
 	"io"
 	"net"
 	"strconv"
-	"sync"
 	"time"
 
 	"github.com/9seconds/mtg/antireplay"
@@ -110,25 +109,7 @@ func (c *ClientProtocol) cloakHost(clientConn io.ReadWriteCloser) {
 		return
 	}
 
-	defer hostConn.Close()
-
-	wg := &sync.WaitGroup{}
-	wg.Add(2)
-
-	go c.pipe(hostConn, clientConn, wg)
-
-	go c.pipe(clientConn, hostConn, wg)
-
-	wg.Wait()
-}
-
-func (c *ClientProtocol) pipe(dst io.WriteCloser, src io.Reader, wg *sync.WaitGroup) {
-	defer func() {
-		wg.Done()
-		dst.Close()
-	}()
-
-	io.Copy(dst, src) // nolint: errcheck
+	cloak(clientConn, hostConn)
 }
 
 func MakeClientProtocol() protocol.ClientProtocol {

faketls/cloak.go

@@ -0,0 +1,62 @@
+package faketls
+
+import (
+	"context"
+	"io"
+	"sync"
+	"time"
+
+	"github.com/9seconds/mtg/wrappers/rwc"
+)
+
+const cloakTimeout = 5 * time.Second
+
+func cloak(one, another io.ReadWriteCloser) {
+	defer func() {
+		one.Close()
+		another.Close()
+	}()
+
+	channelPing := make(chan struct{}, 1)
+	ctx, cancel := context.WithCancel(context.Background())
+	one = rwc.NewPing(ctx, one, channelPing)
+	another = rwc.NewPing(ctx, another, channelPing)
+	wg := &sync.WaitGroup{}
+
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		io.Copy(one, another) // nolint: errcheck
+	}()
+
+	go func() {
+		defer wg.Done()
+		io.Copy(another, one) // nolint: errcheck
+	}()
+
+	go func() {
+		wg.Wait()
+		cancel()
+	}()
+
+	go func() {
+		timer := time.NewTimer(cloakTimeout)
+		defer timer.Stop()
+
+		for {
+			select {
+			case <-channelPing:
+				timer.Stop()
+				timer = time.NewTimer(cloakTimeout)
+			case <-ctx.Done():
+				return
+			case <-timer.C:
+				cancel()
+				return
+			}
+		}
+	}()
+
+	<-ctx.Done()
+}

wrappers/rwc/ping.go

@@ -0,0 +1,48 @@
+package rwc
+
+import (
+	"context"
+	"io"
+)
+
+type wrapperPing struct {
+	parent      io.ReadWriteCloser
+	ctx         context.Context
+	channelPing chan<- struct{}
+}
+
+func (w *wrapperPing) Read(p []byte) (int, error) {
+	n, err := w.parent.Read(p)
+	if err == nil {
+		select {
+		case <-w.ctx.Done():
+		case w.channelPing <- struct{}{}:
+		}
+	}
+
+	return n, err
+}
+
+func (w *wrapperPing) Write(p []byte) (int, error) {
+	n, err := w.parent.Write(p)
+	if err == nil {
+		select {
+		case <-w.ctx.Done():
+		case w.channelPing <- struct{}{}:
+		}
+	}
+
+	return n, err
+}
+
+func (w *wrapperPing) Close() error {
+	return w.parent.Close()
+}
+
+func NewPing(ctx context.Context, parent io.ReadWriteCloser, channelPing chan<- struct{}) io.ReadWriteCloser {
+	return &wrapperPing{
+		parent:      parent,
+		ctx:         ctx,
+		channelPing: channelPing,
+	}
+}