From 87da922aff666918e8f46b50946790cbf99248e8 Mon Sep 17 00:00:00 2001
From: Majd Alfhaily <majd@alfhaily.me>
Date: Sat, 29 Aug 2026 21:02:00 +0200
Subject: [PATCH] fix: retry transient authentication responses

---
 pkg/appstore/appstore.go            |  4 ++
 pkg/appstore/appstore_login.go      | 55 ++++++++++++++++++++++++-
 pkg/appstore/appstore_login_test.go | 63 +++++++++++++++++++++++++++--
 3 files changed, 117 insertions(+), 5 deletions(-)

diff --git a/pkg/appstore/appstore.go b/pkg/appstore/appstore.go
index 33d66bd7..6acadfed 100644
--- a/pkg/appstore/appstore.go
+++ b/pkg/appstore/appstore.go
@@ -1,6 +1,8 @@
 package appstore
 
 import (
+	"time"
+
 	"github.com/majd/ipatool/v2/pkg/http"
 	"github.com/majd/ipatool/v2/pkg/keychain"
 	"github.com/majd/ipatool/v2/pkg/util/machine"
@@ -46,6 +48,7 @@ type appstore struct {
 	ownedAppsClient     http.Client[[]byte]
 	httpClient          http.Client[interface{}]
 	actionSignerFactory ActionSignerFactory
+	authRetrySleep      func(time.Duration)
 	machine             machine.Machine
 	os                  operatingsystem.OperatingSystem
 }
@@ -79,6 +82,7 @@ func NewAppStore(args Args) AppStore {
 		ownedAppsClient:     http.NewClient[[]byte](clientArgs),
 		httpClient:          http.NewClient[interface{}](clientArgs),
 		actionSignerFactory: actionSignerFactory,
+		authRetrySleep:      time.Sleep,
 		machine:             args.Machine,
 		os:                  args.OperatingSystem,
 	}
diff --git a/pkg/appstore/appstore_login.go b/pkg/appstore/appstore_login.go
index e15fa739..d700f0f2 100644
--- a/pkg/appstore/appstore_login.go
+++ b/pkg/appstore/appstore_login.go
@@ -7,6 +7,7 @@ import (
 	gohttp "net/http"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/majd/ipatool/v2/pkg/http"
 	"github.com/majd/ipatool/v2/pkg/util"
@@ -16,6 +17,11 @@ var (
 	ErrAuthCodeRequired = errors.New("auth code is required")
 )
 
+const (
+	maxAuthenticationRequestAttempts = 3
+	authenticationRetryDelay         = 250 * time.Millisecond
+)
+
 type LoginInput struct {
 	Email    string
 	Password string
@@ -119,7 +125,7 @@ func (t *appstore) login(email, password, authCode, guid, endpoint string, signe
 
 		request := t.loginRequest(email, password, authCode, guid, endpoint, requestAttempt, signer)
 		request.URL, _ = util.IfEmpty(redirect, request.URL), ""
-		res, err = t.loginClient.Send(request)
+		res, err = t.sendAuthenticationRequest(request)
 
 		if err != nil {
 			return Account{}, fmt.Errorf("request failed: %w", err)
@@ -168,6 +174,53 @@ func (t *appstore) login(email, password, authCode, guid, endpoint string, signe
 	return acc, nil
 }
 
+func (t *appstore) sendAuthenticationRequest(request http.Request) (http.Result[loginResult], error) {
+	statuses := make([]string, 0, maxAuthenticationRequestAttempts)
+
+	sleep := t.authRetrySleep
+	if sleep == nil {
+		sleep = time.Sleep
+	}
+
+	for attempt := 1; ; attempt++ {
+		result, err := t.loginClient.Send(request)
+
+		status, retry := retryableAuthenticationError(err)
+		if !retry {
+			if err != nil {
+				return result, fmt.Errorf("%w", err)
+			}
+
+			return result, nil
+		}
+
+		statuses = append(statuses, strconv.Itoa(status))
+
+		if attempt == maxAuthenticationRequestAttempts {
+			return result, fmt.Errorf(
+				"authentication request failed after %d attempts (HTTP %s): %w",
+				maxAuthenticationRequestAttempts, strings.Join(statuses, ", "), err,
+			)
+		}
+
+		sleep(time.Duration(attempt) * authenticationRetryDelay)
+	}
+}
+
+func retryableAuthenticationError(err error) (int, bool) {
+	var responseErr *http.UnexpectedResponseError
+	if !errors.As(err, &responseErr) {
+		return 0, false
+	}
+
+	status := responseErr.StatusCode
+	retry := status == gohttp.StatusNoContent ||
+		status == gohttp.StatusNotFound ||
+		status/100 == 5
+
+	return status, retry
+}
+
 func (t *appstore) parseLoginResponse(res *http.Result[loginResult], attempt int, authCode string) (bool, string, error) {
 	var (
 		retry    bool
diff --git a/pkg/appstore/appstore_login_test.go b/pkg/appstore/appstore_login_test.go
index 29d91ec5..a5d24fdb 100644
--- a/pkg/appstore/appstore_login_test.go
+++ b/pkg/appstore/appstore_login_test.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/majd/ipatool/v2/pkg/http"
 	"github.com/majd/ipatool/v2/pkg/keychain"
@@ -41,10 +42,11 @@ var _ = Describe("AppStore (Login)", func() {
 		mockMachine = machine.NewMockMachine(ctrl)
 		signer = &stubActionSigner{}
 		as = &appstore{
-			keychain:    mockKeychain,
-			loginClient: mockClient,
-			bagClient:   mockBagClient,
-			machine:     mockMachine,
+			keychain:       mockKeychain,
+			loginClient:    mockClient,
+			bagClient:      mockBagClient,
+			machine:        mockMachine,
+			authRetrySleep: func(time.Duration) {},
 			actionSignerFactory: func(config SAPConfig, machineID []byte) (ActionSigner, error) {
 				Expect(config).To(Equal(validSAPConfig()))
 				Expect(machineID).To(Equal([]byte{0, 0, 0, 0, 0, 0}))
@@ -58,6 +60,59 @@ var _ = Describe("AppStore (Login)", func() {
 		ctrl.Finish()
 	})
 
+	Describe("transient authentication responses", func() {
+		It("retries the same request until it succeeds", func() {
+			request := as.loginRequest(testEmail, testPassword, "", "guid", testAuthEndpoint, 1, signer)
+			responses := []struct {
+				result http.Result[loginResult]
+				err    error
+			}{
+				{err: &http.UnexpectedResponseError{StatusCode: 204}},
+				{err: &http.UnexpectedResponseError{StatusCode: 404}},
+				{result: http.Result[loginResult]{StatusCode: 200}},
+			}
+			call := 0
+			mockClient.EXPECT().
+				Send(gomock.Any()).
+				DoAndReturn(func(actual http.Request) (http.Result[loginResult], error) {
+					Expect(actual.URL).To(Equal(request.URL))
+					Expect(actual.Payload.(*http.XMLPayload).Content).To(HaveKeyWithValue("attempt", "1"))
+					response := responses[call]
+					call++
+
+					return response.result, response.err
+				}).
+				Times(3)
+
+			result, err := as.sendAuthenticationRequest(request)
+
+			Expect(err).ToNot(HaveOccurred())
+			Expect(result.StatusCode).To(Equal(200))
+		})
+
+		DescribeTable("bounds retries to transient statuses",
+			func(status, expectedCalls int) {
+				request := as.loginRequest(testEmail, testPassword, "", "guid", testAuthEndpoint, 1, signer)
+				responseErr := &http.UnexpectedResponseError{StatusCode: status}
+				mockClient.EXPECT().
+					Send(gomock.Any()).
+					Return(http.Result[loginResult]{}, responseErr).
+					Times(expectedCalls)
+
+				_, err := as.sendAuthenticationRequest(request)
+
+				Expect(errors.Is(err, responseErr)).To(BeTrue())
+				if expectedCalls == maxAuthenticationRequestAttempts {
+					Expect(err.Error()).To(ContainSubstring("after 3 attempts"))
+					Expect(err.Error()).To(ContainSubstring(fmt.Sprintf("HTTP %d, %d, %d", status, status, status)))
+				}
+			},
+			Entry("retries and then stops for HTTP 204", 204, maxAuthenticationRequestAttempts),
+			Entry("retries and then stops for HTTP 503", 503, maxAuthenticationRequestAttempts),
+			Entry("does not retry HTTP 403", 403, 1),
+		)
+	})
+
 	When("fails to read Machine's MAC address", func() {
 		BeforeEach(func() {
 			mockMachine.EXPECT().