Fix and refactor the state store

internal/pkg/agent/application/actions/handlers/handler_action_cancel.go

@@ -37,11 +37,13 @@ func (h *Cancel) Handle(ctx context.Context, a fleetapi.Action, acker acker.Acke
 	if !ok {
 		return fmt.Errorf("invalid type, expected ActionCancel and received %T", a)
 	}
-	n := h.c.Cancel(action.TargetID)
+	n := h.c.Cancel(action.Data.TargetID)
 	if n == 0 {
-		h.log.Debugf("Cancel action id: %s target id: %s found no actions in queue.", action.ActionID, action.TargetID)
+		h.log.Debugf("Cancel action id: %s target id: %s found no actions in queue.",
+			action.ActionID, action.Data.TargetID)
 		return nil
 	}
-	h.log.Infof("Cancel action id: %s target id: %s removed %d action(s) from queue.", action.ActionID, action.TargetID, n)
+	h.log.Infof("Cancel action id: %s target id: %s removed %d action(s) from queue.",
+		action.ActionID, action.Data.TargetID, n)
 	return nil
 }

internal/pkg/agent/application/actions/handlers/handler_action_diagnostics.go

@@ -196,7 +196,7 @@ func (h *Diagnostics) runHooks(ctx context.Context, action *fleetapi.ActionDiagn
 	// Currently CPU is the only additional metric we can collect.
 	// If this changes we would need to change how we scan AdditionalMetrics.
 	collectCPU := false
-	for _, metric := range action.AdditionalMetrics {
+	for _, metric := range action.Data.AdditionalMetrics {
 		if metric == "CPU" {
 			h.log.Debug("Diagnostics will collect CPU profile.")
 			collectCPU = true
@@ -290,7 +290,7 @@ func (h *Diagnostics) diagComponents(ctx context.Context, action *fleetapi.Actio
 		h.log.Debugf("Component diagnostics complete. Took: %s", time.Since(startTime))
 	}()
 	additionalMetrics := []cproto.AdditionalDiagnosticRequest{}
-	for _, metric := range action.AdditionalMetrics {
+	for _, metric := range action.Data.AdditionalMetrics {
 		if metric == "CPU" {
 			additionalMetrics = append(additionalMetrics, cproto.AdditionalDiagnosticRequest_CPU)
 		}

internal/pkg/agent/application/actions/handlers/handler_action_policy_change.go

@@ -98,7 +98,7 @@ func (h *PolicyChangeHandler) Handle(ctx context.Context, a fleetapi.Action, ack
 	// // Cache signature validation key for the next policy handling
 	// h.signatureValidationKey = signatureValidationKey
 
-	c, err := config.NewConfigFrom(action.Policy)
+	c, err := config.NewConfigFrom(action.Data.Policy)
 	if err != nil {
 		return errors.New(err, "could not parse the configuration from the policy", errors.TypeConfig)
 	}

internal/pkg/agent/application/actions/handlers/handler_action_policy_change_test.go

@@ -44,7 +44,9 @@ func TestPolicyChange(t *testing.T) {
 		action := &fleetapi.ActionPolicyChange{
 			ActionID:   "abc123",
 			ActionType: "POLICY_CHANGE",
-			Policy:     conf,
+			Data: fleetapi.ActionPolicyChangeData{
+				Policy: conf,
+			},
 		}
 
 		cfg := configuration.DefaultConfiguration()
@@ -73,7 +75,9 @@ func TestPolicyAcked(t *testing.T) {
 		action := &fleetapi.ActionPolicyChange{
 			ActionID:   actionID,
 			ActionType: "POLICY_CHANGE",
-			Policy:     config,
+			Data: fleetapi.ActionPolicyChangeData{
+				Policy: config,
+			},
 		}
 
 		cfg := configuration.DefaultConfiguration()

internal/pkg/agent/application/actions/handlers/handler_action_settings.go

@@ -45,17 +45,18 @@ func (h *Settings) Handle(ctx context.Context, a fleetapi.Action, acker acker.Ac
 		return fmt.Errorf("invalid type, expected ActionSettings and received %T", a)
 	}
 
-	if !isSupportedLogLevel(action.LogLevel) {
-		return fmt.Errorf("invalid log level, expected debug|info|warning|error and received '%s'", action.LogLevel)
+	if !isSupportedLogLevel(action.Data.LogLevel) {
+		return fmt.Errorf("invalid log level, expected debug|info|warning|error and received '%s'",
+			action.Data.LogLevel)
 	}
 
 	lvl := logp.InfoLevel
-	err := lvl.Unpack(action.LogLevel)
+	err := lvl.Unpack(action.Data.LogLevel)
 	if err != nil {
 		return fmt.Errorf("failed to unpack log level: %w", err)
 	}
 
-	if err := h.agentInfo.SetLogLevel(ctx, action.LogLevel); err != nil {
+	if err := h.agentInfo.SetLogLevel(ctx, action.Data.LogLevel); err != nil {
 		return fmt.Errorf("failed to update log level: %w", err)
 	}
 

internal/pkg/agent/application/actions/handlers/handler_action_unenroll.go

@@ -22,11 +22,11 @@ const (
 )
 
 type stateStore interface {
-	Add(fleetapi.Action)
+	SetAction(fleetapi.Action)
 	AckToken() string
 	SetAckToken(ackToken string)
 	Save() error
-	Actions() []fleetapi.Action
+	Action() fleetapi.Action
 }
 
 // Unenroll results in  running agent entering idle state, non managed non standalone.
@@ -94,7 +94,7 @@ func (h *Unenroll) Handle(ctx context.Context, a fleetapi.Action, acker acker.Ac
 
 	if h.stateStore != nil {
 		// backup action for future start to avoid starting fleet gateway loop
-		h.stateStore.Add(a)
+		h.stateStore.SetAction(a)
 		if err := h.stateStore.Save(); err != nil {
 			h.log.Warnf("Failed to update state store: %v", err)
 		}

internal/pkg/agent/application/actions/handlers/handler_action_upgrade.go

@@ -74,9 +74,9 @@ func (h *Upgrade) Handle(ctx context.Context, a fleetapi.Action, ack acker.Acker
 	}
 
 	go func() {
-		h.log.Infof("starting upgrade to version %s in background", action.Version)
-		if err := h.coord.Upgrade(asyncCtx, action.Version, action.SourceURI, action, false, false); err != nil {
-			h.log.Errorf("upgrade to version %s failed: %v", action.Version, err)
+		h.log.Infof("starting upgrade to version %s in background", action.Data.Version)
+		if err := h.coord.Upgrade(asyncCtx, action.Data.Version, action.Data.SourceURI, action, false, false); err != nil {
+			h.log.Errorf("upgrade to version %s failed: %v", action.Data.Version, err)
 			// If context is cancelled in getAsyncContext, the actions are acked there
 			if !errors.Is(asyncCtx.Err(), context.Canceled) {
 				h.bkgMutex.Lock()
@@ -125,14 +125,17 @@ func (h *Upgrade) getAsyncContext(ctx context.Context, action fleetapi.Action, a
 		h.log.Errorf("invalid type, expected ActionUpgrade and received %T", action)
 		return nil, false
 	}
-	if (upgradeAction.Version == bkgAction.Version) && (upgradeAction.SourceURI == bkgAction.SourceURI) {
-		h.log.Infof("Duplicate upgrade to version %s received", bkgAction.Version)
+	if (upgradeAction.Data.Version == bkgAction.Data.Version) &&
+		(upgradeAction.Data.SourceURI == bkgAction.Data.SourceURI) {
+		h.log.Infof("Duplicate upgrade to version %s received",
+			bkgAction.Data.Version)
 		h.bkgActions = append(h.bkgActions, action)
 		return nil, false
 	}
 
 	// Versions must be different, cancel the first upgrade and run the new one
-	h.log.Infof("Canceling upgrade to version %s received", bkgAction.Version)
+	h.log.Infof("Canceling upgrade to version %s received",
+		bkgAction.Data.Version)
 	h.bkgCancel()
 
 	// Ack here because we have the lock, and we need to clear out the saved actions

internal/pkg/agent/application/actions/handlers/handler_action_upgrade_test.go

@@ -81,7 +81,8 @@ func TestUpgradeHandler(t *testing.T) {
 	go c.Run(ctx)
 
 	u := NewUpgrade(log, c)
-	a := fleetapi.ActionUpgrade{Version: "8.3.0", SourceURI: "http://localhost"}
+	a := fleetapi.ActionUpgrade{Data: fleetapi.ActionUpgradeData{
+		Version: "8.3.0", SourceURI: "http://localhost"}}
 	ack := noopacker.New()
 	err := u.Handle(ctx, &a, ack)
 	require.NoError(t, err)
@@ -114,7 +115,8 @@ func TestUpgradeHandlerSameVersion(t *testing.T) {
 	go c.Run(ctx)
 
 	u := NewUpgrade(log, c)
-	a := fleetapi.ActionUpgrade{Version: "8.3.0", SourceURI: "http://localhost"}
+	a := fleetapi.ActionUpgrade{Data: fleetapi.ActionUpgradeData{
+		Version: "8.3.0", SourceURI: "http://localhost"}}
 	ack := noopacker.New()
 	err1 := u.Handle(ctx, &a, ack)
 	err2 := u.Handle(ctx, &a, ack)
@@ -149,8 +151,10 @@ func TestUpgradeHandlerNewVersion(t *testing.T) {
 	go c.Run(ctx)
 
 	u := NewUpgrade(log, c)
-	a1 := fleetapi.ActionUpgrade{Version: "8.2.0", SourceURI: "http://localhost"}
-	a2 := fleetapi.ActionUpgrade{Version: "8.5.0", SourceURI: "http://localhost"}
+	a1 := fleetapi.ActionUpgrade{Data: fleetapi.ActionUpgradeData{
+		Version: "8.2.0", SourceURI: "http://localhost"}}
+	a2 := fleetapi.ActionUpgrade{Data: fleetapi.ActionUpgradeData{
+		Version: "8.5.0", SourceURI: "http://localhost"}}
 	ack := noopacker.New()
 	err1 := u.Handle(ctx, &a1, ack)
 	require.NoError(t, err1)

internal/pkg/agent/application/dispatcher/dispatcher.go

@@ -304,7 +304,7 @@ func (ad *ActionDispatcher) handleExpired(
 			version := "unknown"
 			expiration := "unknown"
 			if upgrade, ok := e.(*fleetapi.ActionUpgrade); ok {
-				version = upgrade.Version
+				version = upgrade.Data.Version
 				expiration = upgrade.ActionExpiration
 			}
 			ad.lastUpgradeDetails = details.NewDetails(version, details.StateFailed, e.ID())
@@ -356,7 +356,7 @@ func (ad *ActionDispatcher) reportNextScheduledUpgrade(input []fleetapi.Action,
 	}
 
 	upgradeDetails := details.NewDetails(
-		nextUpgrade.Version,
+		nextUpgrade.Data.Version,
 		details.StateScheduled,
 		nextUpgrade.ID())
 	startTime, err := nextUpgrade.StartTime()

internal/pkg/agent/application/dispatcher/dispatcher_test.go

@@ -675,7 +675,9 @@ func TestReportNextScheduledUpgrade(t *testing.T) {
 			actions: []fleetapi.Action{
 				&fleetapi.ActionUpgrade{
 					ActionID: "action1",
-					Version:  "8.12.3",
+					Data: fleetapi.ActionUpgradeData{
+						Version: "8.12.3",
+					},
 				},
 			},
 			expectedErrLogMsg: "failed to get start time for scheduled upgrade action [id = action1]",
@@ -685,7 +687,9 @@ func TestReportNextScheduledUpgrade(t *testing.T) {
 				&fleetapi.ActionUpgrade{
 					ActionID:        "action2",
 					ActionStartTime: later.Format(time.RFC3339),
-					Version:         "8.13.0",
+					Data: fleetapi.ActionUpgradeData{
+						Version: "8.13.0",
+					},
 				},
 			},
 			expectedDetails: &details.Details{
@@ -702,12 +706,16 @@ func TestReportNextScheduledUpgrade(t *testing.T) {
 				&fleetapi.ActionUpgrade{
 					ActionID:        "action3",
 					ActionStartTime: muchLater.Format(time.RFC3339),
-					Version:         "8.14.1",
+					Data: fleetapi.ActionUpgradeData{
+						Version: "8.14.1",
+					},
 				},
 				&fleetapi.ActionUpgrade{
 					ActionID:        "action4",
 					ActionStartTime: later.Format(time.RFC3339),
-					Version:         "8.13.5",
+					Data: fleetapi.ActionUpgradeData{
+						Version: "8.13.5",
+					},
 				},
 			},
 			expectedDetails: &details.Details{
@@ -723,8 +731,10 @@ func TestReportNextScheduledUpgrade(t *testing.T) {
 			actions: []fleetapi.Action{
 				&fleetapi.ActionUpgrade{
 					ActionID:        "action1",
-					Version:         "8.13.2",
 					ActionStartTime: "invalid",
+					Data: fleetapi.ActionUpgradeData{
+						Version: "8.13.2",
+					},
 				},
 			},
 			expectedErrLogMsg: "failed to get start time for scheduled upgrade action [id = action1]",

internal/pkg/agent/application/gateway/fleet/fleet_gateway.go

@@ -58,11 +58,9 @@ type agentInfo interface {
 }
 
 type stateStore interface {
-	Add(fleetapi.Action)
 	AckToken() string
 	SetAckToken(ackToken string)
 	Save() error
-	Actions() []fleetapi.Action
 }
 
 type FleetGateway struct {

internal/pkg/agent/application/managed_mode.go

@@ -157,14 +157,14 @@ func (m *managedConfigManager) Run(ctx context.Context) error {
 		close(retrierRun)
 	}()
 
-	actions := m.stateStore.Actions()
+	action := m.stateStore.Action()
 	stateRestored := false
-	if len(actions) > 0 && !m.wasUnenrolled() {
+	if action != nil && !m.wasUnenrolled() {
 		// TODO(ph) We will need an improvement on fleet, if there is an error while dispatching a
 		// persisted action on disk we should be able to ask Fleet to get the latest configuration.
 		// But at the moment this is not possible because the policy change was acked.
 		m.log.Info("restoring current policy from disk")
-		m.dispatcher.Dispatch(ctx, m.coord.SetUpgradeDetails, actionAcker, actions...)
+		m.dispatcher.Dispatch(ctx, m.coord.SetUpgradeDetails, actionAcker, action)
 		stateRestored = true
 	}
 
@@ -268,13 +268,8 @@ func (m *managedConfigManager) Watch() <-chan coordinator.ConfigChange {
 }
 
 func (m *managedConfigManager) wasUnenrolled() bool {
-	actions := m.stateStore.Actions()
-	for _, a := range actions {
-		if a.Type() == "UNENROLL" {
-			return true
-		}
-	}
-	return false
+	return m.stateStore.Action() != nil &&
+		m.stateStore.Action().Type() == fleetapi.ActionTypeUnenroll
 }
 
 func (m *managedConfigManager) initFleetServer(ctx context.Context, cfg *configuration.FleetServerConfig) error {

internal/pkg/agent/application/paths/files.go

@@ -23,13 +23,20 @@ const defaultAgentFleetFile = "fleet.enc"
 // defaultAgentEnrollFile is a name of file used to enroll agent on first-start
 const defaultAgentEnrollFile = "enroll.yml"
 
-// defaultAgentActionStoreFile is the file that will contain the action that can be replayed after restart.
+// defaultAgentActionStoreFile is the file that will contain the action that can
+// be replayed after restart.
+// It's deprecated and kept for migration purposes.
+// Deprecated.
 const defaultAgentActionStoreFile = "action_store.yml"
 
-// defaultAgentStateStoreYmlFile is the file that will contain the action that can be replayed after restart.
+// defaultAgentStateStoreYmlFile is the file that will contain the action that
+// can be replayed after restart.
+// It's deprecated and kept for migration purposes.
+// Deprecated.
 const defaultAgentStateStoreYmlFile = "state.yml"
 
-// defaultAgentStateStoreFile is the file that will contain the action that can be replayed after restart encrypted.
+// defaultAgentStateStoreFile is the file that will contain the encrypted state
+// store.
 const defaultAgentStateStoreFile = "state.enc"
 
 // defaultInputDPath return the location of the inputs.d.

internal/pkg/agent/application/upgrade/step_mark.go

@@ -70,8 +70,8 @@ func convertToMarkerAction(a *fleetapi.ActionUpgrade) *MarkerActionUpgrade {
 	return &MarkerActionUpgrade{
 		ActionID:   a.ActionID,
 		ActionType: a.ActionType,
-		Version:    a.Version,
-		SourceURI:  a.SourceURI,
+		Version:    a.Data.Version,
+		SourceURI:  a.Data.SourceURI,
 	}
 }
 
@@ -82,8 +82,10 @@ func convertToActionUpgrade(a *MarkerActionUpgrade) *fleetapi.ActionUpgrade {
 	return &fleetapi.ActionUpgrade{
 		ActionID:   a.ActionID,
 		ActionType: a.ActionType,
-		Version:    a.Version,
-		SourceURI:  a.SourceURI,
+		Data: fleetapi.ActionUpgradeData{
+			Version:   a.Version,
+			SourceURI: a.SourceURI,
+		},
 	}
 }
 

internal/pkg/agent/cmd/run.go

@@ -205,7 +205,9 @@ func runElasticAgent(ctx context.Context, cancel context.CancelFunc, override cf
 	}
 
 	// the encrypted state does not exist but the unencrypted file does
-	err = migration.MigrateToEncryptedConfig(ctx, l, paths.AgentStateStoreYmlFile(), paths.AgentStateStoreFile())
+	err = migration.MigrateToEncryptedConfig(ctx, l,
+		paths.AgentStateStoreYmlFile(),
+		paths.AgentStateStoreFile())
 	if err != nil {
 		return errors.New(err, "error migrating agent state")
 	}

internal/pkg/agent/protection/action.go

@@ -21,7 +21,7 @@ var (
 )
 
 type fleetActionWithAgents struct {
-	ActionID         string          `json:"action_id"` // Note the action_id here, since the signed action uses action_id for id
+	ActionID         string          `json:"id"`
 	ActionType       string          `json:"type,omitempty"`
 	InputType        string          `json:"input_type,omitempty"`
 	Timestamp        string          `json:"@timestamp"`

internal/pkg/agent/protection/action_test.go

@@ -46,9 +46,7 @@ func signAction(action map[string]interface{}, emptyData bool, pk *ecdsa.Private
 		"data":      base64.StdEncoding.EncodeToString(payload),
 		"signature": base64.StdEncoding.EncodeToString(sig),
 	}
-	// Remap the action_id to id same way the fleet server does for checkins
-	action["id"] = action["action_id"]
-	delete(action, "action_id")
+
 	return action, nil
 }
 

internal/pkg/agent/storage/encrypted_disk_store.go

@@ -108,8 +108,13 @@ func (d *EncryptedDiskStore) ensureKey(ctx context.Context) error {
 	return nil
 }
 
-// Save will write the encrypted storage to disk.
-// Specifically it will write to a .tmp file then rotate the file to the target name to ensure that an error does not corrupt the previously written file.
+// Save will read 'in' and write its contents encrypted to disk.
+// If EncryptedDiskStore.Load() was called, the io.ReadCloser it returns MUST be
+// closed before Save() can be called. It is so because Save() writes to a .tmp
+// file then rotate the file to the target name to ensure that an error does not
+// corrupt the previously written file.
+// Specially on windows systems, if the original files is still open because of
+// Load(), Save() would fail.
 func (d *EncryptedDiskStore) Save(in io.Reader) error {
 	// Ensure has agent key
 	err := d.ensureKey(d.ctx)