fix: ensure we assign ABM profiles for modified hosts (#13275)

for #12958 and #13110

changes/12958-profiles-assignment

@@ -0,0 +1 @@
+* Ensure DEP profiles are assigned even for devices that already exist and have an op type = "modified"

server/datastore/mysql/hosts.go

@@ -4162,8 +4162,8 @@ func amountHostsByOsqueryVersionDB(ctx context.Context, db sqlx.QueryerContext)
 	return counts, nil
 }
 
-func (ds *Datastore) GetMatchingHostSerials(ctx context.Context, serials []string) (map[string]struct{}, error) {
-	result := map[string]struct{}{}
+func (ds *Datastore) GetMatchingHostSerials(ctx context.Context, serials []string) (map[string]*fleet.Host, error) {
+	result := map[string]*fleet.Host{}
 	if len(serials) == 0 {
 		return result, nil
 	}
@@ -4172,17 +4172,17 @@ func (ds *Datastore) GetMatchingHostSerials(ctx context.Context, serials []strin
 	for _, serial := range serials {
 		args = append(args, serial)
 	}
-	stmt, args, err := sqlx.In("SELECT hardware_serial FROM hosts WHERE hardware_serial IN (?)", args)
+	stmt, args, err := sqlx.In("SELECT hardware_serial, team_id FROM hosts WHERE hardware_serial IN (?)", args)
 	if err != nil {
 		return nil, ctxerr.Wrap(ctx, err, "building IN statement for matching hosts")
 	}
-	var matchingSerials []string
-	if err := sqlx.SelectContext(ctx, ds.reader(ctx), &matchingSerials, stmt, args...); err != nil {
+	var matchingHosts []*fleet.Host
+	if err := sqlx.SelectContext(ctx, ds.reader(ctx), &matchingHosts, stmt, args...); err != nil {
 		return nil, err
 	}
 
-	for _, serial := range matchingSerials {
-		result[serial] = struct{}{}
+	for _, host := range matchingHosts {
+		result[host.HardwareSerial] = host
 	}
 
 	return result, nil

server/datastore/mysql/hosts_test.go

@@ -7132,7 +7132,15 @@ func testHostsListHostsLiteByUUIDs(t *testing.T, ds *Datastore) {
 func testGetMatchingHostSerials(t *testing.T, ds *Datastore) {
 	ctx := context.Background()
 	serials := []string{"foo", "bar", "baz"}
+	team, err := ds.NewTeam(context.Background(), &fleet.Team{
+		Name: "team1",
+	})
+	require.NoError(t, err)
 	for i, serial := range serials {
+		var tmID *uint
+		if serial == "bar" {
+			tmID = &team.ID
+		}
 		_, err := ds.NewHost(ctx, &fleet.Host{
 			DetailUpdatedAt: time.Now(),
 			LabelUpdatedAt:  time.Now(),
@@ -7145,6 +7153,7 @@ func testGetMatchingHostSerials(t *testing.T, ds *Datastore) {
 			PrimaryIP:       "192.168.1.1",
 			PrimaryMac:      "30-65-EC-6F-C4-58",
 			HardwareSerial:  serial,
+			TeamID:          tmID,
 		})
 		require.NoError(t, err)
 	}
@@ -7152,13 +7161,29 @@ func testGetMatchingHostSerials(t *testing.T, ds *Datastore) {
 	cases := []struct {
 		name string
 		in   []string
-		want map[string]struct{}
+		want map[string]*fleet.Host
 		err  string
 	}{
-		{"no serials provided", []string{}, map[string]struct{}{}, ""},
-		{"no matching serials", []string{"oof", "rab"}, map[string]struct{}{}, ""},
-		{"partial matches", []string{"foo", "rab"}, map[string]struct{}{"foo": {}}, ""},
-		{"all matching", []string{"foo", "bar", "baz"}, map[string]struct{}{"foo": {}, "bar": {}, "baz": {}}, ""},
+		{"no serials provided", []string{}, map[string]*fleet.Host{}, ""},
+		{"no matching serials", []string{"oof", "rab"}, map[string]*fleet.Host{}, ""},
+		{
+			"partial matches",
+			[]string{"foo", "rab"},
+			map[string]*fleet.Host{
+				"foo": {HardwareSerial: "foo", TeamID: nil},
+			},
+			"",
+		},
+		{
+			"all matching",
+			[]string{"foo", "bar", "baz"},
+			map[string]*fleet.Host{
+				"foo": {HardwareSerial: "foo", TeamID: nil},
+				"bar": {HardwareSerial: "bar", TeamID: &team.ID},
+				"baz": {HardwareSerial: "baz", TeamID: nil},
+			},
+			"",
+		},
 	}
 
 	for _, tt := range cases {

server/fleet/datastore.go

@@ -980,8 +980,8 @@ type Datastore interface {
 	GetMDMAppleDefaultSetupAssistant(ctx context.Context, teamID *uint) (profileUUID string, updatedAt time.Time, err error)
 
 	// GetMatchingHostSerials receives a list of serial numbers and returns
-	// a map with all the matching serial numbers in the database.
-	GetMatchingHostSerials(ctx context.Context, serials []string) (map[string]struct{}, error)
+	// a map that only contains the serials that have a matching row in the `hosts` table.
+	GetMatchingHostSerials(ctx context.Context, serials []string) (map[string]*Host, error)
 
 	// DeleteHostDEPAssignments marks as deleted entries in
 	// host_dep_assignments for host with matching serials.

server/mdm/apple/apple_mdm.go

@@ -433,7 +433,7 @@ func (d *DEPService) processDeviceResponse(ctx context.Context, depClient *godep
 		return ctxerr.Wrap(ctx, err, "deleting DEP assignments")
 	}
 
-	n, tmID, err := d.ds.IngestMDMAppleDevicesFromDEPSync(ctx, addedDevices)
+	n, defaultABMTeamID, err := d.ds.IngestMDMAppleDevicesFromDEPSync(ctx, addedDevices)
 	switch {
 	case err != nil:
 		level.Error(kitlog.With(d.logger)).Log("err", err)
@@ -446,11 +446,83 @@ func (d *DEPService) processDeviceResponse(ctx context.Context, depClient *godep
 
 	// at this point, the hosts rows are created for the devices, with the
 	// correct team_id, so we know what team-specific profile needs to be applied.
+	//
+	// collect a map of all the profiles => serials we need to assign.
+	profileToSerials := map[string][]string{}
+
+	// each new device should be assigned the DEP profile of the default
+	// ABM team as configured by the IT admin.
+	if len(addedDevices) > 0 {
+		level.Info(kitlog.With(d.logger)).Log("msg", "gathering added serials to assign devices", "len", len(addedDevices))
+		profUUID, err := d.getProfileUUIDForTeam(ctx, defaultABMTeamID)
+		if err != nil {
+			return ctxerr.Wrapf(ctx, err, "getting profile for default team with id: %v", defaultABMTeamID)
+		}
+
+		var addedSerials []string
+		for _, d := range addedDevices {
+			addedSerials = append(addedSerials, d.SerialNumber)
+		}
+		profileToSerials[profUUID] = addedSerials
+	} else {
+		level.Info(kitlog.With(d.logger)).Log("msg", "no added devices to assign DEP profiles")
+	}
+
+	// for all other hosts we received, find out the right DEP profile to assign, based on the team.
+	if len(existingSerials) > 0 {
+		level.Info(kitlog.With(d.logger)).Log("msg", "gathering existing serials to assign devices", "len", len(existingSerials))
+		serialsByTeam := map[*uint][]string{}
+		for _, host := range existingSerials {
+			if serialsByTeam[host.TeamID] == nil {
+				serialsByTeam[host.TeamID] = []string{}
+			}
+			serialsByTeam[host.TeamID] = append(serialsByTeam[host.TeamID], host.HardwareSerial)
+		}
+		for team, serials := range serialsByTeam {
+			profUUID, err := d.getProfileUUIDForTeam(ctx, team)
+			if err != nil {
+				return ctxerr.Wrapf(ctx, err, "getting profile for team with id: %v", team)
+			}
+			if profileToSerials[profUUID] == nil {
+				profileToSerials[profUUID] = []string{}
+			}
+			profileToSerials[profUUID] = append(profileToSerials[profUUID], serials...)
+
+		}
+	} else {
+		level.Info(kitlog.With(d.logger)).Log("msg", "no existing devices to assign DEP profiles")
+	}
+
+	for profUUID, serials := range profileToSerials {
+		logger := kitlog.With(d.logger, "profile_uuid", profUUID)
+		level.Info(logger).Log("msg", "calling DEP client to assign profile", "profile_uuid", profUUID)
+		apiResp, err := depClient.AssignProfile(ctx, DEPName, profUUID, serials...)
+		if err != nil {
+			level.Info(logger).Log(
+				"msg", "assign profile",
+				"devices", len(serials),
+				"err", err,
+			)
+			return fmt.Errorf("assign profile: %w", err)
+		}
+
+		logs := []interface{}{
+			"msg", "profile assigned",
+			"devices", len(serials),
+		}
+		logs = append(logs, logCountsForResults(apiResp.Devices)...)
+		level.Info(logger).Log(logs...)
+	}
+
+	return nil
+}
+
+func (d *DEPService) getProfileUUIDForTeam(ctx context.Context, tmID *uint) (string, error) {
 	var appleBMTeam *fleet.Team
 	if tmID != nil {
 		tm, err := d.ds.Team(ctx, *tmID)
 		if err != nil && !fleet.IsNotFound(err) {
-			return ctxerr.Wrap(ctx, err, "get team")
+			return "", ctxerr.Wrap(ctx, err, "get team")
 		}
 		appleBMTeam = tm
 	}
@@ -458,53 +530,16 @@ func (d *DEPService) processDeviceResponse(ctx context.Context, depClient *godep
 	// get profile uuid of team or default
 	profUUID, _, err := d.EnsureCustomSetupAssistantIfExists(ctx, appleBMTeam)
 	if err != nil {
-		return fmt.Errorf("ensure setup assistant for team %v: %w", tmID, err)
+		return "", fmt.Errorf("ensure setup assistant for team %v: %w", tmID, err)
 	}
 	if profUUID == "" {
 		profUUID, _, err = d.EnsureDefaultSetupAssistant(ctx, appleBMTeam)
 		if err != nil {
-			return fmt.Errorf("ensure default setup assistant: %w", err)
+			return "", fmt.Errorf("ensure default setup assistant: %w", err)
 		}
 	}
 
-	if profUUID == "" {
-		level.Debug(d.logger).Log("msg", "empty assigner profile UUID")
-		return nil
-	}
-
-	logger := kitlog.With(d.logger, "profile_uuid", profUUID)
-
-	if len(addedDevices) < 1 {
-		level.Debug(logger).Log(
-			"msg", "no serials to assign",
-			"devices", len(resp.Devices),
-		)
-		return nil
-	}
-
-	var addedSerials []string
-	for _, d := range addedDevices {
-		addedSerials = append(addedSerials, d.SerialNumber)
-	}
-
-	apiResp, err := depClient.AssignProfile(ctx, DEPName, profUUID, addedSerials...)
-	if err != nil {
-		level.Info(logger).Log(
-			"msg", "assign profile",
-			"devices", len(addedSerials),
-			"err", err,
-		)
-		return fmt.Errorf("assign profile: %w", err)
-	}
-
-	logs := []interface{}{
-		"msg", "profile assigned",
-		"devices", len(addedSerials),
-	}
-	logs = append(logs, logCountsForResults(apiResp.Devices)...)
-	level.Info(logger).Log(logs...)
-
-	return nil
+	return profUUID, nil
 }
 
 // logCountsForResults tries to aggregate the result types and log the counts.

server/mock/datastore_mock.go

@@ -644,7 +644,7 @@ type SetMDMAppleDefaultSetupAssistantProfileUUIDFunc func(ctx context.Context, t
 
 type GetMDMAppleDefaultSetupAssistantFunc func(ctx context.Context, teamID *uint) (profileUUID string, updatedAt time.Time, err error)
 
-type GetMatchingHostSerialsFunc func(ctx context.Context, serials []string) (map[string]struct{}, error)
+type GetMatchingHostSerialsFunc func(ctx context.Context, serials []string) (map[string]*fleet.Host, error)
 
 type DeleteHostDEPAssignmentsFunc func(ctx context.Context, serials []string) error
 
@@ -3818,7 +3818,7 @@ func (s *DataStore) GetMDMAppleDefaultSetupAssistant(ctx context.Context, teamID
 	return s.GetMDMAppleDefaultSetupAssistantFunc(ctx, teamID)
 }
 
-func (s *DataStore) GetMatchingHostSerials(ctx context.Context, serials []string) (map[string]struct{}, error) {
+func (s *DataStore) GetMatchingHostSerials(ctx context.Context, serials []string) (map[string]*fleet.Host, error) {
 	s.mu.Lock()
 	s.GetMatchingHostSerialsFuncInvoked = true
 	s.mu.Unlock()

server/service/integration_mdm_test.go

@@ -1161,7 +1161,14 @@ func (s *integrationMDMTestSuite) TestDEPProfileAssignment() {
 		{SerialNumber: uuid.New().String(), Model: "MacBook Mini", OS: "osx", OpType: "modified"},
 	}
 
+	type profileAssignmentReq struct {
+		ProfileUUID string   `json:"profile_uuid"`
+		Devices     []string `json:"devices"`
+	}
+	profileAssignmentReqs := []profileAssignmentReq{}
+
 	runDEPSchedule := func() {
+		profileAssignmentReqs = []profileAssignmentReq{}
 		ch := make(chan bool)
 		s.onDEPScheduleDone = func() { close(ch) }
 		_, err := s.depSchedule.Trigger()
@@ -1220,7 +1227,7 @@ func (s *integrationMDMTestSuite) TestDEPProfileAssignment() {
 			err := encoder.Encode(map[string]string{"auth_session_token": "xyz"})
 			require.NoError(t, err)
 		case "/profile":
-			err := encoder.Encode(godep.ProfileResponse{ProfileUUID: "abc"})
+			err := encoder.Encode(godep.ProfileResponse{ProfileUUID: uuid.New().String()})
 			require.NoError(t, err)
 		case "/server/devices":
 			// This endpoint  is used to get an initial list of
@@ -1233,6 +1240,11 @@ func (s *integrationMDMTestSuite) TestDEPProfileAssignment() {
 			err := encoder.Encode(godep.DeviceResponse{Devices: devices, Cursor: "foo"})
 			require.NoError(t, err)
 		case "/profile/devices":
+			b, err := io.ReadAll(r.Body)
+			require.NoError(t, err)
+			var prof profileAssignmentReq
+			require.NoError(t, json.Unmarshal(b, &prof))
+			profileAssignmentReqs = append(profileAssignmentReqs, prof)
 			_, _ = w.Write([]byte(`{}`))
 		default:
 			_, _ = w.Write([]byte(`{}`))
@@ -1261,6 +1273,12 @@ func (s *integrationMDMTestSuite) TestDEPProfileAssignment() {
 		require.NoError(t, err)
 	}
 	require.ElementsMatch(t, wantSerials, gotSerials)
+	// called two times:
+	// - one when we get the initial list of devices (/server/devices)
+	// - one when we do the device sync (/device/sync)
+	require.Len(t, profileAssignmentReqs, 2)
+	require.Len(t, profileAssignmentReqs[0].Devices, 1)
+	require.Len(t, profileAssignmentReqs[1].Devices, len(devices))
 
 	// create a new host
 	nonDEPHost := createHostAndDeviceToken(t, s.ds, "not-dep")
@@ -1313,6 +1331,23 @@ func (s *integrationMDMTestSuite) TestDEPProfileAssignment() {
 	}
 	require.True(t, found)
 
+	// add devices[1].SerialNumber to a team
+	teamName := t.Name() + "team1"
+	team := &fleet.Team{
+		Name:        teamName,
+		Description: "desc team1",
+	}
+	var createTeamResp teamResponse
+	s.DoJSON("POST", "/api/latest/fleet/teams", team, http.StatusOK, &createTeamResp)
+	require.NotZero(t, createTeamResp.Team.ID)
+	team = createTeamResp.Team
+	for _, h := range listHostsRes.Hosts {
+		if h.HardwareSerial == devices[1].SerialNumber {
+			err = s.ds.AddHostsToTeam(ctx, &team.ID, []uint{h.ID})
+			require.NoError(t, err)
+		}
+	}
+
 	// modify the response and trigger another sync to include:
 	//
 	// 1. A repeated device with "added"
@@ -1348,6 +1383,21 @@ func (s *integrationMDMTestSuite) TestDEPProfileAssignment() {
 		}
 	}
 	require.ElementsMatch(t, wantSerials, gotSerials)
+	require.Len(t, profileAssignmentReqs, 3)
+
+	// first request to get a list of profiles
+	// TODO: seems like we're doing this request on each loop?
+	require.Len(t, profileAssignmentReqs[0].Devices, 1)
+	require.Equal(t, devices[0].SerialNumber, profileAssignmentReqs[0].Devices[0])
+	// - existing device with "added"
+	// - new device with "added"
+	require.Len(t, profileAssignmentReqs[1].Devices, 2)
+	require.Equal(t, devices[0].SerialNumber, profileAssignmentReqs[1].Devices[0])
+	require.Equal(t, addedSerial, profileAssignmentReqs[1].Devices[1])
+
+	// - existing device with "modified" and a different team (thus different profile request)
+	require.Len(t, profileAssignmentReqs[2].Devices, 1)
+	require.Equal(t, devices[1].SerialNumber, profileAssignmentReqs[2].Devices[0])
 
 	// entries for all hosts except for the one with OpType = "deleted"
 	assignment, err := s.ds.GetHostDEPAssignment(ctx, deletedHostID)