From 7cd178efec9062a3625d524ea62b913a4f14f733 Mon Sep 17 00:00:00 2001
From: Victor Lyuboslavsky <victor@fleetdm.com>
Date: Fri, 10 Jan 2025 16:13:57 -0600
Subject: [PATCH] Fix issue when identical MDM commands are sent twice to the
 same device when replica DB is being used.

---
 changes/24816-fix-double-mdm-commands         |  1 +
 server/mdm/nanomdm/service/nanomdm/service.go |  2 +-
 .../nanomdm/service/nanomdm/service_test.go   | 58 +++++++++++++++++++
 server/service/integration_mdm_test.go        |  2 +
 4 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 changes/24816-fix-double-mdm-commands
 create mode 100644 server/mdm/nanomdm/service/nanomdm/service_test.go

diff --git a/changes/24816-fix-double-mdm-commands b/changes/24816-fix-double-mdm-commands
new file mode 100644
index 000000000000..6b3313128e7b
--- /dev/null
+++ b/changes/24816-fix-double-mdm-commands
@@ -0,0 +1 @@
+Fix issue when identical MDM commands are sent twice to the same device when replica DB is being used.
diff --git a/server/mdm/nanomdm/service/nanomdm/service.go b/server/mdm/nanomdm/service/nanomdm/service.go
index 227df7fa1b2c..4b6c45fea144 100644
--- a/server/mdm/nanomdm/service/nanomdm/service.go
+++ b/server/mdm/nanomdm/service/nanomdm/service.go
@@ -256,7 +256,7 @@ func (s *Service) CommandAndReportResults(r *mdm.Request, results *mdm.CommandRe
 	}
 	if results.Status != "Idle" {
 		// If the host is not idle, we use primary DB since we just wrote results of previous command.
-		ctxdb.RequirePrimary(r.Context, true)
+		r.Context = ctxdb.RequirePrimary(r.Context, true)
 	}
 	cmd, err := s.store.RetrieveNextCommand(r, results.Status == "NotNow")
 	if err != nil {
diff --git a/server/mdm/nanomdm/service/nanomdm/service_test.go b/server/mdm/nanomdm/service/nanomdm/service_test.go
new file mode 100644
index 000000000000..4e5c1ee673e9
--- /dev/null
+++ b/server/mdm/nanomdm/service/nanomdm/service_test.go
@@ -0,0 +1,58 @@
+package nanomdm
+
+import (
+	"context"
+	"testing"
+
+	"github.com/fleetdm/fleet/v4/server/contexts/ctxdb"
+	"github.com/fleetdm/fleet/v4/server/mdm/nanomdm/mdm"
+	mock "github.com/fleetdm/fleet/v4/server/mock/mdm"
+	"github.com/micromdm/nanolib/log"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCommandAndReportResultsPrimaryDBUse(t *testing.T) {
+	ds := new(mock.MDMAppleStore)
+
+	enrollID := &mdm.EnrollID{
+		ID:   "1",
+		Type: mdm.Device,
+	}
+	s := Service{
+		logger: log.NopLogger,
+		store:  ds,
+		normalizer: func(e *mdm.Enrollment) *mdm.EnrollID {
+			return enrollID
+		},
+	}
+	ds.StoreCommandReportFunc = func(r *mdm.Request, report *mdm.CommandResults) error {
+		return nil
+	}
+	var primaryRequired bool
+	ds.RetrieveNextCommandFunc = func(r *mdm.Request, skipNotNow bool) (*mdm.CommandWithSubtype, error) {
+		assert.Equal(t, primaryRequired, ctxdb.IsPrimaryRequired(r.Context))
+		return nil, nil
+	}
+
+	mdmRequest := &mdm.Request{
+		Context: context.Background(),
+	}
+	mdmCommandResults := &mdm.CommandResults{
+		Status: "Idle",
+	}
+	// We don't use primary DB with "Idle" status because we don't update the status of existing commands
+	cmd, err := s.CommandAndReportResults(mdmRequest, mdmCommandResults)
+	require.NoError(t, err)
+	require.Nil(t, cmd)
+
+	// We use primary DB with non-"Idle" status
+	mdmCommandResults = &mdm.CommandResults{
+		Status: "Acknowledge",
+	}
+	primaryRequired = true
+	cmd, err = s.CommandAndReportResults(mdmRequest, mdmCommandResults)
+	require.NoError(t, err)
+	require.Nil(t, cmd)
+
+}
diff --git a/server/service/integration_mdm_test.go b/server/service/integration_mdm_test.go
index 46c0e36406b3..a24fe947bcf2 100644
--- a/server/service/integration_mdm_test.go
+++ b/server/service/integration_mdm_test.go
@@ -11583,6 +11583,8 @@ func (s *integrationMDMTestSuite) TestVPPApps() {
 			// Simulate successful installation on the host
 			cmd, err = mdmClient.Acknowledge(cmd.CommandUUID)
 			require.NoError(t, err)
+			// No further commands expected
+			assert.Nil(t, cmd)
 
 			listResp = listHostsResponse{}
 			s.DoJSON("GET", "/api/latest/fleet/hosts", nil, http.StatusOK, &listResp, "software_status", "installed", "team_id",