From a5afd3241dbb494315de7b29b5189e14e0581b4c Mon Sep 17 00:00:00 2001
From: Ben Waples <bwaples@specterops.io>
Date: Tue, 2 Apr 2024 13:51:41 -0700
Subject: [PATCH] fix: requestCompleted logic and comments

---
 client/rest/client_test.go | 24 ++++++++++--------------
 client/rest/utils.go       |  2 ++
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/client/rest/client_test.go b/client/rest/client_test.go
index ca8e4c5..85dbb3b 100644
--- a/client/rest/client_test.go
+++ b/client/rest/client_test.go
@@ -18,7 +18,6 @@
 package rest
 
 import (
-	"fmt"
 	"net/http"
 	"net/http/httptest"
 
@@ -50,26 +49,23 @@ func TestClosedConnection(t *testing.T) {
 		if req, err := http.NewRequest(http.MethodGet, testServer.URL, nil); err != nil {
 			t.Fatalf("error creating request %v", err)
 		} else {
-			didSucceed := false
+			requestCompleted := false
 
-			// make request in separate goroutine so we can end it early after the first retry
+			// make request in separate goroutine so its not blocking after we validated the retry
 			go func() {
-				// end request on the second attempt after a closed connection
-				if res, err := client.Send(req); err != nil {
-					fmt.Println(err)
-				} else if res.Status == "200" {
-					didSucceed = true
-				}
+				client.Send(req)
+				requestCompleted = true
 			}()
 
-			for attempt < 3 {
-				if attempt > 1 {
-					return
+			// block until attempt is > 2 or request succeeds
+			for attempt <= 2 {
+				if attempt > 1 || requestCompleted {
+					break
 				}
 			}
 
-			if didSucceed {
-				t.Fatalf("expected an attempted retry but the request succeeded")
+			if requestCompleted {
+				t.Fatalf("expected an attempted retry but the request completed")
 			}
 		}
 	}
diff --git a/client/rest/utils.go b/client/rest/utils.go
index 941a688..e183156 100644
--- a/client/rest/utils.go
+++ b/client/rest/utils.go
@@ -126,6 +126,8 @@ var ClosedConnectionMsg = "An existing connection was forcibly closed by the rem
 
 func IsClosedConnectionErr(err error) bool {
 	closedFromClient := strings.Contains(err.Error(), ClosedConnectionMsg)
+	// Mocking http.Do would require a larger refactor,
+	// so closedFromTestCase is used for testing only.
 	closedFromTestCase := strings.HasSuffix(err.Error(), ": EOF")
 	return closedFromClient || closedFromTestCase
 }