avoid new API key being marked for invalidation

internal/pkg/dl/constants.go

@@ -43,6 +43,7 @@ const (
 	FieldPolicyRevisionIdx             = "policy_revision_idx"
 	FieldRevisionIdx                   = "revision_idx"
 	FieldUnenrolledReason              = "unenrolled_reason"
+	FiledType                          = "type"
 
 	FieldActive           = "active"
 	FieldUpdatedAt        = "updated_at"

internal/pkg/policy/policy_output.go

@@ -74,8 +74,8 @@ func (p *Output) prepareElasticsearch(
 		return ErrNoOutputPerms
 	}
 
-	output, ok := agent.Outputs[p.Name]
-	if !ok {
+	output, foundOutput := agent.Outputs[p.Name]
+	if !foundOutput {
 		if agent.Outputs == nil {
 			agent.Outputs = map[string]*model.PolicyOutput{}
 		}
@@ -120,11 +120,6 @@ func (p *Output) prepareElasticsearch(
 			return fmt.Errorf("failed generate output API key: %w", err)
 		}
 
-		output.Type = OutputTypeElasticsearch
-		output.APIKey = outputAPIKey.Agent()
-		output.APIKeyID = outputAPIKey.ID
-		output.PermissionsHash = p.Role.Sha2 // for the sake of consistency
-
 		// When a new keys is generated we need to update the Agent record,
 		// this will need to be updated when multiples remote Elasticsearch output
 		// are supported.
@@ -138,6 +133,10 @@ func (p *Output) prepareElasticsearch(
 			dl.FieldPolicyOutputAPIKeyID:        outputAPIKey.ID,
 			dl.FieldPolicyOutputPermissionsHash: p.Role.Sha2,
 		}
+
+		if !foundOutput {
+			fields[dl.FiledType] = OutputTypeElasticsearch
+		}
 		if output.APIKeyID != "" {
 			fields[dl.FieldPolicyOutputToRetireAPIKeyIDs] = model.ToRetireAPIKeyIdsItems{
 				ID:        output.APIKeyID,
@@ -155,6 +154,15 @@ func (p *Output) prepareElasticsearch(
 			zlog.Error().Err(err).Msg("fail update agent record")
 			return fmt.Errorf("fail update agent record: %w", err)
 		}
+
+		// Now that all is done, we can update the output on the agent variable
+		// Right not it's more for consistency and to ensure the in-memory agent
+		// data is correct and in sync with ES, so it can be safely used after
+		// this method returns.
+		output.Type = OutputTypeElasticsearch
+		output.APIKey = outputAPIKey.Agent()
+		output.APIKeyID = outputAPIKey.ID
+		output.PermissionsHash = p.Role.Sha2 // for the sake of consistency
 	}
 
 	// Always insert the `api_key` as part of the output block, this is required

internal/pkg/policy/policy_output_integration_test.go

@@ -13,12 +13,14 @@ import (
 	"time"
 
 	"github.com/gofrs/uuid"
+	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"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/model"
+	"github.com/elastic/fleet-server/v7/internal/pkg/smap"
 	ftesting "github.com/elastic/fleet-server/v7/internal/pkg/testing"
 )
 
@@ -125,3 +127,74 @@ func TestRenderUpdatePainlessScript(t *testing.T) {
 		})
 	}
 }
+
+func TestPolicyOutputESPrepareRealES(t *testing.T) {
+	index, bulker := ftesting.SetupCleanIndex(context.Background(), t, dl.FleetAgents)
+
+	agentID := createAgent(t, index, bulker)
+	agent, err := dl.FindAgent(
+		context.Background(), bulker, dl.QueryAgentByID, dl.FieldID, agentID, dl.WithIndexName(index))
+	if err != nil {
+		require.NoError(t, err, "failed to find agent ID %q", agentID)
+	}
+
+	output := Output{
+		Type: OutputTypeElasticsearch,
+		Name: "test output",
+		Role: &RoleT{
+			Sha2: "new-hash",
+			Raw:  TestPayload,
+		},
+	}
+	policyMap := smap.Map{
+		"test output": map[string]interface{}{},
+	}
+
+	err = output.prepareElasticsearch(
+		context.Background(), zerolog.Nop(), bulker, &agent, policyMap)
+	require.NoError(t, err)
+
+	// need to wait a bit before querying the agent again
+	// TODO: find a better way to query the updated agent
+	time.Sleep(time.Second)
+
+	got, err := dl.FindAgent(
+		context.Background(), bulker, dl.QueryAgentByID, dl.FieldID, agentID, dl.WithIndexName(index))
+	if err != nil {
+		require.NoError(t, err, "failed to find agent ID %q", agentID)
+	}
+
+	gotOutput, ok := got.Outputs[output.Name]
+	require.True(t, ok, "no '%s' output fouled on agent document", output.Name)
+
+	assert.Empty(t, gotOutput.ToRetireAPIKeyIds)
+	assert.Equal(t, gotOutput.Type, OutputTypeElasticsearch)
+	assert.Equal(t, gotOutput.PermissionsHash, output.Role.Sha2)
+	assert.NotEmpty(t, gotOutput.APIKey)
+	assert.NotEmpty(t, gotOutput.APIKeyID)
+}
+
+func createAgent(t *testing.T, index string, bulker bulk.Bulk) string {
+	const nowStr = "2022-08-12T16:50:05Z"
+
+	agentID := uuid.Must(uuid.NewV4()).String()
+	policyID := uuid.Must(uuid.NewV4()).String()
+
+	agentModel := model.Agent{
+		PolicyID:          policyID,
+		Active:            true,
+		LastCheckin:       nowStr,
+		LastCheckinStatus: "",
+		UpdatedAt:         nowStr,
+		EnrolledAt:        nowStr,
+	}
+
+	body, err := json.Marshal(agentModel)
+	require.NoError(t, err)
+
+	_, err = bulker.Create(
+		context.Background(), index, agentID, body, bulk.WithRefresh())
+	require.NoError(t, err)
+
+	return agentID
+}