From 7aa3fee45b2e3e8aca137ea6f549be48177d01fc Mon Sep 17 00:00:00 2001
From: Victor Lyuboslavsky <victor.lyuboslavsky@gmail.com>
Date: Fri, 10 Jan 2025 16:50:22 -0600
Subject: [PATCH] Fix issue when identical MDM commands are sent twice to the
 same device when replica DB is being used. (#25355)

For #24816
Fix issue when identical MDM commands are sent twice to the same device
when replica DB is being used.

Root cause was that ctxdb.RequirePrimary wasn't used correctly, and
proper test was missing.

# Checklist for submitter

- [x] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
- [x] Added/updated automated tests
- [x] Manual QA for all new/changed functionality
---
 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 0000000000..6b3313128e
--- /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 227df7fa1b..4b6c45fea1 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 0000000000..4e5c1ee673
--- /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 46c0e36406..a24fe947bc 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",