From a874af9c122dc1272616ba7ade389dcec617ee0f Mon Sep 17 00:00:00 2001
From: Anton Petruhin <anton.petruhin@gmail.com>
Date: Mon, 11 Sep 2023 13:34:06 +0300
Subject: [PATCH] sleep between unsuccessful `proxy` attempts

---
 connect.go      | 20 ++++++++++----------
 connect_test.go |  7 +++----
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/connect.go b/connect.go
index 3b61ac7..ab66957 100644
--- a/connect.go
+++ b/connect.go
@@ -59,16 +59,18 @@ func (t *Tunnel) keepConnected(ctx context.Context) {
 			return
 		default:
 			t.gwConn, err = connect(t.address, t.serverName, t.token, t.config)
+			if err == nil {
+				err = proxy(ctx, t.gwConn)
+				_ = t.gwConn.Close()
+			}
 			if err != nil {
-				d := b.Duration()
 				klog.Errorln(err)
-				klog.Errorf("reconnecting to %s in %.0fs", t.address, d.Seconds())
+				d := b.Duration()
+				klog.Infof("reconnecting to %s in %.0fs", t.address, d.Seconds())
 				time.Sleep(d)
 				continue
 			}
 			b.Reset()
-			proxy(ctx, t.gwConn)
-			_ = t.gwConn.Close()
 		}
 	}
 }
@@ -220,25 +222,23 @@ func connect(gwAddr, serverName, token string, config []byte) (net.Conn, error)
 	return gwConn, nil
 }
 
-func proxy(ctx context.Context, gwConn net.Conn) {
+func proxy(ctx context.Context, gwConn net.Conn) error {
 	cfg := yamux.DefaultConfig()
 	cfg.KeepAliveInterval = time.Second
 	cfg.LogOutput = io.Discard
 	session, err := yamux.Server(gwConn, cfg)
 	if err != nil {
-		klog.Errorln("failed to start a TCP multiplexing server:", err)
-		return
+		return fmt.Errorf("failed to start a TCP multiplexing server: %s", err)
 	}
 	defer session.Close()
 	for {
 		select {
 		case <-ctx.Done():
-			return
+			return nil
 		default:
 			gwStream, err := session.Accept()
 			if err != nil {
-				klog.Errorf("failed to accept a stream: %s", err)
-				return
+				return fmt.Errorf("failed to accept a stream: %s", err)
 			}
 			go func(c net.Conn) {
 				defer c.Close()
diff --git a/connect_test.go b/connect_test.go
index 27590cc..8722389 100644
--- a/connect_test.go
+++ b/connect_test.go
@@ -10,7 +10,6 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"io"
-	"io/ioutil"
 	"net"
 	"net/http"
 	"net/http/httptest"
@@ -83,7 +82,7 @@ func TestProxy(t *testing.T) {
 
 	gwConn, err := connect(addr, "", token, []byte("config_data"))
 	go func() {
-		proxy(context.Background(), gwConn)
+		require.NoError(t, proxy(context.Background(), gwConn))
 	}()
 
 	session := <-sessionChan
@@ -107,7 +106,7 @@ func TestProxy(t *testing.T) {
 
 	res, err := client.Get("http://any/-/healthy")
 	require.NoError(t, err)
-	data, err := ioutil.ReadAll(res.Body)
+	data, err := io.ReadAll(res.Body)
 	require.NoError(t, err)
 	res.Body.Close()
 	assert.Equal(t, "Prometheus is Healthy.", string(data))
@@ -131,7 +130,7 @@ func TestProxy(t *testing.T) {
 
 	res, err = client.Get("http://any/-/healthy")
 	require.NoError(t, err)
-	data, err = ioutil.ReadAll(res.Body)
+	data, err = io.ReadAll(res.Body)
 	require.NoError(t, err)
 	res.Body.Close()
 	assert.Equal(t, "Pyroscope is Healthy.", string(data))