From 9f0404fbe15d0a08072fe8e8f1a5eed1a9570132 Mon Sep 17 00:00:00 2001
From: Sarah Gillespie <73313222+gillespi314@users.noreply.github.com>
Date: Tue, 18 Jun 2024 18:05:44 -0500
Subject: [PATCH] Add fallback checks for MDM unenrollment during migration
 flow (#19860)

Follow up to address additional issues uncovered during QA of #19512
---
 orbit/pkg/useraction/mdm_migration_darwin.go  | 47 ++++++++++++++-----
 .../useraction/mdm_migration_darwin_test.go   | 19 +++++++-
 server/fleet/hosts.go                         |  9 +++-
 3 files changed, 62 insertions(+), 13 deletions(-)

diff --git a/orbit/pkg/useraction/mdm_migration_darwin.go b/orbit/pkg/useraction/mdm_migration_darwin.go
index 323fc847faea..ba4c574d6ec5 100644
--- a/orbit/pkg/useraction/mdm_migration_darwin.go
+++ b/orbit/pkg/useraction/mdm_migration_darwin.go
@@ -14,6 +14,7 @@ import (
 	"text/template"
 	"time"
 
+	"github.com/fleetdm/fleet/v4/orbit/pkg/profiles"
 	"github.com/fleetdm/fleet/v4/pkg/file"
 	"github.com/fleetdm/fleet/v4/pkg/retry"
 	"github.com/rs/zerolog/log"
@@ -190,10 +191,13 @@ type swiftDialogMDMMigrator struct {
 	lastShownMu sync.RWMutex
 	showCh      chan struct{}
 
-	// testEnrollmentCheckFn is used in tests to mock the call to verify
+	// testEnrollmentCheckFileFn is used in tests to mock the call to verify
 	// the enrollment status of the host
-	testEnrollmentCheckFn     func() (bool, error)
-	unenrollmentRetryInterval time.Duration
+	testEnrollmentCheckFileFn func() (bool, error)
+	// testEnrollmentCheckStatusFn is used in tests to mock the call to verify
+	// the enrollment status of the host
+	testEnrollmentCheckStatusFn func() (bool, string, error)
+	unenrollmentRetryInterval   time.Duration
 }
 
 /**
@@ -273,21 +277,42 @@ func (m *swiftDialogMDMMigrator) renderError() (chan swiftDialogExitCode, chan e
 // unenroll, an error is returned.
 func (m *swiftDialogMDMMigrator) waitForUnenrollment() error {
 	maxRetries := int(mdmUnenrollmentTotalWaitTime.Seconds() / m.unenrollmentRetryInterval.Seconds())
-	fn := m.testEnrollmentCheckFn
-	if fn == nil {
-		fn = func() (bool, error) {
+	checkFileFn := m.testEnrollmentCheckFileFn
+	if checkFileFn == nil {
+		checkFileFn = func() (bool, error) {
 			return file.Exists(mdmEnrollmentFile)
 		}
 	}
+	checkStatusFn := m.testEnrollmentCheckStatusFn
+	if checkStatusFn == nil {
+		checkStatusFn = func() (bool, string, error) {
+			return profiles.IsEnrolledInMDM()
+		}
+	}
 	return retry.Do(func() error {
-		enrolled, err := fn()
-		if err != nil {
-			log.Error().Err(err).Msgf("checking enrollment status in migration modal, will retry in %s", m.unenrollmentRetryInterval)
-			return err
+		var unenrolled bool
 
+		fileExists, fileErr := checkFileFn()
+		if fileErr != nil {
+			log.Error().Err(fileErr).Msg("checking for existence of cloudConfigProfileInstalled in migration modal")
+		} else if fileExists {
+			log.Info().Msg("checking for existence of cloudConfigProfileInstalled in migration modal: found")
+		} else {
+			log.Info().Msg("checking for existence of cloudConfigProfileInstalled in migration modal: not found")
+			unenrolled = true
+		}
+
+		statusEnrolled, serverURL, statusErr := checkStatusFn()
+		if statusErr != nil {
+			log.Error().Err(statusErr).Msgf("checking profiles status in migration modal")
+		} else if statusEnrolled {
+			log.Info().Msgf("checking profiles status in migration modal: enrolled to %s", serverURL)
+		} else {
+			log.Info().Msg("checking profiles status in migration modal: not enrolled")
+			unenrolled = true
 		}
 
-		if enrolled {
+		if !unenrolled {
 			log.Info().Msgf("device is still enrolled, waiting %s", m.unenrollmentRetryInterval)
 			return errors.New("host didn't unenroll from MDM")
 		}
diff --git a/orbit/pkg/useraction/mdm_migration_darwin_test.go b/orbit/pkg/useraction/mdm_migration_darwin_test.go
index ff13512dfd16..76be50f9644e 100644
--- a/orbit/pkg/useraction/mdm_migration_darwin_test.go
+++ b/orbit/pkg/useraction/mdm_migration_darwin_test.go
@@ -39,7 +39,7 @@ func TestWaitForUnenrollment(t *testing.T) {
 	for _, c := range cases {
 		t.Run(c.name, func(t *testing.T) {
 			tries := 0
-			m.testEnrollmentCheckFn = func() (bool, error) {
+			m.testEnrollmentCheckFileFn = func() (bool, error) {
 				if tries >= c.unenrollAfterNTries {
 					return false, c.enrollErr
 				}
@@ -47,6 +47,10 @@ func TestWaitForUnenrollment(t *testing.T) {
 				return true, c.enrollErr
 			}
 
+			m.testEnrollmentCheckStatusFn = func() (bool, string, error) {
+				return true, "example.com", nil
+			}
+
 			outErr := m.waitForUnenrollment()
 			if c.wantErr {
 				require.Error(t, outErr)
@@ -56,4 +60,17 @@ func TestWaitForUnenrollment(t *testing.T) {
 			}
 		})
 	}
+
+	t.Run("fallback to enrollment check file", func(t *testing.T) {
+		m.testEnrollmentCheckFileFn = func() (bool, error) {
+			return true, nil
+		}
+
+		m.testEnrollmentCheckStatusFn = func() (bool, string, error) {
+			return false, "", nil
+		}
+
+		outErr := m.waitForUnenrollment()
+		require.NoError(t, outErr)
+	})
 }
diff --git a/server/fleet/hosts.go b/server/fleet/hosts.go
index 2e8c385d3d8e..013b8d0f35cb 100644
--- a/server/fleet/hosts.go
+++ b/server/fleet/hosts.go
@@ -708,7 +708,14 @@ func (h *Host) IsEligibleForDEPMigration(isConnectedToFleetMDM bool) bool {
 func (h *Host) NeedsDEPEnrollment(isConnectedToFleetMDM bool) bool {
 	return h.MDMInfo != nil &&
 		!h.MDMInfo.Enrolled &&
-		!isConnectedToFleetMDM &&
+		// as a special case for migration with user interaction, we
+		// also check the information stored in host_mdm, and assume
+		// the host needs migration if it's not Fleet
+		//
+		// this is because we can't always rely on nano setting
+		// `nano_enrollment.active = 1` since sometimes Fleet won't get
+		// the checkout message from the host.
+		(!isConnectedToFleetMDM || h.MDMInfo.Name != WellKnownMDMFleet) &&
 		h.IsDEPAssignedToFleet()
 }