Add formal checks to manual rollback arguments

Author: pchila

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

@@ -57,11 +57,13 @@ var agentArtifact = artifact.Artifact{
 }
 
 var (
-	ErrWatcherNotStarted  = errors.New("watcher did not start in time")
-	ErrUpgradeSameVersion = errors.New("upgrade did not occur because it is the same version")
-	ErrNonFipsToFips      = errors.New("cannot switch to fips mode when upgrading")
-	ErrFipsToNonFips      = errors.New("cannot switch to non-fips mode when upgrading")
-	ErrNilUpdateMarker    = errors.New("loaded a nil update marker")
+	ErrWatcherNotStarted    = errors.New("watcher did not start in time")
+	ErrUpgradeSameVersion   = errors.New("upgrade did not occur because it is the same version")
+	ErrNonFipsToFips        = errors.New("cannot switch to fips mode when upgrading")
+	ErrFipsToNonFips        = errors.New("cannot switch to non-fips mode when upgrading")
+	ErrNilUpdateMarker      = errors.New("loaded a nil update marker")
+	ErrEmptyRollbackVersion = errors.New("rollback version is empty")
+	ErrNoRollbacksAvailable = errors.New("no rollbacks available")
 )
 
 func init() {
@@ -222,7 +224,7 @@ func checkUpgrade(log *logger.Logger, currentVersion, newVersion agentVersion, m
 func (u *Upgrader) Upgrade(ctx context.Context, version string, rollback bool, sourceURI string, action *fleetapi.ActionUpgrade, det *details.Details, skipVerifyOverride bool, skipDefaultPgp bool, pgpBytes ...string) (_ reexec.ShutdownCallbackFn, err error) {
 
 	if rollback {
-		return u.forceRollbackToPreviousVersion(ctx, paths.Top(), version, action, det)
+		return u.forceRollbackToPreviousVersion(ctx, paths.Top(), time.Now(), version, action)
 	}
 
 	u.log.Infow("Upgrading agent", "version", version, "source_uri", sourceURI)
@@ -406,7 +408,11 @@ func (u *Upgrader) Upgrade(ctx context.Context, version string, rollback bool, s
 	return cb, nil
 }
 
-func (u *Upgrader) forceRollbackToPreviousVersion(ctx context.Context, topDir string, version string, action *fleetapi.ActionUpgrade, d *details.Details) (reexec.ShutdownCallbackFn, error) {
+func (u *Upgrader) forceRollbackToPreviousVersion(ctx context.Context, topDir string, now time.Time, version string, action *fleetapi.ActionUpgrade) (reexec.ShutdownCallbackFn, error) {
+	if version == "" {
+		return nil, ErrEmptyRollbackVersion
+	}
+
 	// check that the upgrade marker exists and is accessible
 	updateMarkerPath := markerFilePath(paths.DataFrom(topDir))
 	_, err := os.Stat(updateMarkerPath)
@@ -419,18 +425,53 @@ func (u *Upgrader) forceRollbackToPreviousVersion(ctx context.Context, topDir st
 	// 2. there has been at least the first restart with the new agent (i.e. we are not still downloading/extracting/rotating)
 	// 3. upgrade marker exists
 	// these should be revalidated after taking over watcher
-	updateMarker, err := u.persistManualRollback(ctx, topDir)
-	if err != nil {
-		return nil, fmt.Errorf("persisting rollback in update marker: %w", err)
-	}
 
-	previous, current, err := extractAgentInstallsFromMarker(updateMarker)
+	watcherExecutable := ""
+	err = withTakeOverWatcher(ctx, u.log, topDir, u.watcherHelper, func() error {
+		// read the upgrade marker
+		updateMarker, err := LoadMarker(paths.DataFrom(topDir))
+		if err != nil {
+			return fmt.Errorf("loading marker: %w", err)
+		}
+
+		if updateMarker == nil {
+			return ErrNilUpdateMarker
+		}
+
+		if updateMarker.Details == nil || len(updateMarker.Details.Metadata.RollbacksAvailable) == 0 {
+			return ErrNoRollbacksAvailable
+		}
+		var selectedRollback *details.RollbackAvailable
+		for _, rollback := range updateMarker.Details.Metadata.RollbacksAvailable {
+			if rollback.Version == version && now.Before(rollback.ValidUntil) {
+				selectedRollback = &rollback
+				break
+			}
+		}
+		if selectedRollback == nil {
+			return fmt.Errorf("version %q not listed among the available rollbacks: %w", version, ErrNoRollbacksAvailable)
+		}
+
+		// write the desired outcome of the upgrade
+		err = u.persistManualRollback(topDir, updateMarker)
+		if err != nil {
+			return fmt.Errorf("persisting rollback in update marker: %w", err)
+		}
+
+		// extract the agent installs involved in the upgrade and select the most appropriate watcher executable
+		previous, current, err := extractAgentInstallsFromMarker(updateMarker)
+		if err != nil {
+			return fmt.Errorf("extracting current and previous install details: %w", err)
+		}
+		watcherExecutable = u.watcherHelper.SelectWatcherExecutable(topDir, previous, current)
+		return nil
+	})
+
 	if err != nil {
-		return nil, fmt.Errorf("extracting current and previous install details: %w", err)
+		return nil, err
 	}
 
-	// Invoke watcher again
-	watcherExecutable := u.watcherHelper.SelectWatcherExecutable(topDir, previous, current)
+	// Invoke watcher again (now that we released the watcher applocks)
 	_, err = u.watcherHelper.InvokeWatcher(u.log, watcherExecutable)
 	if err != nil {
 		return nil, fmt.Errorf("invoking watcher: %w", err)
@@ -440,6 +481,21 @@ func (u *Upgrader) forceRollbackToPreviousVersion(ctx context.Context, topDir st
 
 }
 
+func withTakeOverWatcher(ctx context.Context, log *logger.Logger, topDir string, watcherHelper WatcherHelper, f func() error) error {
+	watcherApplock, err := watcherHelper.TakeOverWatcher(ctx, log, topDir)
+	if err != nil {
+		return fmt.Errorf("taking over watcher processes: %w", err)
+	}
+	defer func(watcherApplock *filelock.AppLocker) {
+		releaseWatcherAppLockerErr := watcherApplock.Unlock()
+		if releaseWatcherAppLockerErr != nil {
+			log.Warnw("error releasing watcher applock", "error", releaseWatcherAppLockerErr)
+		}
+	}(watcherApplock)
+
+	return f()
+}
+
 func extractAgentInstallsFromMarker(updateMarker *UpdateMarker) (previous agentInstall, current agentInstall, err error) {
 	previousParsedVersion, err := agtversion.ParseVersion(updateMarker.PrevVersion)
 	if err != nil {
@@ -466,35 +522,14 @@ func extractAgentInstallsFromMarker(updateMarker *UpdateMarker) (previous agentI
 	return previous, current, nil
 }
 
-func (u *Upgrader) persistManualRollback(ctx context.Context, topDir string) (*UpdateMarker, error) {
-	watcherApplock, err := u.watcherHelper.TakeOverWatcher(ctx, u.log, topDir)
-	if err != nil {
-		return nil, fmt.Errorf("taking over watcher processes: %w", err)
-	}
-	defer func(watcherApplock *filelock.AppLocker) {
-		releaseWatcherAppLockerErr := watcherApplock.Unlock()
-		if releaseWatcherAppLockerErr != nil {
-			u.log.Warnw("error releasing watcher applock", "error", releaseWatcherAppLockerErr)
-		}
-	}(watcherApplock)
-
-	// read the upgrade marker
-	updateMarker, err := LoadMarker(paths.DataFrom(topDir))
-	if err != nil {
-		return nil, fmt.Errorf("loading marker: %w", err)
-	}
-
-	if updateMarker == nil {
-		return nil, ErrNilUpdateMarker
-	}
-
+func (u *Upgrader) persistManualRollback(topDir string, updateMarker *UpdateMarker) error {
 	updateMarker.DesiredOutcome = OUTCOME_ROLLBACK
-	err = SaveMarker(paths.DataFrom(topDir), updateMarker, true)
+	err := SaveMarker(paths.DataFrom(topDir), updateMarker, true)
 	if err != nil {
-		return updateMarker, fmt.Errorf("saving marker: %w", err)
+		return fmt.Errorf("saving marker: %w", err)
 	}
 
-	return updateMarker, nil
+	return nil
 }
 
 // Ack acks last upgrade action

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

@@ -1038,6 +1038,23 @@ func TestIsSameReleaseVersion(t *testing.T) {
 }
 
 func TestManualRollback(t *testing.T) {
+	const updatemarkerwatching456NoRollbackAvailable = `
+    version: 4.5.6
+    hash: newver
+    versioned_home: data/elastic-agent-4.5.6-newver
+    updated_on: 2025-07-11T10:11:12.131415Z
+    prev_version: 1.2.3
+    prev_hash: oldver
+    prev_versioned_home: data/elastic-agent-1.2.3-oldver
+    acked: false
+    action: null
+    details:
+        target_version: 4.5.6
+        state: UPG_WATCHING
+        metadata:
+            retry_until: null
+    desired_outcome: UPGRADE
+    `
 	const updatemarkerwatching456 = `
     version: 4.5.6
     hash: newver
@@ -1059,6 +1076,7 @@ func TestManualRollback(t *testing.T) {
                   valid_until: 2025-07-18T10:11:12.131415Z
     desired_outcome: UPGRADE
     `
+
 	parsed123Version, err := agtversion.ParseVersion("1.2.3")
 	require.NoError(t, err)
 	parsed456Version, err := agtversion.ParseVersion("4.5.6")
@@ -1078,19 +1096,40 @@ func TestManualRollback(t *testing.T) {
 		versionedHome: "data/elastic-agent-4.5.6-newver",
 	}
 
+	// this is the updated_on timestamp in the example
+	nowBeforeTTL, err := time.Parse(time.RFC3339, `2025-07-11T10:11:12Z`)
+	require.NoError(t, err, "error parsing nowBeforeTTL")
+
+	// the update marker yaml assume 7d TLL for rollbacks, let's make an extra day pass
+	nowAfterTTL := nowBeforeTTL.Add(8 * 24 * time.Hour)
+
 	type setupF func(t *testing.T, topDir string, agent *infomocks.Agent, watcherHelper *MockWatcherHelper)
 	type postRollbackAssertionsF func(t *testing.T, topDir string)
 	type testcase struct {
 		name              string
 		setup             setupF
 		artifactSettings  *artifact.Config
 		upgradeSettings   *configuration.UpgradeConfig
+		now               time.Time
 		version           string
 		wantErr           assert.ErrorAssertionFunc
 		additionalAsserts postRollbackAssertionsF
 	}
 
 	testcases := []testcase{
+		{
+			name: "no rollback version - rollback fails",
+			setup: func(t *testing.T, topDir string, agent *infomocks.Agent, watcherHelper *MockWatcherHelper) {
+				//do not setup anything here, let the rollback fail
+			},
+			artifactSettings: artifact.DefaultConfig(),
+			upgradeSettings:  configuration.DefaultUpgradeConfig(),
+			version:          "",
+			wantErr: func(t assert.TestingT, err error, i ...interface{}) bool {
+				return assert.ErrorIs(t, err, ErrEmptyRollbackVersion)
+			},
+			additionalAsserts: nil,
+		},
 		{
 			name: "no update marker - rollback fails",
 			setup: func(t *testing.T, topDir string, agent *infomocks.Agent, watcherHelper *MockWatcherHelper) {
@@ -1104,6 +1143,85 @@ func TestManualRollback(t *testing.T) {
 			},
 			additionalAsserts: nil,
 		},
+		{
+			name: "update marker ok but rollback available is empty - error",
+			setup: func(t *testing.T, topDir string, agent *infomocks.Agent, watcherHelper *MockWatcherHelper) {
+				err := os.WriteFile(markerFilePath(paths.DataFrom(topDir)), []byte(updatemarkerwatching456NoRollbackAvailable), 0600)
+				require.NoError(t, err, "error setting up update marker")
+				locker := filelock.NewAppLocker(topDir, "watcher.lock")
+				err = locker.TryLock()
+				require.NoError(t, err, "error locking initial watcher AppLocker")
+				watcherHelper.EXPECT().TakeOverWatcher(t.Context(), mock.Anything, topDir).Return(locker, nil)
+			},
+			artifactSettings: artifact.DefaultConfig(),
+			upgradeSettings:  configuration.DefaultUpgradeConfig(),
+			version:          "2.3.4-unknown",
+			wantErr: func(t assert.TestingT, err error, i ...interface{}) bool {
+				return assert.ErrorIs(t, err, ErrNoRollbacksAvailable)
+			},
+			additionalAsserts: func(t *testing.T, topDir string) {
+				// marker should be untouched
+				filePath := markerFilePath(paths.DataFrom(topDir))
+				require.FileExists(t, filePath)
+				markerFileBytes, readMarkerErr := os.ReadFile(filePath)
+				require.NoError(t, readMarkerErr)
+
+				assert.YAMLEq(t, updatemarkerwatching456NoRollbackAvailable, string(markerFileBytes), "update marker should be untouched")
+			},
+		},
+		{
+			name: "update marker ok but version is not available for rollback - error",
+			setup: func(t *testing.T, topDir string, agent *infomocks.Agent, watcherHelper *MockWatcherHelper) {
+				err := os.WriteFile(markerFilePath(paths.DataFrom(topDir)), []byte(updatemarkerwatching456), 0600)
+				require.NoError(t, err, "error setting up update marker")
+				locker := filelock.NewAppLocker(topDir, "watcher.lock")
+				err = locker.TryLock()
+				require.NoError(t, err, "error locking initial watcher AppLocker")
+				watcherHelper.EXPECT().TakeOverWatcher(t.Context(), mock.Anything, topDir).Return(locker, nil)
+			},
+			artifactSettings: artifact.DefaultConfig(),
+			upgradeSettings:  configuration.DefaultUpgradeConfig(),
+			version:          "2.3.4-unknown",
+			wantErr: func(t assert.TestingT, err error, i ...interface{}) bool {
+				return assert.ErrorIs(t, err, ErrNoRollbacksAvailable)
+			},
+			additionalAsserts: func(t *testing.T, topDir string) {
+				// marker should be untouched
+				filePath := markerFilePath(paths.DataFrom(topDir))
+				require.FileExists(t, filePath)
+				markerFileBytes, readMarkerErr := os.ReadFile(filePath)
+				require.NoError(t, readMarkerErr)
+
+				assert.YAMLEq(t, updatemarkerwatching456, string(markerFileBytes), "update marker should be untouched")
+			},
+		},
+		{
+			name: "update marker ok but rollback is expired - error",
+			setup: func(t *testing.T, topDir string, agent *infomocks.Agent, watcherHelper *MockWatcherHelper) {
+				err := os.WriteFile(markerFilePath(paths.DataFrom(topDir)), []byte(updatemarkerwatching456), 0600)
+				require.NoError(t, err, "error setting up update marker")
+				locker := filelock.NewAppLocker(topDir, "watcher.lock")
+				err = locker.TryLock()
+				require.NoError(t, err, "error locking initial watcher AppLocker")
+				watcherHelper.EXPECT().TakeOverWatcher(t.Context(), mock.Anything, topDir).Return(locker, nil)
+			},
+			artifactSettings: artifact.DefaultConfig(),
+			upgradeSettings:  configuration.DefaultUpgradeConfig(),
+			now:              nowAfterTTL,
+			version:          "1.2.3",
+			wantErr: func(t assert.TestingT, err error, i ...interface{}) bool {
+				return assert.ErrorIs(t, err, ErrNoRollbacksAvailable)
+			},
+			additionalAsserts: func(t *testing.T, topDir string) {
+				// marker should be untouched
+				filePath := markerFilePath(paths.DataFrom(topDir))
+				require.FileExists(t, filePath)
+				markerFileBytes, readMarkerErr := os.ReadFile(filePath)
+				require.NoError(t, readMarkerErr)
+
+				assert.YAMLEq(t, updatemarkerwatching456, string(markerFileBytes), "update marker should be untouched")
+			},
+		},
 		{
 			name: "update marker ok - takeover watcher, persist rollback and restart most recent watcher",
 			setup: func(t *testing.T, topDir string, agent *infomocks.Agent, watcherHelper *MockWatcherHelper) {
@@ -1119,6 +1237,7 @@ func TestManualRollback(t *testing.T) {
 			},
 			artifactSettings: artifact.DefaultConfig(),
 			upgradeSettings:  configuration.DefaultUpgradeConfig(),
+			now:              nowBeforeTTL,
 			version:          "1.2.3",
 			wantErr:          assert.NoError,
 			additionalAsserts: func(t *testing.T, topDir string) {
@@ -1148,8 +1267,7 @@ func TestManualRollback(t *testing.T) {
 
 			upgrader, err := NewUpgrader(log, tc.artifactSettings, tc.upgradeSettings, mockAgentInfo, mockWatcherHelper)
 			require.NoError(t, err, "error instantiating upgrader")
-
-			_, err = upgrader.forceRollbackToPreviousVersion(t.Context(), topDir, tc.version, nil, nil)
+			_, err = upgrader.forceRollbackToPreviousVersion(t.Context(), topDir, tc.now, tc.version, nil)
 			tc.wantErr(t, err, "unexpected error returned by forceRollbackToPreviousVersion()")
 			if tc.additionalAsserts != nil {
 				tc.additionalAsserts(t, topDir)