Alter checkin API to remove attributes set by audit/unenroll API (#3827)

Alter the checkin wrapper's flush method to include bulk update calls to remove the attributes for agents that have a non-empty audit_unenrolled_reason.

Author: michel-laterman

changelog/fragments/1724283895-Checkin-after-audit-unenroll-API.yaml

@@ -0,0 +1,34 @@
+# 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: enhancement
+
+# Change summary; a 80ish characters long description of the change.
+summary: Checkin after audit/unenroll API
+
+# 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: |
+  Fix the behaviour of the checkin handler if the audit/unenroll endpoint was previously called.
+  Calling the checkin handler will now result in the removal of the audit_unenroll and unenroll_time attributes from the agent doc.
+
+# Affected component; a word indicating the component this changeset affects.
+component:
+
+# 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/3827
+
+# 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/484

internal/pkg/api/handleCheckin.go

@@ -317,7 +317,7 @@ func (ct *CheckinT) ProcessRequest(zlog zerolog.Logger, w http.ResponseWriter, r
 	defer longPoll.Stop()
 
 	// Initial update on checkin, and any user fields that might have changed
-	err = ct.bc.CheckIn(agent.Id, string(req.Status), req.Message, rawMeta, rawComponents, seqno, ver, unhealthyReason)
+	err = ct.bc.CheckIn(agent.Id, string(req.Status), req.Message, rawMeta, rawComponents, seqno, ver, unhealthyReason, agent.AuditUnenrolledReason != "")
 	if err != nil {
 		zlog.Error().Err(err).Str(logger.AgentID, agent.Id).Msg("checkin failed")
 	}
@@ -371,7 +371,7 @@ func (ct *CheckinT) ProcessRequest(zlog zerolog.Logger, w http.ResponseWriter, r
 				zlog.Trace().Msg("fire long poll")
 				break LOOP
 			case <-tick.C:
-				err := ct.bc.CheckIn(agent.Id, string(req.Status), req.Message, nil, rawComponents, nil, ver, unhealthyReason)
+				err := ct.bc.CheckIn(agent.Id, string(req.Status), req.Message, nil, rawComponents, nil, ver, unhealthyReason, false)
 				if err != nil {
 					zlog.Error().Err(err).Str(logger.AgentID, agent.Id).Msg("checkin failed")
 				}

internal/pkg/checkin/bulk.go

@@ -7,19 +7,27 @@ package checkin
 
 import (
 	"context"
+	_ "embed"
 	"encoding/json"
+	"errors"
+	"fmt"
 	"sync"
 	"time"
 
 	"github.com/elastic/fleet-server/v7/internal/pkg/bulk"
 	"github.com/elastic/fleet-server/v7/internal/pkg/dl"
 	"github.com/elastic/fleet-server/v7/internal/pkg/sqn"
 
+	estypes "github.com/elastic/go-elasticsearch/v8/typedapi/types"
+	"github.com/elastic/go-elasticsearch/v8/typedapi/types/enums/scriptlanguage"
 	"github.com/rs/zerolog"
 )
 
 const defaultFlushInterval = 10 * time.Second
 
+//go:embed deleteAuditFieldsOnCheckin.painless
+var deleteAuditAttributesScript string
+
 type optionsT struct {
 	flushInterval time.Duration
 }
@@ -33,10 +41,11 @@ func WithFlushInterval(d time.Duration) Opt {
 }
 
 type extraT struct {
-	meta       []byte
-	seqNo      sqn.SeqNo
-	ver        string
-	components []byte
+	meta        []byte
+	seqNo       sqn.SeqNo
+	ver         string
+	components  []byte
+	deleteAudit bool
 }
 
 // Minimize the size of this structure.
@@ -102,17 +111,18 @@ func (bc *Bulk) timestamp() string {
 // The pending agents are sent to elasticsearch as a bulk update at each flush interval.
 // NOTE: If Checkin is called after Run has returned it will just add the entry to the pending map and not do any operations, this may occur when the fleet-server is shutting down.
 // WARNING: Bulk will take ownership of fields, so do not use after passing in.
-func (bc *Bulk) CheckIn(id string, status string, message string, meta []byte, components []byte, seqno sqn.SeqNo, newVer string, unhealthyReason *[]string) error {
+func (bc *Bulk) CheckIn(id string, status string, message string, meta []byte, components []byte, seqno sqn.SeqNo, newVer string, unhealthyReason *[]string, deleteAudit bool) error {
 	// Separate out the extra data to minimize
 	// the memory footprint of the 90% case of just
 	// updating the timestamp.
 	var extra *extraT
-	if meta != nil || seqno.IsSet() || newVer != "" || components != nil {
+	if meta != nil || seqno.IsSet() || newVer != "" || components != nil || deleteAudit {
 		extra = &extraT{
-			meta:       meta,
-			seqNo:      seqno,
-			ver:        newVer,
-			components: components,
+			meta:        meta,
+			seqNo:       seqno,
+			ver:         newVer,
+			components:  components,
+			deleteAudit: deleteAudit,
 		}
 	}
 
@@ -182,7 +192,6 @@ func (bc *Bulk) flush(ctx context.Context) error {
 		// JSON body containing just the timestamp updates.
 		var body []byte
 		if pendingData.extra == nil {
-
 			var ok bool
 			body, ok = simpleCache[pendingData]
 			if !ok {
@@ -198,8 +207,27 @@ func (bc *Bulk) flush(ctx context.Context) error {
 				}
 				simpleCache[pendingData] = body
 			}
+		} else if pendingData.extra.deleteAudit {
+			// Use a script instead of a partial doc to update if attributes need to be removed
+			params, err := encodeParams(nowTimestamp, pendingData)
+			if err != nil {
+				return err
+			}
+			action := &estypes.UpdateAction{
+				Script: &estypes.InlineScript{
+					Lang:   &scriptlanguage.Painless,
+					Source: deleteAuditAttributesScript,
+					Params: params,
+				},
+			}
+			body, err = json.Marshal(&action)
+			if err != nil {
+				return fmt.Errorf("could not marshall script action: %w", err)
+			}
+			if pendingData.extra.seqNo.IsSet() {
+				needRefresh = true
+			}
 		} else {
-
 			fields := bulk.UpdateFields{
 				dl.FieldLastCheckin:        pendingData.ts,      // Set the checkin timestamp
 				dl.FieldUpdatedAt:          nowTimestamp,        // Set "updated_at" to the current timestamp
@@ -265,3 +293,58 @@ func (bc *Bulk) flush(ctx context.Context) error {
 
 	return err
 }
+
+func encodeParams(now string, data pendingT) (map[string]json.RawMessage, error) {
+	var (
+		tsNow      json.RawMessage
+		ts         json.RawMessage
+		status     json.RawMessage
+		message    json.RawMessage
+		reason     json.RawMessage
+		ver        json.RawMessage
+		meta       json.RawMessage
+		components json.RawMessage
+		isSet      json.RawMessage
+		seqNo      json.RawMessage
+		err        error
+	)
+	tsNow, err = json.Marshal(now)
+	Err := errors.Join(err)
+	ts, err = json.Marshal(data.ts)
+	Err = errors.Join(Err, err)
+	status, err = json.Marshal(data.status)
+	Err = errors.Join(Err, err)
+	message, err = json.Marshal(data.message)
+	Err = errors.Join(Err, err)
+	reason, err = json.Marshal(data.unhealthyReason)
+	Err = errors.Join(Err, err)
+	ver, err = json.Marshal(data.extra.ver)
+	Err = errors.Join(Err, err)
+	isSet, err = json.Marshal(data.extra.seqNo.IsSet())
+	Err = errors.Join(Err, err)
+	seqNo, err = json.Marshal(data.extra.seqNo)
+	Err = errors.Join(Err, err)
+	if data.extra.meta != nil {
+		meta, err = json.Marshal(data.extra.meta)
+		Err = errors.Join(Err, err)
+	}
+	if data.extra.components != nil {
+		components, err = json.Marshal(data.extra.components)
+		Err = errors.Join(Err, err)
+	}
+	if Err != nil {
+		return nil, Err
+	}
+	return map[string]json.RawMessage{
+		"Now":             tsNow,
+		"TS":              ts,
+		"Status":          status,
+		"Message":         message,
+		"UnhealthyReason": reason,
+		"Ver":             ver,
+		"Meta":            meta,
+		"Components":      components,
+		"SeqNoSet":        isSet,
+		"SeqNo":           seqNo,
+	}, nil
+}

internal/pkg/checkin/bulk_test.go

@@ -193,7 +193,7 @@ func TestBulkSimple(t *testing.T) {
 			mockBulk.On("MUpdate", mock.Anything, mock.MatchedBy(matchOp(t, c, start)), mock.Anything).Return([]bulk.BulkIndexerResponseItem{}, nil).Once()
 			bc := NewBulk(mockBulk)
 
-			if err := bc.CheckIn(c.id, c.status, c.message, c.meta, c.components, c.seqno, c.ver, c.unhealthyReason); err != nil {
+			if err := bc.CheckIn(c.id, c.status, c.message, c.meta, c.components, c.seqno, c.ver, c.unhealthyReason, false); err != nil {
 				t.Fatal(err)
 			}
 
@@ -228,7 +228,7 @@ func benchmarkBulk(n int, b *testing.B) {
 	b.ReportAllocs()
 	for i := 0; i < b.N; i++ {
 		for _, id := range ids {
-			err := bc.CheckIn(id, "", "", nil, nil, nil, "", nil)
+			err := bc.CheckIn(id, "", "", nil, nil, nil, "", nil, false)
 			if err != nil {
 				b.Fatal(err)
 			}
@@ -254,7 +254,7 @@ func benchmarkFlush(n int, b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		b.StopTimer()
 		for _, id := range ids {
-			err := bc.CheckIn(id, "", "", nil, nil, nil, "", nil)
+			err := bc.CheckIn(id, "", "", nil, nil, nil, "", nil, false)
 			if err != nil {
 				b.Fatal(err)
 			}

internal/pkg/checkin/deleteAuditFieldsOnCheckin.painless

@@ -0,0 +1,20 @@
+ctx._source.last_checkin = params.TS;
+ctx._source.updated_at = params.Now;
+ctx._source.last_checkin_status = params.Status;
+ctx._source.last_checkin_message = params.Message;
+ctx._source.unhealthy_reason = params.UnhealthyReason;
+if (params.Ver != "") {
+  ctx._source.agent.version = params.Ver;
+}
+if (params.Meta != null) {
+  ctx._source.local_metadata = params.Meta;
+}
+if (params.Components != null) {
+  ctx._source.components = params.Components;
+}
+if (params.SeqNoSet) {
+    ctx._source.action_seq_no = params.SeqNo;
+}
+ctx._source.remove('audit_unenrolled_reason');
+ctx._source.remove('audit_unenrolled_time');
+ctx._source.remove('dl.unenrolled_at');

internal/pkg/server/fleet_integration_test.go

@@ -1490,4 +1490,46 @@ func Test_SmokeTest_AuditUnenroll(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, http.StatusConflict, res.StatusCode)
 	res.Body.Close()
+
+	t.Logf("Fake a checkin for agent %s", id)
+	req, err = http.NewRequestWithContext(ctx, "POST", srv.baseURL()+"/api/fleet/agents/"+id+"/checkin", strings.NewReader(checkinBody))
+	require.NoError(t, err)
+	req.Header.Set("Authorization", "ApiKey "+key)
+	req.Header.Set("User-Agent", "elastic agent "+serverVersion)
+	req.Header.Set("Content-Type", "application/json")
+	res, err = cli.Do(req)
+	require.NoError(t, err)
+
+	require.Equal(t, http.StatusOK, res.StatusCode)
+	t.Log("Checkin successful, verify body")
+	p, _ := io.ReadAll(res.Body)
+	res.Body.Close()
+	var obj map[string]interface{}
+	err = json.Unmarshal(p, &obj)
+	require.NoError(t, err)
+
+	require.Eventually(t, func() bool {
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://localhost:9200/.fleet-agents/_doc/"+id, nil)
+		require.NoError(t, err)
+		req.SetBasicAuth("elastic", "changeme")
+		res, err := cli.Do(req)
+		require.NoError(t, err)
+		defer res.Body.Close()
+		if res.StatusCode != http.StatusOK {
+			return false
+		}
+		p, err := io.ReadAll(res.Body)
+		require.NoError(t, err)
+		var tmp map[string]interface{}
+		err = json.Unmarshal(p, &tmp)
+		require.NoError(t, err)
+		o, ok := tmp["_source"]
+		require.Truef(t, ok, "expected to find _source in: %v", tmp)
+		obj, ok := o.(map[string]interface{})
+		require.Truef(t, ok, "expected _source to be an object, was: %T", o)
+		_, ok = obj["audit_unenrolled_reason"]
+		return !ok
+	}, time.Second*20, time.Second, "agent document should not have audit_unenrolled_reason attribute")
+	cancel()
+	srv.waitExit() //nolint:errcheck // test case
 }

testing/e2e/api_version/client_api_current.go

@@ -394,4 +394,28 @@ func (tester *ClientAPITester) TestEnrollAuditUnenroll() {
 	tester.T().Logf("audit/unenroll endpoint for agent: %s should return conflict", agentID)
 	status = tester.AuditUnenroll(ctx, agentKey, agentID, api.Uninstall, now)
 	tester.Require().Equal(http.StatusConflict, status)
+
+	tester.T().Logf("test checkin agent: %s", agentID)
+	_, _, statusCode := tester.Checkin(ctx, agentKey, agentID, nil, nil, nil)
+	tester.Require().Equal(http.StatusOK, statusCode, "Expected status code 200 for successful checkin")
+
+	// verify that audit_unenrolled_reason attribute does not exist in agent doc
+	tester.Require().Eventually(func() bool {
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://"+tester.ESHosts+"/.fleet-agents/_doc/"+agentID, nil)
+		tester.Require().NoError(err)
+		req.SetBasicAuth(tester.ElasticUser, tester.ElasticPass)
+		res, err := tester.Client.Do(req)
+		tester.Require().NoError(err)
+		defer res.Body.Close()
+		if res.StatusCode != http.StatusOK {
+			return false
+		}
+		var obj struct {
+			Source map[string]interface{} `json:"_source"`
+		}
+		err = json.NewDecoder(res.Body).Decode(&obj)
+		tester.Require().NoError(err)
+		_, ok := obj.Source["audit_unenrolled_reason"]
+		return !ok
+	}, time.Second*20, time.Second, "agent document in elasticsearch should not have audit_unenrolled_reason attribute")
 }