From 0069222f74ac6df278717cb6c032c8e4136346aa Mon Sep 17 00:00:00 2001
From: Abhishek Pandey <abhishek@alcion.ai>
Date: Fri, 5 Jan 2024 19:25:04 -0800
Subject: [PATCH] Retry 400 from resp

---
 src/pkg/services/m365/api/graph/middleware.go | 26 +++++++++++++------
 .../m365/api/graph/middleware_test.go         | 20 +++++++-------
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/src/pkg/services/m365/api/graph/middleware.go b/src/pkg/services/m365/api/graph/middleware.go
index 8b23e04dd..31888a627 100644
--- a/src/pkg/services/m365/api/graph/middleware.go
+++ b/src/pkg/services/m365/api/graph/middleware.go
@@ -161,8 +161,7 @@ func (mw RetryMiddleware) Intercept(
 
 	retriable := IsErrTimeout(err) ||
 		IsErrConnectionReset(err) ||
-		mw.isRetriableRespCode(ctx, resp) ||
-		IsErrInvalidRequest(err, resp)
+		mw.isRetriableRespCode(ctx, resp)
 
 	if !retriable {
 		return resp, stackReq(ctx, req, resp, err).OrNil()
@@ -208,11 +207,7 @@ func (mw RetryMiddleware) retryRequest(
 	// 3, the request method is retriable.
 	// 4, we haven't already hit maximum retries.
 
-	if resp != nil && resp.StatusCode == http.StatusBadRequest && IsErrInvalidRequest(priorErr, resp) {
-		logger.Ctx(ctx).Infof("retrying 400_invalid_request: %s", getRespDump(ctx, resp, true))
-	}
-
-	shouldRetry := (priorErr != nil || mw.isRetriableRespCode(ctx, resp) || IsErrInvalidRequest(priorErr, resp)) &&
+	shouldRetry := (priorErr != nil || mw.isRetriableRespCode(ctx, resp)) &&
 		mw.isRetriableRequest(req) &&
 		executionCount < mw.MaxRetries
 
@@ -220,6 +215,10 @@ func (mw RetryMiddleware) retryRequest(
 		return resp, stackReq(ctx, req, resp, priorErr).OrNil()
 	}
 
+	if resp != nil && resp.StatusCode == http.StatusBadRequest {
+		logger.Ctx(ctx).Info("retrying 400_invalid_request")
+	}
+
 	executionCount++
 
 	delay := mw.getRetryDelay(req, resp, exponentialBackoff)
@@ -273,6 +272,10 @@ var retryableRespCodes = []int{
 	http.StatusBadGateway,
 }
 
+const (
+	invalidRequestMsg = "invalidRequest"
+)
+
 func (mw RetryMiddleware) isRetriableRespCode(ctx context.Context, resp *http.Response) bool {
 	if resp == nil {
 		return false
@@ -290,9 +293,16 @@ func (mw RetryMiddleware) isRetriableRespCode(ctx context.Context, resp *http.Re
 
 	// not a status code, but the message itself might indicate a connectivity issue that
 	// can be retried independent of the status code.
-	return strings.Contains(
+	first := strings.Contains(
 		strings.ToLower(getRespDump(ctx, resp, true)),
 		strings.ToLower(string(IOErrDuringRead)))
+
+	second := resp.StatusCode == http.StatusBadRequest &&
+		strings.Contains(
+			strings.ToLower(getRespDump(ctx, resp, true)),
+			strings.ToLower(invalidRequestMsg))
+
+	return first || second
 }
 
 func (mw RetryMiddleware) isRetriableRequest(req *http.Request) bool {
diff --git a/src/pkg/services/m365/api/graph/middleware_test.go b/src/pkg/services/m365/api/graph/middleware_test.go
index 8c7bd0dc8..97a05febd 100644
--- a/src/pkg/services/m365/api/graph/middleware_test.go
+++ b/src/pkg/services/m365/api/graph/middleware_test.go
@@ -44,7 +44,8 @@ func newMWReturns(code int, body []byte, err error) mwReturns {
 
 	resp := &http.Response{
 		StatusCode: code,
-		Body:       brc,
+		//Status:     "",
+		Body: brc,
 	}
 
 	if code == 0 {
@@ -167,13 +168,13 @@ func (suite *RetryMWIntgSuite) TestRetryMiddleware_Intercept_byStatusCode() {
 		mw               testMW
 		expectErr        assert.ErrorAssertionFunc
 	}{
-		{
-			name:             "200, no retries",
-			status:           http.StatusOK,
-			providedErr:      nil,
-			expectRetryCount: 0,
-			expectErr:        assert.NoError,
-		},
+		// {
+		// 	name:             "200, no retries",
+		// 	status:           http.StatusOK,
+		// 	providedErr:      nil,
+		// 	expectRetryCount: 0,
+		// 	expectErr:        assert.NoError,
+		// },
 		{
 			name:             "400, no retries",
 			status:           http.StatusBadRequest,
@@ -252,10 +253,11 @@ func (suite *RetryMWIntgSuite) TestRetryMiddleware_Intercept_byStatusCode() {
 			ctx, flush := tester.NewContext(t)
 			defer flush()
 
+			body := "HTTP/1.1 400 Bad Request\r\nTransfer-Encoding: chunked\r\nCache-Control: max-age=0, private\r\nClient-Request-Id: 9546c859-8a3d-446f-8606-9d85bc3de41d\r\nContent-Type: application/json\r\nDate: Thu, 07 Dec 2023 10:06:17 GMT\r\nRequest-Id: ed7106b8-9ee7-42ff-907a-c3d2df29cf55\r\nStrict-Transport-Security: max-age=31536000\r\nVary: Accept-Encoding\r\nX-Ms-Ags-Diagnostic: {\"ServerInfo\":{\"DataCenter\":\"East US\",\"Slice\":\"E\",\"Ring\":\"5\",\"ScaleUnit\":\"004\",\"RoleInstance\":\"BL02EPF00007CD5\"}}\r\n\r\nda\r\n{\"error\":{\"code\":\"invalidRequest\",\"message\":\"Invalid request\",\"innerError\":{\"date\":\"2023-12-07T10:06:17\",\"request-id\":\"ed7106b8-9ee7-42ff-907a-c3d2df29cf55\",\"client-request-id\":\"9546c859-8a3d-446f-8606-9d85bc3de41d\"}}}\r\n0\r\n\r\n"
 			called := 0
 			mw := newTestMW(
 				func(*http.Request) { called++ },
-				newMWReturns(test.status, nil, test.providedErr))
+				newMWReturns(test.status, []byte(body), test.providedErr))
 			mw.repeatReturn0 = true
 
 			// Add a large timeout of 100 seconds to ensure that the ctx deadline