From 3cfe583ea02b760f1a1ec6331eab3af63f0bebbb Mon Sep 17 00:00:00 2001
From: Roberto Dip <dip.jesusr@gmail.com>
Date: Mon, 19 Aug 2024 12:02:43 -0300
Subject: [PATCH] fix issue with disk encryption banner (#21385)

for #21381

# 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] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/Committing-Changes.md#changes-files)
for more information.
- [x] Manual QA for all new/changed functionality
- For Orbit and Fleet Desktop changes:
- [x] Orbit runs on macOS, Linux and Windows. Check if the orbit
feature/bugfix should only apply to one platform (`runtime.GOOS`).
- [ ] Manual QA must be performed in the three main OSs, macOS, Windows
and Linux.
- [x] Auto-update manual QA, from released version of component to new
version (see [tools/tuf/test](../tools/tuf/test/README.md)).
---
 orbit/changes/21381-fv              |  1 +
 orbit/cmd/orbit/orbit.go            |  7 ++++++-
 orbit/pkg/update/disk_encryption.go | 28 +++++++++++++++++++++++++---
 3 files changed, 32 insertions(+), 4 deletions(-)
 create mode 100644 orbit/changes/21381-fv

diff --git a/orbit/changes/21381-fv b/orbit/changes/21381-fv
new file mode 100644
index 0000000000..942f2f435f
--- /dev/null
+++ b/orbit/changes/21381-fv
@@ -0,0 +1 @@
+* Fixed an issue that would display a disk encryption modal with MDM configured and FileVault enabled if the user hadn't escrowed the key in the past.
diff --git a/orbit/cmd/orbit/orbit.go b/orbit/cmd/orbit/orbit.go
index d74ab7efeb..334d69be85 100644
--- a/orbit/cmd/orbit/orbit.go
+++ b/orbit/cmd/orbit/orbit.go
@@ -1223,7 +1223,12 @@ func main() {
 			if orbitClient.GetServerCapabilities().Has(fleet.CapabilityEscrowBuddy) {
 				orbitClient.RegisterConfigReceiver(update.NewEscrowBuddyRunner(updateRunner, 5*time.Minute))
 			} else {
-				orbitClient.RegisterConfigReceiver(update.ApplyDiskEncryptionRunnerMiddleware())
+				orbitClient.RegisterConfigReceiver(
+					update.ApplyDiskEncryptionRunnerMiddleware(
+						orbitClient.GetServerCapabilities,
+						orbitClient.TriggerOrbitRestart,
+					),
+				)
 			}
 		}
 
diff --git a/orbit/pkg/update/disk_encryption.go b/orbit/pkg/update/disk_encryption.go
index ae09f386d6..1f6497265d 100644
--- a/orbit/pkg/update/disk_encryption.go
+++ b/orbit/pkg/update/disk_encryption.go
@@ -1,6 +1,7 @@
 package update
 
 import (
+	"errors"
 	"sync/atomic"
 
 	"github.com/fleetdm/fleet/v4/orbit/pkg/useraction"
@@ -11,16 +12,37 @@ import (
 const maxRetries = 2
 
 type DiskEncryptionRunner struct {
-	isRunning atomic.Bool
+	isRunning           atomic.Bool
+	capabilitiesFetcher func() fleet.CapabilityMap
+	triggerOrbitRestart func(reason string)
 }
 
-func ApplyDiskEncryptionRunnerMiddleware() fleet.OrbitConfigReceiver {
-	return &DiskEncryptionRunner{}
+func ApplyDiskEncryptionRunnerMiddleware(
+	capabilitiesFetcher func() fleet.CapabilityMap,
+	triggerOrbitRestart func(reason string),
+) fleet.OrbitConfigReceiver {
+	return &DiskEncryptionRunner{
+		capabilitiesFetcher: capabilitiesFetcher,
+		triggerOrbitRestart: triggerOrbitRestart,
+	}
 }
 
 func (d *DiskEncryptionRunner) Run(cfg *fleet.OrbitConfig) error {
 	log.Debug().Msgf("running disk encryption fetcher middleware, notification: %v, isIdle: %v", cfg.Notifications.RotateDiskEncryptionKey, d.isRunning.Load())
 
+	if d.capabilitiesFetcher == nil {
+		return errors.New("disk encryption runner needs a capabilitites fetcher configured")
+	}
+
+	if d.triggerOrbitRestart == nil {
+		return errors.New("disk encryption runner needs a function to trigger orbit restarts configured")
+	}
+
+	if d.capabilitiesFetcher().Has(fleet.CapabilityEscrowBuddy) {
+		d.triggerOrbitRestart("server has Escrow Buddy capability but old disk encryption fetcher was running")
+		return nil
+	}
+
 	if cfg.Notifications.RotateDiskEncryptionKey && !d.isRunning.Swap(true) {
 		go func() {
 			defer d.isRunning.Store(false)