test: TestClosedConnection

Author: benwaples

client/rest/client.go

@@ -255,10 +255,9 @@ func (s *restClient) send(req *http.Request) (*http.Response, error) {
 
 			// Try the request
 			if res, err = s.http.Do(req); err != nil {
-				// TODO: maybe add a test case for this?
-				if strings.Contains(err.Error(), "An existing connection was forcibly closed by the remote host.") {
-
-					fmt.Printf("remote host force closed connection while requesting %s; attempt %d/%d; trying again", req.URL, retry+1, maxRetries)
+				closedConnectionMsg := "An existing connection was forcibly closed by the remote host."
+				if strings.Contains(err.Error(), closedConnectionMsg) || strings.HasSuffix(err.Error(), ": EOF") {
+					fmt.Printf("remote host force closed connection while requesting %s; attempt %d/%d; trying again\n", req.URL, retry+1, maxRetries)
 					backoff := math.Pow(5, float64(retry+1))
 					time.Sleep(time.Second * time.Duration(backoff))
 					continue

client/rest/client_test.go

@@ -18,22 +18,59 @@
 package rest
 
 import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+
 	"testing"
+
+	"github.com/bloodhoundad/azurehound/v2/client/config"
 )
 
 func TestClosedConnection(t *testing.T) {
-	// var s *httptest.Server
-	// var h http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
-	// 	fmt.Println("closing client connections")
-	// 	s.CloseClientConnections()
-	// }
-	// s = httptest.NewServer(h)
-	// defer s.Close()
-
-	// res, err := RestClient.Get(gomock.Any(), s.URL)
-	// if err == nil {
-	// 	t.Fatalf("Something aint right, err should be nil %v", err)
-	// }
-
-	// fmt.Printf("res:%v,err:%v", res, err)
+	attempt := 0
+	var testServer *httptest.Server
+	var mockHandler http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
+		attempt++
+		testServer.CloseClientConnections()
+	}
+
+	testServer = httptest.NewServer(mockHandler)
+	defer testServer.Close()
+
+	defaultConfig := config.Config{
+		Username:  "azurehound",
+		Password:  "we_collect",
+		Authority: testServer.URL,
+	}
+
+	if client, err := NewRestClient(testServer.URL, defaultConfig); err != nil {
+		t.Fatalf("error initializing rest client %v", err)
+	} else {
+		if req, err := http.NewRequest(http.MethodGet, testServer.URL, nil); err != nil {
+			t.Fatalf("error creating request %v", err)
+		} else {
+			didSucceed := false
+
+			// make request in separate goroutine so we can end it early after the first 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
+				}
+			}()
+
+			for attempt < 3 {
+				if attempt > 1 {
+					return
+				}
+			}
+
+			if didSucceed {
+				t.Fatalf("expected an attempted retry but the request succeeded")
+			}
+		}
+	}
 }