Add ability to enroll with defined ID and replace_token (#6498) (#6807)

Allows an Elastic Agent to enroll with a defined ID and replacement token to allow it to replace an existing Elastic Agent. The original Elastic Agent must have also been enrolled with the same --replace-token or it will not be allow to enroll if the --id collides with an existing Elastic Agent.

(cherry picked from commit 8a878fc)

Co-authored-by: Blake Rouse <blake.rouse@elastic.co>

Author: mergify[bot]

changelog/fragments/1738784301-Add-id-and-replace-token-to-enrollment.yaml

@@ -0,0 +1,36 @@
+# 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: feature
+
+# Change summary; a 80ish characters long description of the change.
+summary: Add --id and --replace-token to enrollment
+
+# 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: |
+  Add support for --id and --replace-token to the install and enroll command. Add support for ELASTIC_AGENT_ID
+  and FLEET_REPLACE_TOKEN to the container support the same behavior as the enroll command. Allows the ability to
+  define a specific ID to use for the Elastic Agent when enrolling into Fleet. The replace-token defines the token
+  that must be used to re-enroll an Elastic Agent with the same ID as a replacement of the previous Elastic Agent.
+
+# Affected component; usually one of "elastic-agent", "fleet-server", "filebeat", "metricbeat", "auditbeat", "all", etc.
+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/elastic-agent/pull/6498
+
+# 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/6361

internal/pkg/agent/application/coordinator/diagnostics_test.go

@@ -59,7 +59,8 @@ func TestDiagnosticLocalConfig(t *testing.T) {
 		Fleet: &configuration.FleetAgentConfig{
 			Enabled:             true,
 			AccessAPIKey:        "test-key",
-			EnrollmentTokenHash: "test-hash",
+			EnrollmentTokenHash: "test-enroll-hash",
+			ReplaceTokenHash:    "test-replace-hash",
 			Client: remote.Config{
 				Protocol: "test-protocol",
 			},
@@ -120,7 +121,8 @@ agent:
 fleet:
   enabled: true
   access_api_key: "test-key"
-  enrollment_token_hash: "test-hash"
+  enrollment_token_hash: "test-enroll-hash"
+  replace_token_hash: "test-replace-hash"
   agent:
   protocol: "test-protocol"
 `

internal/pkg/agent/cmd/container.go

@@ -507,6 +507,12 @@ func buildEnrollArgs(cfg setupConfig, token string, policyID string) ([]string,
 	if token != "" {
 		args = append(args, "--enrollment-token", token)
 	}
+	if cfg.Fleet.ID != "" {
+		args = append(args, "--id", cfg.Fleet.ID)
+	}
+	if cfg.Fleet.ReplaceToken != "" {
+		args = append(args, "--replace-token", cfg.Fleet.ReplaceToken)
+	}
 	if cfg.Fleet.DaemonTimeout != 0 {
 		args = append(args, "--daemon-timeout")
 		args = append(args, cfg.Fleet.DaemonTimeout.String())
@@ -1031,7 +1037,9 @@ func shouldFleetEnroll(setupCfg setupConfig) (bool, error) {
 		return true, nil
 	}
 
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	store, err := newEncryptedDiskStore(ctx, agentCfgFilePath)
 	if err != nil {
 		return false, fmt.Errorf("failed to instantiate encrypted disk store: %w", err)
@@ -1052,6 +1060,13 @@ func shouldFleetEnroll(setupCfg setupConfig) (bool, error) {
 		return false, fmt.Errorf("failed to read from disk store: %w", err)
 	}
 
+	// Check if enrolling with a specifically defined Elastic Agent ID.
+	// If the ID's don't match then it needs to enroll.
+	if setupCfg.Fleet.ID != "" && (storedConfig.Fleet.Info == nil || storedConfig.Fleet.Info.ID != setupCfg.Fleet.ID) {
+		// ID is a mismatch
+		return true, nil
+	}
+
 	storedFleetHosts := storedConfig.Fleet.Client.GetHosts()
 	if len(storedFleetHosts) == 0 || !slices.Contains(storedFleetHosts, setupCfg.Fleet.URL) {
 		// The Fleet URL in the setup does not exist in the stored configuration, so enrollment is required.
@@ -1064,7 +1079,7 @@ func shouldFleetEnroll(setupCfg setupConfig) (bool, error) {
 	if len(storedConfig.Fleet.EnrollmentTokenHash) > 0 && len(setupCfg.Fleet.EnrollmentToken) > 0 {
 		enrollmentHashBytes, err := base64.StdEncoding.DecodeString(storedConfig.Fleet.EnrollmentTokenHash)
 		if err != nil {
-			return false, fmt.Errorf("failed to decode hash: %w", err)
+			return false, fmt.Errorf("failed to decode enrollment token hash: %w", err)
 		}
 
 		err = crypto.ComparePBKDF2HashAndPassword(enrollmentHashBytes, []byte(setupCfg.Fleet.EnrollmentToken))
@@ -1073,7 +1088,26 @@ func shouldFleetEnroll(setupCfg setupConfig) (bool, error) {
 			// The stored enrollment token hash does not match the new token, so enrollment is required.
 			return true, nil
 		case err != nil:
-			return false, fmt.Errorf("failed to compare hash: %w", err)
+			return false, fmt.Errorf("failed to compare enrollment token hash: %w", err)
+		}
+	}
+
+	// Evaluate the stored replace token hash against the setup replace token if both are present.
+	// Note that when "upgrading" from an older agent version the replace token hash will not exist
+	// in the stored configuration.
+	if len(storedConfig.Fleet.ReplaceTokenHash) > 0 && len(setupCfg.Fleet.ReplaceToken) > 0 {
+		replaceHashBytes, err := base64.StdEncoding.DecodeString(storedConfig.Fleet.ReplaceTokenHash)
+		if err != nil {
+			return false, fmt.Errorf("failed to decode replace token hash: %w", err)
+		}
+
+		err = crypto.ComparePBKDF2HashAndPassword(replaceHashBytes, []byte(setupCfg.Fleet.ReplaceToken))
+		switch {
+		case errors.Is(err, crypto.ErrMismatchedHashAndPassword):
+			// The stored enrollment token hash does not match the new token, so enrollment is required.
+			return true, nil
+		case err != nil:
+			return false, fmt.Errorf("failed to compare replace token hash: %w", err)
 		}
 	}
 
@@ -1100,16 +1134,33 @@ func shouldFleetEnroll(setupCfg setupConfig) (bool, error) {
 		return false, fmt.Errorf("failed to validate api token: %w", err)
 	}
 
+	saveConfig := false
+
 	// Update the stored enrollment token hash if there is no previous enrollment token hash
 	// (can happen when "upgrading" from an older version of the agent) and setup enrollment token is present.
 	if len(storedConfig.Fleet.EnrollmentTokenHash) == 0 && len(setupCfg.Fleet.EnrollmentToken) > 0 {
 		enrollmentHashBytes, err := crypto.GeneratePBKDF2FromPassword([]byte(setupCfg.Fleet.EnrollmentToken))
 		if err != nil {
-			return false, errors.New("failed to generate enrollment hash")
+			return false, errors.New("failed to generate enrollment token hash")
 		}
 		enrollmentTokenHash := base64.StdEncoding.EncodeToString(enrollmentHashBytes)
 		storedConfig.Fleet.EnrollmentTokenHash = enrollmentTokenHash
+		saveConfig = true
+	}
+
+	// Update the stored replace token hash if there is no previous replace token hash
+	// (can happen when "upgrading" from an older version of the agent) and setup replace token is present.
+	if len(storedConfig.Fleet.ReplaceTokenHash) == 0 && len(setupCfg.Fleet.ReplaceToken) > 0 {
+		replaceHashBytes, err := crypto.GeneratePBKDF2FromPassword([]byte(setupCfg.Fleet.ReplaceToken))
+		if err != nil {
+			return false, errors.New("failed to generate replace token hash")
+		}
+		replaceTokenHash := base64.StdEncoding.EncodeToString(replaceHashBytes)
+		storedConfig.Fleet.ReplaceTokenHash = replaceTokenHash
+		saveConfig = true
+	}
 
+	if saveConfig {
 		data, err := yaml.Marshal(storedConfig)
 		if err != nil {
 			return false, errors.New("could not marshal config")

internal/pkg/agent/cmd/container_test.go

@@ -259,12 +259,19 @@ func TestKibanaFetchToken(t *testing.T) {
 }
 
 func TestShouldEnroll(t *testing.T) {
-	enrollmentToken := "test-token"
+	// enroll token
+	enrollmentToken := "test-enroll-token"
 	enrollmentTokenHash, err := crypto.GeneratePBKDF2FromPassword([]byte(enrollmentToken))
 	require.NoError(t, err)
 	enrollmentTokenHashBase64 := base64.StdEncoding.EncodeToString(enrollmentTokenHash)
+	enrollmentTokenOther := "test-enroll-token-other"
 
-	enrollmentTokenOther := "test-token-other"
+	// replace token
+	replaceToken := "test-replace-token"
+	replaceTokenHash, err := crypto.GeneratePBKDF2FromPassword([]byte(replaceToken))
+	require.NoError(t, err)
+	replaceTokenHashBase64 := base64.StdEncoding.EncodeToString(replaceTokenHash)
+	replaceTokenOther := "test-replace-token-other"
 
 	fleetNetworkErr := errors.New("fleet network error")
 	for name, tc := range map[string]struct {
@@ -289,6 +296,41 @@ func TestShouldEnroll(t *testing.T) {
 			cfg:                  setupConfig{Fleet: fleetConfig{Enroll: true, Force: true}},
 			expectedShouldEnroll: true,
 		},
+		"should enroll on agent id but no existing id": {
+			statFn: func(path string) (os.FileInfo, error) { return nil, nil },
+			cfg:    setupConfig{Fleet: fleetConfig{Enroll: true, URL: "host1", ID: "diff-agent-id"}},
+			encryptedDiskStoreFn: func(t *testing.T, savedConfig *configuration.Configuration) storage.Storage {
+				m := mockStorage.NewStorage(t)
+				m.On("Load").Return(io.NopCloser(strings.NewReader(`fleet:
+  enabled: true
+  access_api_key: "test-key"
+  enrollment_token_hash: "test-hash"
+  hosts:
+    - host1
+  agent:
+  protocol: "https"`)), nil).Once()
+				return m
+			},
+			expectedShouldEnroll: true,
+		},
+		"should enroll on agent id but diff agent id": {
+			statFn: func(path string) (os.FileInfo, error) { return nil, nil },
+			cfg:    setupConfig{Fleet: fleetConfig{Enroll: true, URL: "host1", ID: "diff-agent-id"}},
+			encryptedDiskStoreFn: func(t *testing.T, savedConfig *configuration.Configuration) storage.Storage {
+				m := mockStorage.NewStorage(t)
+				m.On("Load").Return(io.NopCloser(strings.NewReader(`fleet:
+  enabled: true
+  access_api_key: "test-key"
+  enrollment_token_hash: "test-hash"
+  hosts:
+    - host1
+  agent:
+    id: "agent-id"
+  protocol: "https"`)), nil).Once()
+				return m
+			},
+			expectedShouldEnroll: true,
+		},
 		"should enroll on fleet url change": {
 			statFn: func(path string) (os.FileInfo, error) { return nil, nil },
 			cfg:    setupConfig{Fleet: fleetConfig{Enroll: true, URL: "host1"}},
@@ -321,6 +363,26 @@ func TestShouldEnroll(t *testing.T) {
     - host2
     - host3
   agent:
+  protocol: "https"`)), nil).Once()
+				return m
+			},
+			expectedShouldEnroll: true,
+		},
+		"should enroll on replace token change": {
+			statFn: func(path string) (os.FileInfo, error) { return nil, nil },
+			cfg:    setupConfig{Fleet: fleetConfig{Enroll: true, URL: "host1", EnrollmentToken: enrollmentToken, ReplaceToken: replaceTokenOther}},
+			encryptedDiskStoreFn: func(t *testing.T, savedConfig *configuration.Configuration) storage.Storage {
+				m := mockStorage.NewStorage(t)
+				m.On("Load").Return(io.NopCloser(strings.NewReader(`fleet:
+  enabled: true
+  access_api_key: "test-key"
+  enrollment_token_hash: "`+enrollmentTokenHashBase64+`"
+  replace_token_hash: "`+replaceTokenHashBase64+`"
+  hosts:
+    - host1
+    - host2
+    - host3
+  agent:
   protocol: "https"`)), nil).Once()
 				return m
 			},
@@ -373,6 +435,44 @@ func TestShouldEnroll(t *testing.T) {
     - host2
     - host3
   agent:
+  protocol: "https"`)), nil).Once()
+				return m
+			},
+			fleetClientFn: func(t *testing.T) client.Sender {
+				tries := 0
+				m := mockFleetClient.NewSender(t)
+				call := m.On("Send", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
+				call.Run(func(args mock.Arguments) {
+					if tries <= 1 {
+						call.Return(nil, fleetNetworkErr)
+					} else {
+						call.Return(&http.Response{
+							StatusCode: http.StatusOK,
+							Body:       io.NopCloser(strings.NewReader(`{"action": "acks", "items":[]}`)),
+						}, nil)
+					}
+					tries++
+				}).Times(3)
+				return m
+			},
+			expectedShouldEnroll: false,
+		},
+		"should not enroll on no changes with agent ID and replace token": {
+			statFn: func(path string) (os.FileInfo, error) { return nil, nil },
+			cfg:    setupConfig{Fleet: fleetConfig{Enroll: true, URL: "host1", ID: "custom-id", EnrollmentToken: enrollmentToken, ReplaceToken: replaceToken}},
+			encryptedDiskStoreFn: func(t *testing.T, savedConfig *configuration.Configuration) storage.Storage {
+				m := mockStorage.NewStorage(t)
+				m.On("Load").Return(io.NopCloser(strings.NewReader(`fleet:
+  enabled: true
+  access_api_key: "test-key"
+  enrollment_token_hash: "`+enrollmentTokenHashBase64+`"
+  replace_token_hash: "`+replaceTokenHashBase64+`"
+  hosts:
+    - host1
+    - host2
+    - host3
+  agent:
+    id: "custom-id"
   protocol: "https"`)), nil).Once()
 				return m
 			},
@@ -433,6 +533,33 @@ func TestShouldEnroll(t *testing.T) {
     - host2
     - host3
   agent:
+  protocol: "https"`)), nil).Once()
+				return m
+			},
+			fleetClientFn: func(t *testing.T) client.Sender {
+				m := mockFleetClient.NewSender(t)
+				m.On("Send", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+					Return(&http.Response{
+						StatusCode: http.StatusOK,
+						Body:       io.NopCloser(strings.NewReader(`{"action": "acks", "items":[]}`)),
+					}, nil).Once()
+				return m
+			},
+			expectedShouldEnroll: false,
+		},
+		"should not update the replace token hash if it does not exist in setup configuration": {
+			statFn: func(path string) (os.FileInfo, error) { return nil, nil },
+			cfg:    setupConfig{Fleet: fleetConfig{Enroll: true, URL: "host1", EnrollmentToken: "", ReplaceToken: ""}},
+			encryptedDiskStoreFn: func(t *testing.T, savedConfig *configuration.Configuration) storage.Storage {
+				m := mockStorage.NewStorage(t)
+				m.On("Load").Return(io.NopCloser(strings.NewReader(`fleet:
+  enabled: true
+  access_api_key: "test-key"
+  hosts:
+    - host1
+    - host2
+    - host3
+  agent:
   protocol: "https"`)), nil).Once()
 				return m
 			},
@@ -486,6 +613,48 @@ func TestShouldEnroll(t *testing.T) {
 			},
 			expectedShouldEnroll: false,
 		},
