Add agent_policy_id and policy_revision_idx to checkin attributes

Allow the agents to add their currently running policy_id and
revision_idx attributes to the checkin request bodies. These attributes,
if included and different from the agent doc will be used when updating
the agent doc in the pre-poll checkin. If the agent's policy id does not
match the expected policy id from the server a reassign is detected and
a new policy change action will be sent. If the checkin ID is greater
than what was previously recorded or the policy id changes from what was
previously recoreded, then the api keys will be managed.

Author: michel-laterman

changelog/fragments/1757633911-Add-agent_policy_id-and-policy_revision_idx-to-checkin-requests.yaml

@@ -0,0 +1,38 @@
+# Kind can be one of:
+# - breaking-change: a change to previously-documented behavior
+# - deprecation: functionality that is being removed in a later release
+# - bug-fix: fixes a problem in a previous version
+# - enhancement: extends functionality but does not break or fix existing behavior
+# - feature: new functionality
+# - known-issue: problems that we are aware of in a given version
+# - security: impacts on the security of a product or a user’s deployment.
+# - upgrade: important information for someone upgrading from a prior version
+# - other: does not fit into any of the other categories
+kind: feature
+
+# Change summary; a 80ish characters long description of the change.
+summary: Add agent_policy_id and policy_revision_idx to checkin requests
+
+# Long description; in case the summary is not enough to describe the change
+# this field accommodate a description without length limits.
+# NOTE: This field will be rendered only for breaking-change and known-issue kinds at the moment.
+description: |
+  Add the agent_policy_id and policy_revision_idx attributes to checkin
+  request bodies so an agent is able to inform fleet-server of its exact
+  policy. These details will replace the need for an ack on
+  policy_change actions, and will be used to determine when to send a
+  policy change when there is a new revision available, or when the
+  agent is reassigned to a different policy.
+
+# Affected component; usually one of "elastic-agent", "fleet-server", "filebeat", "metricbeat", "auditbeat", "all", etc.
+component: fleet-server
+
+# PR URL; optional; the PR number that added the changeset.
+# If not present is automatically filled by the tooling finding the PR where this changelog fragment has been added.
+# NOTE: the tooling supports backports, so it's able to fill the original PR number instead of the backport PR number.
+# Please provide it if you are adding a fragment for a different PR.
+pr: https://github.com/elastic/fleet-server/pull/5501
+
+# Issue URL; optional; the GitHub issue related to this changeset (either closes or is part of).
+# If not present is automatically filled by the tooling with the issue linked to the PR number.
+issue: https://github.com/elastic/elastic-agent/issues/6446

internal/pkg/api/handleAck.go

