From 5ea213e8756d49455e79e3df7429680ce17f772b Mon Sep 17 00:00:00 2001
From: Roberto Dip <dip.jesusr@gmail.com>
Date: Tue, 16 Jul 2024 13:16:00 -0300
Subject: [PATCH] improve VPP API error handling (#20446)

reading the [docs][1] I realized we're missing some recommendations for
error management.

the docs also note that certain operations like assignments happen
asynchronously and you must subscribe to events to get those errors.

this part wasn't estimated nor considered.

[1]:
https://developer.apple.com/documentation/devicemanagement/app_and_book_management/handling_error_responses

# Checklist for submitter

If some of the following don't apply, delete the relevant line.

<!-- Note that API documentation changes are now addressed by the
product design team. -->

- [x] Added/updated tests
---
 server/mdm/apple/vpp/api.go      | 37 ++++++++++++++++++-
 server/mdm/apple/vpp/api_test.go | 63 ++++++++++++++++++++++++++++++++
 2 files changed, 99 insertions(+), 1 deletion(-)

diff --git a/server/mdm/apple/vpp/api.go b/server/mdm/apple/vpp/api.go
index de4dbe4fe6..80d48808c6 100644
--- a/server/mdm/apple/vpp/api.go
+++ b/server/mdm/apple/vpp/api.go
@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/fleetdm/fleet/v4/pkg/fleethttp"
+	"github.com/fleetdm/fleet/v4/pkg/retry"
 )
 
 // Asset is a product in the store.
@@ -196,11 +197,45 @@ func do[T any](req *http.Request, token string, dest *T) error {
 		return fmt.Errorf("reading response body from Apple VPP endpoint: %w", err)
 	}
 
+	// For HTTP 5xx server error responses, a Retry-After header indicates
+	// how long the client must wait before making additional requests.
+	//
+	// https://developer.apple.com/documentation/devicemanagement/app_and_book_management/handling_error_responses#3742679
+	retryAfter := resp.Header.Get("Retry-After")
+	if resp.StatusCode == http.StatusInternalServerError && retryAfter != "" {
+		seconds, err := strconv.ParseInt(retryAfter, 10, 0)
+		if err != nil {
+			return fmt.Errorf("parsing retry-after header: %w", err)
+		}
+
+		ticker := time.NewTicker(time.Duration(seconds) * time.Second)
+		defer ticker.Stop()
+		<-ticker.C
+		return do(req, token, dest)
+	}
+
 	// For some reason, Apple returns 200 OK even if you pass an invalid token in the Auth header.
 	// We will need to parse the response and check to see if it contains an error.
 	var errResp ErrorResponse
 	if err := json.Unmarshal(body, &errResp); err == nil && (errResp.ErrorMessage != "" || errResp.ErrorNumber != 0) {
-		return &errResp
+		switch errResp.ErrorNumber {
+		// 9646: There are too many requests for the current
+		// Organization and the request has been rejected, either due
+		// to high server volume or an MDM issue. Use an
+		// incremental/exponential backoff strategy to retry the
+		// request until successful.
+		//
+		// https://developer.apple.com/documentation/devicemanagement/app_and_book_management/handling_error_responses#3783126
+		case 9646:
+			return retry.Do(
+				func() error { return do(req, token, dest) },
+				retry.WithBackoffMultiplier(3),
+				retry.WithInterval(5*time.Second),
+				retry.WithMaxAttempts(3),
+			)
+		default:
+			return &errResp
+		}
 	}
 
 	if resp.StatusCode != http.StatusOK {
diff --git a/server/mdm/apple/vpp/api_test.go b/server/mdm/apple/vpp/api_test.go
index 5421826dfd..c3fda13c3d 100644
--- a/server/mdm/apple/vpp/api_test.go
+++ b/server/mdm/apple/vpp/api_test.go
@@ -8,6 +8,7 @@ import (
 	"net/http/httptest"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 )
@@ -230,6 +231,68 @@ func TestGetAssets(t *testing.T) {
 	}
 }
 
+func TestDoRetryAfter(t *testing.T) {
+	tests := []struct {
+		name      string
+		handler   http.HandlerFunc
+		wantCalls int
+		wantErr   bool
+	}{
+		{
+			name: "no retry-after header",
+			handler: func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusInternalServerError)
+				_, err := w.Write([]byte("{}"))
+				require.NoError(t, err)
+			},
+			wantCalls: 1,
+			wantErr:   true,
+		},
+		{
+			name: "invalid retry-after header",
+			handler: func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Add("Retry-After", "foo")
+				w.WriteHeader(http.StatusInternalServerError)
+				_, err := w.Write([]byte("{}"))
+				require.NoError(t, err)
+			},
+			wantCalls: 1,
+			wantErr:   true,
+		},
+		{
+			name: "three retries",
+			handler: func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Add("Retry-After", "1")
+				w.WriteHeader(http.StatusInternalServerError)
+				_, err := w.Write([]byte("{}"))
+				require.NoError(t, err)
+			},
+			wantCalls: 3,
+			wantErr:   false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var calls int
+			setupFakeServer(t, func(w http.ResponseWriter, r *http.Request) {
+				calls++
+				if calls < tt.wantCalls {
+					tt.handler(w, r)
+					return
+				}
+			})
+
+			start := time.Now()
+			req, err := http.NewRequest(http.MethodGet, os.Getenv("FLEET_DEV_VPP_URL"), nil)
+			require.NoError(t, err)
+			err = do[any](req, "test-token", nil)
+			require.NoError(t, err)
+			require.Equal(t, tt.wantCalls, calls)
+			require.WithinRange(t, time.Now(), start, start.Add(time.Duration(tt.wantCalls)*time.Second))
+		})
+	}
+}
+
 func TestGetBaseURL(t *testing.T) {
 	t.Run("Default URL", func(t *testing.T) {
 		os.Setenv("FLEET_DEV_VPP_URL", "")