From 067530cdbf853a6060d7890b246acfde71559818 Mon Sep 17 00:00:00 2001
From: Michael Li <michael.li@hashicorp.com>
Date: Tue, 1 Nov 2022 15:05:22 -0400
Subject: [PATCH] fix(test): Retry websocket connection due to flaky failure
 (#2584)

We have observed several unit tests occasionally fail when trying to read a handshake due to the following error:

failed to read protobuf message: failed to get reader: received close frame: status = StatusProtocolError and
reason = "received header with unexpected rsv bits set: false:false:true"

After an investigation, it's still unclear how this happens. We are opting to just retry to mitigate the flakiness.
---
 .../server/worker_shutdown_reload_test.go     |  4 +-
 internal/tests/helper/testing_helper.go       | 79 +++++++++++++------
 2 files changed, 56 insertions(+), 27 deletions(-)

diff --git a/internal/cmd/commands/server/worker_shutdown_reload_test.go b/internal/cmd/commands/server/worker_shutdown_reload_test.go
index 014c6e59d..f43f15887 100644
--- a/internal/cmd/commands/server/worker_shutdown_reload_test.go
+++ b/internal/cmd/commands/server/worker_shutdown_reload_test.go
@@ -99,7 +99,9 @@ func TestServer_ShutdownWorker(t *testing.T) {
 	sConn.TestSendRecvAll(t)
 
 	// Shutdown the worker and close the connection, as the worker will otherwise wait for it to close.
-	sConn.Close()
+	err = sConn.Close()
+	require.NoError(err)
+
 	workerCmd.ShutdownCh <- struct{}{}
 	if <-workerCodeChan != 0 {
 		output := workerCmd.UI.(*cli.MockUi).ErrorWriter.String() + workerCmd.UI.(*cli.MockUi).OutputWriter.String()
diff --git a/internal/tests/helper/testing_helper.go b/internal/tests/helper/testing_helper.go
index 0267ee1b2..29b0563ea 100644
--- a/internal/tests/helper/testing_helper.go
+++ b/internal/tests/helper/testing_helper.go
@@ -13,6 +13,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/cenkalti/backoff/v4"
 	"github.com/hashicorp/boundary/api/targets"
 	"github.com/hashicorp/boundary/globals"
 	"github.com/hashicorp/boundary/internal/cmd/commands/connect"
@@ -107,35 +108,61 @@ func NewTestSession(
 func (s *TestSession) connect(ctx context.Context, t *testing.T) net.Conn {
 	t.Helper()
 	require := require.New(t)
-	conn, resp, err := websocket.Dial(
-		ctx,
-		fmt.Sprintf("wss://%s/v1/proxy", s.WorkerAddr),
-		&websocket.DialOptions{
-			HTTPClient: &http.Client{
-				Transport: s.Transport,
-			},
-			Subprotocols: []string{globals.TcpProxyV1},
+
+	var conn *websocket.Conn
+	var resp *http.Response
+	var handshakeResult proxy.HandshakeResult
+	// A retry was added here to mitigate some flakiness.
+	// Occasionally, `wspb.Read` would throw an error due to "received header with unexpected
+	// rsv bits set: false:false:true". It was unclear what the cause of this was, so we
+	// resorted to retrying the connection.
+	err := backoff.RetryNotify(
+		func() error {
+			var err error
+			conn, resp, err = websocket.Dial(
+				ctx,
+				fmt.Sprintf("wss://%s/v1/proxy", s.WorkerAddr),
+				&websocket.DialOptions{
+					HTTPClient: &http.Client{
+						Transport: s.Transport,
+					},
+					Subprotocols: []string{globals.TcpProxyV1},
+				},
+			)
+			if err != nil {
+				return backoff.Permanent(err)
+			}
+			if resp.Header.Get("Sec-WebSocket-Protocol") != globals.TcpProxyV1 {
+				return backoff.Permanent(errors.New(
+					fmt.Sprintf("Unexpected header. Expected: %s, Actual: %s",
+						globals.TcpProxyV1,
+						resp.Header.Get("Sec-WebSocket-Protocol"),
+					)))
+			}
+
+			// Send and receive/check the handshake.
+			if s.tofuToken == "" {
+				s.tofuToken, err = base62.Random(20)
+				if err != nil {
+					return backoff.Permanent(errors.New("Error creating token"))
+				}
+			}
+			handshake := proxy.ClientHandshake{TofuToken: s.tofuToken}
+			t.Logf("Using token: %s", s.tofuToken)
+
+			err = wspb.Write(ctx, conn, &handshake)
+			if err != nil {
+				return backoff.Permanent(err)
+			}
+
+			return wspb.Read(ctx, conn, &handshakeResult)
+		},
+		backoff.WithMaxRetries(backoff.NewConstantBackOff(time.Second), 1),
+		func(err error, td time.Duration) {
+			t.Logf("Issue with reading websocket: %s. Retrying...", err)
 		},
 	)
 	require.NoError(err)
-	require.NotNil(conn)
-	require.NotNil(resp)
-	require.Equal(resp.Header.Get("Sec-WebSocket-Protocol"), globals.TcpProxyV1)
-
-	// Send the handshake.
-	if s.tofuToken == "" {
-		s.tofuToken, err = base62.Random(20)
-	}
-
-	require.NoError(err)
-	handshake := proxy.ClientHandshake{TofuToken: s.tofuToken}
-	err = wspb.Write(ctx, conn, &handshake)
-	require.NoError(err)
-
-	// Receive/check the handshake
-	var handshakeResult proxy.HandshakeResult
-	err = wspb.Read(ctx, conn, &handshakeResult)
-	require.NoError(err)
 
 	// This is just a cursory check to make sure that the handshake is
 	// populated. We could check connections remaining too, but that