More fixes to support users with hosts in same team and hosts in different teams (#17789)

#17441

Author: lucasmrod

cmd/fleet/calendar_cron.go

@@ -125,7 +125,7 @@ func cronCalendarEventsForTeam(
 	for _, policy := range policies {
 		policyIDs = append(policyIDs, policy.ID)
 	}
-	hosts, err := ds.GetHostsPolicyMemberships(ctx, domain, policyIDs)
+	hosts, err := ds.GetTeamHostsPolicyMemberships(ctx, domain, team.ID, policyIDs)
 	if err != nil {
 		return fmt.Errorf("get team hosts failing policies: %w", err)
 	}
@@ -150,22 +150,28 @@ func cronCalendarEventsForTeam(
 	}
 	level.Debug(logger).Log(
 		"msg", "summary",
+		"team_id", team.ID,
 		"passing_hosts", len(passingHosts),
 		"failing_hosts", len(failingHosts),
 		"failing_hosts_without_associated_email", len(failingHostsWithoutAssociatedEmail),
 	)
 
+	// Remove calendar events from hosts that are passing the calendar policies.
+	//
+	// We execute this first to remove any calendar events for a user that is now passing
+	// policies on one of its hosts, and possibly create a new calendar event if they have
+	// another failing host on the same team.
+	if err := removeCalendarEventsFromPassingHosts(ctx, ds, calendar, passingHosts); err != nil {
+		level.Info(logger).Log("msg", "removing calendar events from passing hosts", "err", err)
+	}
+
+	// Process hosts that are failing calendar policies.
 	if err := processCalendarFailingHosts(
 		ctx, ds, calendar, orgName, failingHosts, logger,
 	); err != nil {
 		level.Info(logger).Log("msg", "processing failing hosts", "err", err)
 	}
 
-	// Remove calendar events from hosts that are passing the policies.
-	if err := removeCalendarEventsFromPassingHosts(ctx, ds, calendar, passingHosts); err != nil {
-		level.Info(logger).Log("msg", "removing calendar events from passing hosts", "err", err)
-	}
-
 	// At last we want to log the hosts that are failing and don't have an associated email.
 	logHostsWithoutAssociatedEmail(
 		domain,
@@ -184,14 +190,26 @@ func processCalendarFailingHosts(
 	hosts []fleet.HostPolicyMembershipData,
 	logger kitlog.Logger,
 ) error {
+	hosts = filterHostsWithSameEmail(hosts)
+
 	for _, host := range hosts {
 		logger := log.With(logger, "host_id", host.HostID)
 
-		hostCalendarEvent, calendarEvent, err := ds.GetHostCalendarEvent(ctx, host.HostID)
+		hostCalendarEvent, calendarEvent, err := ds.GetHostCalendarEventByEmail(ctx, host.Email)
 
 		expiredEvent := false
 		webhookAlreadyFiredThisMonth := false
 		if err == nil {
+			if hostCalendarEvent.HostID != host.HostID {
+				// This calendar event belongs to another host with this associated email,
+				// thus we skip this entry.
+				continue // continue with next host
+			}
+			if hostCalendarEvent.WebhookStatus == fleet.CalendarWebhookStatusPending {
+				// This can happen if the host went offline (and never returned results)
+				// after setting the webhook as pending.
+				continue // continue with next host
+			}
 			now := time.Now()
 			webhookAlreadyFired := hostCalendarEvent.WebhookStatus == fleet.CalendarWebhookStatusSent
 			if webhookAlreadyFired && sameDate(now, calendarEvent.StartTime) {
@@ -200,7 +218,7 @@ func processCalendarFailingHosts(
 				continue // continue with next host
 			}
 			webhookAlreadyFiredThisMonth = webhookAlreadyFired && sameMonth(now, calendarEvent.StartTime)
-			if calendarEvent.EndTime.Before(time.Now()) {
+			if calendarEvent.EndTime.Before(now) {
 				expiredEvent = true
 			}
 		}
@@ -232,6 +250,25 @@ func processCalendarFailingHosts(
 	return nil
 }
 
+func filterHostsWithSameEmail(hosts []fleet.HostPolicyMembershipData) []fleet.HostPolicyMembershipData {
+	minHostPerEmail := make(map[string]fleet.HostPolicyMembershipData)
+	for _, host := range hosts {
+		minHost, ok := minHostPerEmail[host.Email]
+		if !ok {
+			minHostPerEmail[host.Email] = host
+			continue
+		}
+		if host.HostID < minHost.HostID {
+			minHostPerEmail[host.Email] = host
+		}
+	}
+	filtered := make([]fleet.HostPolicyMembershipData, 0, len(minHostPerEmail))
+	for _, host := range minHostPerEmail {
+		filtered = append(filtered, host)
+	}
+	return filtered
+}
+
 func processFailingHostExistingCalendarEvent(
 	ctx context.Context,
 	ds fleet.Datastore,
@@ -416,10 +453,13 @@ func removeCalendarEventsFromPassingHosts(
 	hosts []fleet.HostPolicyMembershipData,
 ) error {
 	for _, host := range hosts {
-		calendarEvent, err := ds.GetCalendarEvent(ctx, host.Email)
+		hostCalendarEvent, calendarEvent, err := ds.GetHostCalendarEventByEmail(ctx, host.Email)
 		switch {
 		case err == nil:
-			// OK
+			if hostCalendarEvent.HostID != host.HostID {
+				// This calendar event belongs to another host, thus we skip this entry.
+				continue
+			}
 		case fleet.IsNotFound(err):
 			continue
 		default:

server/datastore/mysql/calendar_events.go

@@ -167,6 +167,30 @@ func (ds *Datastore) GetHostCalendarEvent(ctx context.Context, hostID uint) (*fl
 	return &hostCalendarEvent, &calendarEvent, nil
 }
 
+func (ds *Datastore) GetHostCalendarEventByEmail(ctx context.Context, email string) (*fleet.HostCalendarEvent, *fleet.CalendarEvent, error) {
+	const calendarEventsQuery = `
+		SELECT * FROM calendar_events WHERE email = ?
+	`
+	var calendarEvent fleet.CalendarEvent
+	if err := sqlx.GetContext(ctx, ds.reader(ctx), &calendarEvent, calendarEventsQuery, email); err != nil {
+		if err == sql.ErrNoRows {
+			return nil, nil, ctxerr.Wrap(ctx, notFound("CalendarEvent").WithMessage(fmt.Sprintf("email: %s", email)))
+		}
+		return nil, nil, ctxerr.Wrap(ctx, err, "get calendar event")
+	}
+	const hostCalendarEventsQuery = `
+		SELECT * FROM host_calendar_events WHERE calendar_event_id = ?
+	`
+	var hostCalendarEvent fleet.HostCalendarEvent
+	if err := sqlx.GetContext(ctx, ds.reader(ctx), &hostCalendarEvent, hostCalendarEventsQuery, calendarEvent.ID); err != nil {
+		if err == sql.ErrNoRows {
+			return nil, nil, ctxerr.Wrap(ctx, notFound("HostCalendarEvent").WithID(calendarEvent.ID))
+		}
+		return nil, nil, ctxerr.Wrap(ctx, err, "get host calendar event")
+	}
+	return &hostCalendarEvent, &calendarEvent, nil
+}
+
 func (ds *Datastore) UpdateHostCalendarWebhookStatus(ctx context.Context, hostID uint, status fleet.CalendarWebhookStatus) error {
 	const calendarEventsQuery = `
 		UPDATE host_calendar_events SET

server/datastore/mysql/policies.go

@@ -1172,8 +1172,12 @@ func (ds *Datastore) GetCalendarPolicies(ctx context.Context, teamID uint) ([]fl
 }
 
 // TODO(lucas): Must be tested at scale.
-// TODO(lucas): Filter out hosts with team_id == NULL
-func (ds *Datastore) GetHostsPolicyMemberships(ctx context.Context, domain string, policyIDs []uint) ([]fleet.HostPolicyMembershipData, error) {
+func (ds *Datastore) GetTeamHostsPolicyMemberships(
+	ctx context.Context,
+	domain string,
+	teamID uint,
+	policyIDs []uint,
+) ([]fleet.HostPolicyMembershipData, error) {
 	query := `
 	SELECT 
 		COALESCE(sh.email, '') AS email,
@@ -1188,18 +1192,17 @@ func (ds *Datastore) GetHostsPolicyMemberships(ctx context.Context, domain strin
 		GROUP BY host_id
 	) pm
 	LEFT JOIN (
-		SELECT MIN(h.host_id) as host_id, h.email as email
-		FROM (
-			SELECT host_id, MIN(email) AS email
-			FROM host_emails WHERE email LIKE CONCAT('%@', ?)
-			GROUP BY host_id
-		) h GROUP BY h.email
+		SELECT host_id, MIN(email) AS email
+		FROM host_emails
+		JOIN hosts ON host_emails.host_id=hosts.id
+		WHERE email LIKE CONCAT('%@', ?) AND team_id = ? 
+		GROUP BY host_id
 	) sh ON sh.host_id = pm.host_id
 	JOIN hosts h ON h.id = pm.host_id
 	LEFT JOIN host_display_names hdn ON hdn.host_id = pm.host_id;
 `
 
-	query, args, err := sqlx.In(query, policyIDs, domain)
+	query, args, err := sqlx.In(query, policyIDs, domain, teamID)
 	if err != nil {
 		return nil, ctxerr.Wrapf(ctx, err, "build select get team hosts policy memberships query")
 	}