+		"should not enroll on no changes and update the stored enrollment and replace token hash": {
+			statFn: func(path string) (os.FileInfo, error) { return nil, nil },
+			cfg:    setupConfig{Fleet: fleetConfig{Enroll: true, URL: "host1", EnrollmentToken: enrollmentToken, ReplaceToken: replaceToken}},
+			encryptedDiskStoreFn: func(t *testing.T, savedConfig *configuration.Configuration) storage.Storage {
+				m := mockStorage.NewStorage(t)
+				m.On("Load").Return(io.NopCloser(strings.NewReader(`fleet:
+  enabled: true
+  access_api_key: "test-key"
+  hosts:
+    - host1
+    - host2
+    - host3
+  agent:
+  protocol: "https"`)), nil).Once()
+				m.On("Save", mock.Anything).Run(func(args mock.Arguments) {
+					reader := args.Get(0).(io.Reader)
+					data, _ := io.ReadAll(reader)
+					_ = yaml.Unmarshal(data, savedConfig)
+				}).Return(nil).Times(0)
+				return m
+			},
+			fleetClientFn: func(t *testing.T) client.Sender {
+				m := mockFleetClient.NewSender(t)
+				m.On("Send", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+					Return(&http.Response{
+						StatusCode: http.StatusOK,
+						Body:       io.NopCloser(strings.NewReader(`{"action": "acks", "items":[]}`)),
+					}, nil).Once()
+				return m
+			},
+			expectedSavedConfig: func(t *testing.T, savedConfig *configuration.Configuration) {
+				require.NotNil(t, savedConfig)
+				require.NotNil(t, savedConfig.Fleet)
+				enrollmentTokenHash, err := base64.StdEncoding.DecodeString(savedConfig.Fleet.EnrollmentTokenHash)
+				require.NoError(t, err)
+				require.NoError(t, crypto.ComparePBKDF2HashAndPassword(enrollmentTokenHash, []byte(enrollmentToken)))
+				replaceTokenHash, err := base64.StdEncoding.DecodeString(savedConfig.Fleet.ReplaceTokenHash)
+				require.NoError(t, err)
+				require.NoError(t, crypto.ComparePBKDF2HashAndPassword(replaceTokenHash, []byte(replaceToken)))
+			},
+			expectedShouldEnroll: false,
+		},
 	} {
 		t.Run(name, func(t *testing.T) {
 			savedConfig := &configuration.Configuration{}