From e360e7e614bbc3d966f211b4b340153b6835d267 Mon Sep 17 00:00:00 2001
From: Victor Lyuboslavsky <victor@fleetdm.com>
Date: Fri, 13 Jun 2025 13:55:44 -0500
Subject: [PATCH] Fix osquery_perf so it doesn't miss DeclarativeManagement
 commands. (#29975)

Fixes #29973

Tests are failing due to infra issues with https://proxy.golang.org

# Checklist for submitter
- [x] Manual QA for all new/changed functionality
---
 cmd/osquery-perf/agent.go | 42 +++++++++++++++++++++++++++------------
 1 file changed, 29 insertions(+), 13 deletions(-)

diff --git a/cmd/osquery-perf/agent.go b/cmd/osquery-perf/agent.go
index 470fd835d5..6647a418e7 100644
--- a/cmd/osquery-perf/agent.go
+++ b/cmd/osquery-perf/agent.go
@@ -866,8 +866,8 @@ func (a *agent) runMacosMDMLoop() {
 		for mdmCommandPayload != nil {
 			a.stats.IncrementMDMCommandsReceived()
 
-			if mdmCommandPayload.Command.RequestType == "InstallProfile" {
-
+			switch mdmCommandPayload.Command.RequestType {
+			case "InstallProfile":
 				if a.mdmProfileFailureProb > 0.0 && rand.Float64() <= a.mdmProfileFailureProb {
 					errChain := []mdm.ErrorChain{
 						{
@@ -877,21 +877,37 @@ func (a *agent) runMacosMDMLoop() {
 						},
 					}
 					mdmCommandPayload, err = a.macMDMClient.Err(mdmCommandPayload.CommandUUID, errChain)
+					if err != nil {
+						log.Printf("MDM Error request failed: %s", err)
+						a.stats.IncrementMDMErrors()
+						break INNER_FOR_LOOP
+					}
 				} else {
 					mdmCommandPayload, err = a.macMDMClient.Acknowledge(mdmCommandPayload.CommandUUID)
+					if err != nil {
+						log.Printf("MDM Acknowledge request failed: %s", err)
+						a.stats.IncrementMDMErrors()
+						break INNER_FOR_LOOP
+					}
 				}
-
-			} else {
-				mdmCommandPayload, err = a.macMDMClient.Acknowledge(mdmCommandPayload.CommandUUID)
-			}
-
-			if err != nil {
-				log.Printf("MDM Acknowledge request failed: %s", err)
-				a.stats.IncrementMDMErrors()
-				break INNER_FOR_LOOP
-			}
-			if mdmCommandPayload != nil && mdmCommandPayload.Command.RequestType == "DeclarativeManagement" {
+			case "DeclarativeManagement":
+				// Device immediately responds with Acknowledged status and then contacts the Declarations endpoints.
+				nextMdmCommandPayload, err := a.macMDMClient.Acknowledge(mdmCommandPayload.CommandUUID)
+				if err != nil {
+					log.Printf("MDM Acknowledge request failed: %s", err)
+					a.stats.IncrementMDMErrors()
+					break INNER_FOR_LOOP
+				}
+				// Note: Declarative management could happen async while other MDM commands proceed. This is a potential enhancement.
 				a.doDeclarativeManagement(mdmCommandPayload)
+				mdmCommandPayload = nextMdmCommandPayload
+			default:
+				mdmCommandPayload, err = a.macMDMClient.Acknowledge(mdmCommandPayload.CommandUUID)
+				if err != nil {
+					log.Printf("MDM Acknowledge request failed: %s", err)
+					a.stats.IncrementMDMErrors()
+					break INNER_FOR_LOOP
+				}
 			}
 		}
 	}