@@ -442,19 +442,28 @@ func (ack *AckT) updateAPIKey(ctx context.Context,
 	agentID string,
 	apiKeyID, permissionHash string,
 	toRetireAPIKeyIDs []model.ToRetireAPIKeyIdsItems, outputName string) error {
-	bulk := ack.bulk
+	return updateAPIKey(ctx, zlog, ack.bulk, agentID, apiKeyID, permissionHash, toRetireAPIKeyIDs, outputName)
+}
+
+func updateAPIKey(ctx context.Context,
+	zlog zerolog.Logger,
+	bulk bulk.Bulk,
+	agentID string,
+	apiKeyID, permissionHash string,
+	toRetireAPIKeyIDs []model.ToRetireAPIKeyIdsItems, outputName string) error {
 	// use output bulker if exists
+	outBulk := bulk
 	if outputName != "" {
-		outputBulk := ack.bulk.GetBulker(outputName)
+		outputBulk := bulk.GetBulker(outputName)
 		if outputBulk != nil {
 			zlog.Debug().Str(ecs.PolicyOutputName, outputName).Msg("Using output bulker in updateAPIKey")
-			bulk = outputBulk
+			outBulk = outputBulk
 		}
 	}
 	if apiKeyID != "" {
-		res, err := bulk.APIKeyRead(ctx, apiKeyID, true)
+		res, err := outBulk.APIKeyRead(ctx, apiKeyID, true)
 		if err != nil {
-			if isAgentActive(ctx, zlog, ack.bulk, agentID) {
+			if isAgentActive(ctx, zlog, outBulk, agentID) {
 				zlog.Warn().
 					Err(err).
 					Str(LogAPIKeyID, apiKeyID).
@@ -480,7 +489,7 @@ func (ack *AckT) updateAPIKey(ctx context.Context,
 					Str(LogAPIKeyID, apiKeyID).
 					Msg("Failed to cleanup roles")
 			} else if removedRolesCount > 0 {
-				if err := bulk.APIKeyUpdate(ctx, apiKeyID, permissionHash, clean); err != nil {
+				if err := outBulk.APIKeyUpdate(ctx, apiKeyID, permissionHash, clean); err != nil {
 					zlog.Error().Err(err).RawJSON("roles", clean).Str(LogAPIKeyID, apiKeyID).Str(ecs.PolicyOutputName, outputName).Msg("Failed to update API Key")
 				} else {
 					zlog.Debug().
@@ -493,7 +502,7 @@ func (ack *AckT) updateAPIKey(ctx context.Context,
 				}
 			}
 		}
-		ack.invalidateAPIKeys(ctx, zlog, toRetireAPIKeyIDs, apiKeyID)
+		invalidateAPIKeys(ctx, zlog, bulk, toRetireAPIKeyIDs, apiKeyID)
 	}
 
 	return nil

internal/pkg/api/handleCheckin.go

@@ -278,16 +278,34 @@ func (ct *CheckinT) ProcessRequest(zlog zerolog.Logger, w http.ResponseWriter, r
 		return fmt.Errorf("failed to update upgrade_details: %w", err)
 	}
 
+	initialOpts := []checkin.Option{
+		checkin.WithStatus(string(req.Status)),
+		checkin.WithMessage(req.Message),
+		checkin.WithMeta(rawMeta),
+		checkin.WithComponents(rawComponents),
+		checkin.WithSeqNo(seqno),
+		checkin.WithVer(ver),
+		checkin.WithUnhealthyReason(unhealthyReason),
+		checkin.WithDeleteAudit(agent.AuditUnenrolledReason != "" || agent.UnenrolledAt != ""),
+	}
+
+	revID, opts, err := ct.processPolicyDetails(r.Context(), zlog, agent, req)
+	if err != nil {
+		return fmt.Errorf("failed to update policy details: %w", err)
+	}
+	if len(opts) > 0 {
+		initialOpts = append(initialOpts, opts...)
+	}
+
 	// Subscribe to actions dispatcher
 	aSub := ct.ad.Subscribe(zlog, agent.Id, seqno)
 	defer ct.ad.Unsubscribe(zlog, aSub)
 	actCh := aSub.Ch()
 
-	// use revision_idx=0 if the agent has a single output where no API key is defined
-	// This will force the policy monitor to emit a new policy to regerate API keys
-	revID := agent.PolicyRevisionIdx
 	for _, output := range agent.Outputs {
 		if output.APIKey == "" {
+			// use revision_idx=0 if the agent has a single output where no API key is defined
+			// This will force the policy monitor to emit a new policy to regerate API keys
 			revID = 0
 			break
 		}
@@ -327,7 +345,7 @@ func (ct *CheckinT) ProcessRequest(zlog zerolog.Logger, w http.ResponseWriter, r
 	// Initial update on checkin, and any user fields that might have changed
 	// Run a script to remove audit_unenrolled_* and unenrolled_at attributes if one is set on checkin.
 	// 8.16.x releases would incorrectly set unenrolled_at
-	err = ct.bc.CheckIn(agent.Id, checkin.WithStatus(string(req.Status)), checkin.WithMessage(req.Message), checkin.WithMeta(rawMeta), checkin.WithComponents(rawComponents), checkin.WithSeqNo(seqno), checkin.WithVer(ver), checkin.WithUnhealthyReason(unhealthyReason), checkin.WithDeleteAudit(agent.AuditUnenrolledReason != "" || agent.UnenrolledAt != ""))
+	err = ct.bc.CheckIn(agent.Id, initialOpts...)
 	if err != nil {
 		zlog.Error().Err(err).Str(ecs.AgentID, agent.Id).Msg("checkin failed")
 	}
@@ -1123,3 +1141,49 @@ func calcPollDuration(zlog zerolog.Logger, pollDuration, setupDuration, jitterDu
 
 	return pollDuration, jitter
 }
+
+// processPolicyDetails handles the agent_policy_id and revision_idx included in the checkin request.
+// The API keys will be managed if the agent reports a new policy id from its last checkin, or if the revision is greater than what the last checkin reported.
+// It returns the revision idx that should be used when subscribing for new POLICY_CHANGE actons and optional args to use when doing the non-tick checkin.
+func (ct *CheckinT) processPolicyDetails(ctx context.Context, zlog zerolog.Logger, agent *model.Agent, req *CheckinRequest) (int64, []checkin.Option, error) {
+	// no details specified
+	if req == nil || req.PolicyRevisionIdx == nil || req.AgentPolicyId == nil {
+		return agent.PolicyRevisionIdx, nil, nil
+	}
+	policyID := *req.AgentPolicyId
+	revisionIDX := *req.PolicyRevisionIdx
+
+	span, ctx := apm.StartSpan(ctx, "Process policy details", "process")
+	span.Context.SetLabel("agent_id", agent.Agent.ID)
+	span.Context.SetLabel(dl.FieldAgentPolicyID, policyID)
+	span.Context.SetLabel(dl.FieldPolicyRevisionIdx, revisionIDX)
+	defer span.End()
+
+	// update agent doc if policy id or revision idx does not match
+	var opts []checkin.Option
+	if policyID != agent.PolicyID || revisionIDX != agent.PolicyRevisionIdx {
+		opts = []checkin.Option{
+			checkin.WithAgentPolicyID(policyID),
+			checkin.WithPolicyRevisionIDX(revisionIDX),
+		}
+	}
+	// Policy reassign, subscribe to policy with revision 0
+	if policyID != agent.PolicyID {
+		zlog.Debug().Str(dl.FieldAgentPolicyID, policyID).Str("new_policy_id", agent.PolicyID).Msg("Policy ID mismatch detected, reassigning agent.")
+		return 0, opts, nil
+	}
+
+	// Update API keys if the policy has changed, or if the revision increments.
+	if policyID != agent.AgentPolicyID || revisionIDX > agent.PolicyRevisionIdx {
+		for outputName, output := range agent.Outputs {
+			if output.Type != policy.OutputTypeElasticsearch {
+				continue
+			}
+			if err := updateAPIKey(ctx, zlog, ct.bulker, agent.Id, output.APIKeyID, output.PermissionsHash, output.ToRetireAPIKeyIds, outputName); err != nil {
+				// Only returns ErrUpdatingInactiveAgent
+				return 0, nil, err
+			}
+		}
+	}
+	return revisionIDX, opts, nil
+}

internal/pkg/api/handleCheckin_test.go

@@ -1118,3 +1118,131 @@ func TestValidateCheckinRequest(t *testing.T) {
 		})
 	}
 }
+
+func TestProcessPolicyDetails(t *testing.T) {
+	policyID := "policy-id"
+	revIDX2 := int64(2)
+	tests := []struct {
+		name       string
+		agent      *model.Agent
+		req        *CheckinRequest
+		revIDX     int64
+		returnsOps bool
+		err        error
+	}{{
+		name: "request has no policy details",
+		agent: &model.Agent{
+			PolicyRevisionIdx: 1,
+		},
+		req:        &CheckinRequest{},
+		revIDX:     1,
+		returnsOps: false,
+		err:        nil,
+	}, {
+		name: "policy reassign detected",
+		agent: &model.Agent{
+			Agent: &model.AgentMetadata{
+				ID: "agent-id",
+			},
+			PolicyID:          "new-policy-id",
+			AgentPolicyID:     policyID,
+			PolicyRevisionIdx: 2,
+		},
+		req: &CheckinRequest{
+			AgentPolicyId:     &policyID,
+			PolicyRevisionIdx: &revIDX2,
+		},
+		revIDX:     0,
+		returnsOps: true,
+		err:        nil,
+	}, {
+		name: "revision updated",
+		agent: &model.Agent{
+			Agent: &model.AgentMetadata{
+				ID: "agent-id",
+			},
+			PolicyID:          policyID,
+			AgentPolicyID:     policyID,
+			PolicyRevisionIdx: 1,
+		},
+		req: &CheckinRequest{
+			AgentPolicyId:     &policyID,
+			PolicyRevisionIdx: &revIDX2,
+		},
+		revIDX:     2,
+		returnsOps: true,
+		err:        nil,
+	}, {
+		name: "agent_policy_id has changed",
+		agent: &model.Agent{
+			Agent: &model.AgentMetadata{
+				ID: "agent-id",
+			},
+			PolicyID:          policyID,
+			AgentPolicyID:     "old-policy-id",
+			PolicyRevisionIdx: 1,
+		},
+		req: &CheckinRequest{
+			AgentPolicyId:     &policyID,
+			PolicyRevisionIdx: &revIDX2,
+		},
+		revIDX:     2,
+		returnsOps: true,
+		err:        nil,
+	}, {
+		name: "agent does not have agent_policy_id present",
+		agent: &model.Agent{
+			Agent: &model.AgentMetadata{
+				ID: "agent-id",
+			},
+			PolicyID:          policyID,
+			PolicyRevisionIdx: 2,
+		},
+		req: &CheckinRequest{
+			AgentPolicyId:     &policyID,
+			PolicyRevisionIdx: &revIDX2,
+		},
+		revIDX:     2,
+		returnsOps: false,
+		err:        nil,
+	}, {
+		name: "details present with no changes from agent doc",
+		agent: &model.Agent{
+			Agent: &model.AgentMetadata{
+				ID: "agent-id",
+			},
+			AgentPolicyID:     policyID,
+			PolicyID:          policyID,
+			PolicyRevisionIdx: revIDX2,
+		},
+		req: &CheckinRequest{
+			AgentPolicyId:     &policyID,
+			PolicyRevisionIdx: &revIDX2,
+		},
+		revIDX:     2,
+		returnsOps: false,
+		err:        nil,
+	}}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			logger := testlog.SetLogger(t)
+			checkin := &CheckinT{
+				bulker: ftesting.NewMockBulk(),
+			}
+
+			revIDX, opts, err := checkin.processPolicyDetails(t.Context(), logger, tc.agent, tc.req)
+			assert.Equal(t, tc.revIDX, revIDX)
+			if tc.returnsOps {
+				assert.NotEmpty(t, opts)
+			} else {
+				assert.Empty(t, opts)
+			}
+			if tc.err != nil {
+				assert.ErrorIs(t, tc.err, err)
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}