Handle when agent checks in with revision idx that is too high

Handle the scenario when the agent checks in with a revision_idx value
that is greater than the latest available policy in ES. Add E2E tests
when using policy_id and revision_idx values in checkin.

Author: michel-laterman

internal/pkg/api/handleCheckin.go

@@ -1144,7 +1144,7 @@ func calcPollDuration(zlog zerolog.Logger, pollDuration, setupDuration, jitterDu
 }
 
 // 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.
+// The API keys will be managed if the agent reports a new policy id from its last checkin, or if the revision is different 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
@@ -1162,7 +1162,7 @@ func (ct *CheckinT) processPolicyDetails(ctx context.Context, zlog zerolog.Logge
 
 	// update agent doc if policy id or revision idx does not match
 	var opts []checkin.Option
-	if policyID != agent.PolicyID || revisionIDX != agent.PolicyRevisionIdx {
+	if policyID != agent.AgentPolicyID || revisionIDX != agent.PolicyRevisionIdx {
 		opts = []checkin.Option{
 			checkin.WithAgentPolicyID(policyID),
 			checkin.WithPolicyRevisionIDX(revisionIDX),
@@ -1174,8 +1174,14 @@ func (ct *CheckinT) processPolicyDetails(ctx context.Context, zlog zerolog.Logge
 		return 0, opts, nil
 	}
 
-	// Update API keys if the policy has changed, or if the revision increments.
-	if policyID != agent.AgentPolicyID || revisionIDX > agent.PolicyRevisionIdx {
+	// Check if the checkin revision_idx is greater than the latest available
+	latestRev := ct.pm.LatestRev(ctx, agent.PolicyID)
+	if latestRev != 0 && revisionIDX > latestRev {
+		return 0, opts, nil
+	}
+
+	// Update API keys if the policy has changed, or if the revision differs.
+	if policyID != agent.AgentPolicyID || revisionIDX != agent.PolicyRevisionIdx {
 		for outputName, output := range agent.Outputs {
 			if output.Type != policy.OutputTypeElasticsearch {
 				continue

internal/pkg/api/handleCheckin_test.go

@@ -25,7 +25,6 @@ import (
 	"github.com/elastic/fleet-server/v7/internal/pkg/dl"
 	"github.com/elastic/fleet-server/v7/internal/pkg/es"
 	"github.com/elastic/fleet-server/v7/internal/pkg/model"
-	mockmonitor "github.com/elastic/fleet-server/v7/internal/pkg/monitor/mock"
 	"github.com/elastic/fleet-server/v7/internal/pkg/policy"
 	"github.com/elastic/fleet-server/v7/internal/pkg/sqn"
 	ftesting "github.com/elastic/fleet-server/v7/internal/pkg/testing"
@@ -39,6 +38,30 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+type mockPolicyMonitor struct {
+	mock.Mock
+}
+
+func (m *mockPolicyMonitor) Run(ctx context.Context) error {
+	args := m.Called(ctx)
+	return args.Error(0)
+}
+
+func (m *mockPolicyMonitor) Subscribe(agentID, policyID string, revIDX int64) (policy.Subscription, error) {
+	args := m.Called(agentID, policyID, revIDX)
+	return args.Get(0).(policy.Subscription), args.Error(1)
+}
+
+func (m *mockPolicyMonitor) Unsubscribe(sub policy.Subscription) error {
+	args := m.Called(sub)
+	return args.Error(0)
+}
+
+func (m *mockPolicyMonitor) LatestRev(ctx context.Context, id string) int64 {
+	args := m.Called(ctx, id)
+	return args.Get(0).(int64)
+}
+
 func TestConvertActionData(t *testing.T) {
 	tests := []struct {
 		name   string
@@ -339,14 +362,13 @@ func TestResolveSeqNo(t *testing.T) {
 			cfg := &config.Server{}
 			c, _ := cache.New(config.Cache{NumCounters: 100, MaxCost: 100000})
 			bc := checkin.NewBulk(nil)
-			bulker := ftesting.NewMockBulk()
-			pim := mockmonitor.NewMockMonitor()
-			pm := policy.NewMonitor(bulker, pim, config.ServerLimits{PolicyLimit: config.Limit{Interval: 5 * time.Millisecond, Burst: 1}})
+			pm := &mockPolicyMonitor{}
 			ct, err := NewCheckinT(verCon, cfg, c, bc, pm, nil, nil, nil)
 			assert.NoError(t, err)
 
 			resp, _ := ct.resolveSeqNo(ctx, logger, tc.req, tc.agent)
 			assert.Equal(t, tc.resp, resp)
+			pm.AssertExpectations(t)
 		})
 	}
 
@@ -1123,18 +1145,22 @@ 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             string
+		agent            *model.Agent
+		req              *CheckinRequest
+		getPolicyMonitor func() *mockPolicyMonitor
+		revIDX           int64
+		returnsOps       bool
+		err              error
 	}{{
 		name: "request has no policy details",
 		agent: &model.Agent{
 			PolicyRevisionIdx: 1,
 		},
-		req:        &CheckinRequest{},
+		req: &CheckinRequest{},
+		getPolicyMonitor: func() *mockPolicyMonitor {
+			return &mockPolicyMonitor{}
+		},
 		revIDX:     1,
 		returnsOps: false,
 		err:        nil,
@@ -1152,8 +1178,11 @@ func TestProcessPolicyDetails(t *testing.T) {
 			AgentPolicyId:     &policyID,
 			PolicyRevisionIdx: &revIDX2,
 		},
+		getPolicyMonitor: func() *mockPolicyMonitor {
+			return &mockPolicyMonitor{}
+		},
 		revIDX:     0,
-		returnsOps: true,
+		returnsOps: false,
 		err:        nil,
 	}, {
 		name: "revision updated",
@@ -1169,9 +1198,36 @@ func TestProcessPolicyDetails(t *testing.T) {
 			AgentPolicyId:     &policyID,
 			PolicyRevisionIdx: &revIDX2,
 		},
+		getPolicyMonitor: func() *mockPolicyMonitor {
+			pm := &mockPolicyMonitor{}
+			pm.On("LatestRev", mock.Anything, policyID).Return(int64(2)).Once()
+			return pm
+		},
 		revIDX:     2,
 		returnsOps: true,
 		err:        nil,
+	}, {
+		name: "checkin revision is greater than the policy's latest revision",
+		agent: &model.Agent{
+			Agent: &model.AgentMetadata{
+				ID: "agent-id",
+			},
+			PolicyID:          policyID,
+			AgentPolicyID:     policyID,
+			PolicyRevisionIdx: 1,
+		},
+		req: &CheckinRequest{
+			AgentPolicyId:     &policyID,
+			PolicyRevisionIdx: &revIDX2,
+		},
+		getPolicyMonitor: func() *mockPolicyMonitor {
+			pm := &mockPolicyMonitor{}
+			pm.On("LatestRev", mock.Anything, policyID).Return(int64(1)).Once()
+			return pm
+		},
+		revIDX:     0,
+		returnsOps: true,
+		err:        nil,
 	}, {
 		name: "agent_policy_id has changed",
 		agent: &model.Agent{
@@ -1186,6 +1242,11 @@ func TestProcessPolicyDetails(t *testing.T) {
 			AgentPolicyId:     &policyID,
 			PolicyRevisionIdx: &revIDX2,
 		},
+		getPolicyMonitor: func() *mockPolicyMonitor {
+			pm := &mockPolicyMonitor{}
+			pm.On("LatestRev", mock.Anything, policyID).Return(int64(2)).Once()
+			return pm
+		},
 		revIDX:     2,
 		returnsOps: true,
 		err:        nil,
@@ -1202,8 +1263,13 @@ func TestProcessPolicyDetails(t *testing.T) {
 			AgentPolicyId:     &policyID,
 			PolicyRevisionIdx: &revIDX2,
 		},
+		getPolicyMonitor: func() *mockPolicyMonitor {
+			pm := &mockPolicyMonitor{}
+			pm.On("LatestRev", mock.Anything, policyID).Return(int64(2)).Once()
+			return pm
+		},
 		revIDX:     2,
-		returnsOps: false,
+		returnsOps: true,
 		err:        nil,
 	}, {
 		name: "details present with no changes from agent doc",
@@ -1219,6 +1285,11 @@ func TestProcessPolicyDetails(t *testing.T) {
 			AgentPolicyId:     &policyID,
 			PolicyRevisionIdx: &revIDX2,
 		},
+		getPolicyMonitor: func() *mockPolicyMonitor {
+			pm := &mockPolicyMonitor{}
+			pm.On("LatestRev", mock.Anything, policyID).Return(int64(2)).Once()
+			return pm
+		},
 		revIDX:     2,
 		returnsOps: false,
 		err:        nil,
@@ -1227,8 +1298,10 @@ func TestProcessPolicyDetails(t *testing.T) {
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
 			logger := testlog.SetLogger(t)
+			pm := tc.getPolicyMonitor()
 			checkin := &CheckinT{
 				bulker: ftesting.NewMockBulk(),
+				pm:     pm,
 			}
 
 			revIDX, opts, err := checkin.processPolicyDetails(t.Context(), logger, tc.agent, tc.req)
@@ -1243,6 +1316,7 @@ func TestProcessPolicyDetails(t *testing.T) {
 			} else {
 				assert.NoError(t, err)
 			}
+			pm.AssertExpectations(t)
 		})
 	}
 }

internal/pkg/api/handleStatus_test.go

@@ -34,15 +34,15 @@ func withAuthFunc(authfn AuthFunc) OptFunc {
 	}
 }
 
-type mockPolicyMonitor struct {
+type mockStatusPolicyMonitor struct {
 	state client.UnitState
 }
 
-func (pm *mockPolicyMonitor) Run(ctx context.Context) error {
+func (pm *mockStatusPolicyMonitor) Run(ctx context.Context) error {
 	return nil
 }
 
-func (pm *mockPolicyMonitor) State() client.UnitState {
+func (pm *mockStatusPolicyMonitor) State() client.UnitState {
 	return pm.state
 }
 
@@ -86,7 +86,7 @@ func TestHandleStatus(t *testing.T) {
 					ctx = logger.WithContext(ctx)
 					state := client.UnitState(k)
 					r := apiServer{
-						st: NewStatusT(cfg, nil, c, withAuthFunc(tc.AuthFn), WithSelfMonitor(&mockPolicyMonitor{state}), WithBuildInfo(fbuild.Info{
+						st: NewStatusT(cfg, nil, c, withAuthFunc(tc.AuthFn), WithSelfMonitor(&mockStatusPolicyMonitor{state}), WithBuildInfo(fbuild.Info{
 							Version:   "8.1.0",
 							Commit:    "4eff928",
 							BuildTime: time.Now(),

internal/pkg/policy/monitor.go

@@ -62,6 +62,9 @@ type Monitor interface {
 
 	// Unsubscribe removes the current subscription.
 	Unsubscribe(sub Subscription) error
+
+	// LatestRev returns the latest revision idx for the specified policy.
+	LatestRev(ctx context.Context, policyID string) int64
 }
 
 type policyFetcher func(ctx context.Context, bulker bulk.Bulk, opt ...dl.Option) ([]model.Policy, error)
@@ -557,3 +560,34 @@ func (m *monitorT) Unsubscribe(sub Subscription) error {
 
 	return nil
 }
+
+// LatestRev returns the revision_idx for the passed policy ID.
+// If the policy does not exist in the map, then all policies are foribly reloaded.
+// On an error with the reload, or if the policy does not exist a 0 is returned.
+func (m *monitorT) LatestRev(ctx context.Context, id string) int64 {
+	if id == "" {
+		return 0
+	}
+
+	m.mut.Lock()
+	p, ok := m.policies[id]
+	m.mut.Unlock()
+
+	if !ok {
+		// We've not seen this policy before, force load.
+		err := m.loadPolicies(ctx)
+		if err != nil {
+			m.log.Error().Err(err).Str(ecs.PolicyID, id).Msg("Unable to load policies.")
+			return 0
+		}
+
+		m.mut.Lock()
+		p, ok = m.policies[id]
+		m.mut.Unlock()
+		if !ok {
+			m.log.Warn().Str(ecs.PolicyID, id).Msg("Unable to find policy after load.")
+			return 0
+		}
+	}
+	return p.pp.Policy.RevisionIdx
+}

internal/pkg/policy/monitor_test.go

@@ -9,6 +9,7 @@ package policy
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"sync"
 	"testing"
 	"time"
@@ -549,3 +550,76 @@ LOOP:
 	ms.AssertExpectations(t)
 	mm.AssertExpectations(t)
 }
+
+func TestMonitor_LatestRev(t *testing.T) {
+	t.Run("empty policy id", func(t *testing.T) {
+		pm := &monitorT{}
+		idx := pm.LatestRev(t.Context(), "")
+		assert.Equal(t, int64(0), idx)
+	})
+
+	t.Run("policy load error", func(t *testing.T) {
+		bulker := ftesting.NewMockBulk()
+		mm := mmock.NewMockMonitor()
+		monitor := NewMonitor(bulker, mm, config.ServerLimits{})
+		pm := monitor.(*monitorT)
+		pm.policyF = func(ctx context.Context, bulker bulk.Bulk, opt ...dl.Option) ([]model.Policy, error) {
+			return nil, fmt.Errorf("policy fetch error")
+		}
+
+		idx := pm.LatestRev(t.Context(), "test-id")
+		assert.Equal(t, int64(0), idx)
+	})
+
+	t.Run("policy not found", func(t *testing.T) {
+		bulker := ftesting.NewMockBulk()
+		mm := mmock.NewMockMonitor()
+		monitor := NewMonitor(bulker, mm, config.ServerLimits{})
+		pm := monitor.(*monitorT)
+		pm.policyF = func(ctx context.Context, bulker bulk.Bulk, opt ...dl.Option) ([]model.Policy, error) {
+			return []model.Policy{}, nil
+		}
+		idx := pm.LatestRev(t.Context(), "test-id")
+		assert.Equal(t, int64(0), idx)
+	})
+
+	t.Run("policy found after load", func(t *testing.T) {
+		bulker := ftesting.NewMockBulk()
+		mm := mmock.NewMockMonitor()
+		monitor := NewMonitor(bulker, mm, config.ServerLimits{})
+		pm := monitor.(*monitorT)
+		policyId := uuid.Must(uuid.NewV4()).String()
+		rId := xid.New().String()
+		policy := model.Policy{
+			ESDocument: model.ESDocument{
+				Id:      rId,
+				Version: 1,
+				SeqNo:   1,
+			},
+			PolicyID:    policyId,
+			Data:        policyDataDefault,
+			RevisionIdx: 2,
+		}
+		pm.policyF = func(ctx context.Context, bulker bulk.Bulk, opt ...dl.Option) ([]model.Policy, error) {
+			return []model.Policy{policy}, nil
+		}
+		idx := pm.LatestRev(t.Context(), policyId)
+		assert.Equal(t, int64(2), idx)
+	})
+
+	t.Run("policy found", func(t *testing.T) {
+		pm := &monitorT{
+			policies: map[string]policyT{
+				"test-id": policyT{
+					pp: ParsedPolicy{
+						Policy: model.Policy{
+							RevisionIdx: 1,
+						},
+					},
+				},
+			},
+		}
+		idx := pm.LatestRev(t.Context(), "test-id")
+		assert.Equal(t, int64(1), idx)
+	})
+}