From 65e374c85c32a7dd582aa1d438161663a4abc43c Mon Sep 17 00:00:00 2001
From: Sarah Gillespie <73313222+gillespi314@users.noreply.github.com>
Date: Fri, 18 Oct 2024 13:16:04 -0500
Subject: [PATCH] Fix bug in ABM token renewal flow (#22988)

---
 changes/22955-bugfix-abm-renewal | 1 +
 cmd/fleet/cron.go                | 2 +-
 ee/server/service/mdm.go         | 4 ++--
 server/mdm/apple/apple_bm.go     | 9 ++++++++-
 4 files changed, 12 insertions(+), 4 deletions(-)
 create mode 100644 changes/22955-bugfix-abm-renewal

diff --git a/changes/22955-bugfix-abm-renewal b/changes/22955-bugfix-abm-renewal
new file mode 100644
index 0000000000..b18a359755
--- /dev/null
+++ b/changes/22955-bugfix-abm-renewal
@@ -0,0 +1 @@
+- Fixed bug in ABM renewal process that caused upload of new token to fail.
diff --git a/cmd/fleet/cron.go b/cmd/fleet/cron.go
index 877166d6d2..90dc3fffa9 100644
--- a/cmd/fleet/cron.go
+++ b/cmd/fleet/cron.go
@@ -1165,7 +1165,7 @@ func appleMDMDEPSyncerJob(
 		}
 		if incompleteToken != nil {
 			logger.Log("msg", "migrated ABM token found, updating its metadata")
-			if err := apple_mdm.SetABMTokenMetadata(ctx, incompleteToken, depStorage, ds, logger); err != nil {
+			if err := apple_mdm.SetABMTokenMetadata(ctx, incompleteToken, depStorage, ds, logger, false); err != nil {
 				return ctxerr.Wrap(ctx, err, "updating migrated ABM token metadata")
 			}
 			if err := ds.SaveABMToken(ctx, incompleteToken); err != nil {
diff --git a/ee/server/service/mdm.go b/ee/server/service/mdm.go
index 8e22df5b1b..8fa7980d3f 100644
--- a/ee/server/service/mdm.go
+++ b/ee/server/service/mdm.go
@@ -1201,7 +1201,7 @@ func (svc *Service) UploadABMToken(ctx context.Context, token io.Reader) (*fleet
 		EncryptedToken: encryptedToken,
 	}
 
-	if err := apple_mdm.SetDecryptedABMTokenMetadata(ctx, tok, decryptedToken, svc.depStorage, svc.ds, svc.logger); err != nil {
+	if err := apple_mdm.SetDecryptedABMTokenMetadata(ctx, tok, decryptedToken, svc.depStorage, svc.ds, svc.logger, false); err != nil {
 		return nil, ctxerr.Wrap(ctx, err, "setting ABM token metadata")
 	}
 
@@ -1347,7 +1347,7 @@ func (svc *Service) RenewABMToken(ctx context.Context, token io.Reader, tokenID
 		return nil, ctxerr.Wrap(ctx, err, "decrypting ABM token for renewal")
 	}
 
-	if err := apple_mdm.SetDecryptedABMTokenMetadata(ctx, oldTok, decryptedToken, svc.depStorage, svc.ds, svc.logger); err != nil {
+	if err := apple_mdm.SetDecryptedABMTokenMetadata(ctx, oldTok, decryptedToken, svc.depStorage, svc.ds, svc.logger, true); err != nil {
 		return nil, ctxerr.Wrap(ctx, err, "setting ABM token metadata")
 	}
 
diff --git a/server/mdm/apple/apple_bm.go b/server/mdm/apple/apple_bm.go
index eeb2eeddff..d6d0d5c504 100644
--- a/server/mdm/apple/apple_bm.go
+++ b/server/mdm/apple/apple_bm.go
@@ -24,13 +24,14 @@ func SetABMTokenMetadata(
 	depStorage storage.AllDEPStorage,
 	ds fleet.Datastore,
 	logger kitlog.Logger,
+	renewal bool,
 ) error {
 	decryptedToken, err := assets.ABMToken(ctx, ds, abmToken.OrganizationName)
 	if err != nil {
 		return ctxerr.Wrap(ctx, err, "getting ABM token")
 	}
 
-	return SetDecryptedABMTokenMetadata(ctx, abmToken, decryptedToken, depStorage, ds, logger)
+	return SetDecryptedABMTokenMetadata(ctx, abmToken, decryptedToken, depStorage, ds, logger, renewal)
 }
 
 const UnsavedABMTokenOrgName = "new_abm_token" //nolint:gosec
@@ -42,6 +43,7 @@ func SetDecryptedABMTokenMetadata(
 	depStorage storage.AllDEPStorage,
 	ds fleet.Datastore,
 	logger kitlog.Logger,
+	renewal bool,
 ) error {
 	depClient := NewDEPClient(depStorage, ds, logger)
 
@@ -56,6 +58,11 @@ func SetDecryptedABMTokenMetadata(
 		orgName = UnsavedABMTokenOrgName
 	}
 
+	if renewal {
+		// If we're renewing the token, we need to ensure the new token included in the context.
+		ctx = abmctx.NewContext(ctx, decryptedToken)
+	}
+
 	res, err := depClient.AccountDetail(ctx, orgName)
 	if err != nil {
 		var authErr *depclient.AuthError