From d147dd536ea80c0da774eeb0e45e33d4f8e6d45f Mon Sep 17 00:00:00 2001
From: Blake Rouse <blake.rouse@elastic.co>
Date: Wed, 12 Jan 2022 12:03:37 -0500
Subject: [PATCH] [Elastic Agent] Fix issue with ensureServiceToken. (#29800)

* Fix issue with ensureServiceToken.

* Move ensureServiceToken up to line 273.

(cherry picked from commit cb5f2eb3eb43383121dafd82b03935a6eb2ab62f)
---
 .../elastic-agent/pkg/agent/cmd/container.go  | 20 ++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/x-pack/elastic-agent/pkg/agent/cmd/container.go b/x-pack/elastic-agent/pkg/agent/cmd/container.go
index 33f46728b65..25ff4c15ddc 100644
--- a/x-pack/elastic-agent/pkg/agent/cmd/container.go
+++ b/x-pack/elastic-agent/pkg/agent/cmd/container.go
@@ -199,11 +199,6 @@ func containerCmd(streams *cli.IOStreams, cmd *cobra.Command) error {
 		}
 	}
 
-	err = ensureServiceToken(streams, &cfg)
-	if err != nil {
-		return err
-	}
-
 	// start apm-server legacy process when in cloud mode
 	var wg sync.WaitGroup
 	var apmProc *process.Info
@@ -274,6 +269,12 @@ func runContainerCmd(streams *cli.IOStreams, cmd *cobra.Command, cfg setupConfig
 		return run(streams, logToStderr)
 	}
 
+	if cfg.Kibana.Fleet.Setup || cfg.Fleet.Enroll {
+		err = ensureServiceToken(streams, &cfg)
+		if err != nil {
+			return err
+		}
+	}
 	if cfg.Kibana.Fleet.Setup {
 		client, err = kibanaClient(cfg.Kibana, cfg.Kibana.Headers)
 		if err != nil {
@@ -309,7 +310,10 @@ func runContainerCmd(streams *cli.IOStreams, cmd *cobra.Command, cfg setupConfig
 		if policy != nil {
 			policyID = policy.ID
 		}
-		logInfo(streams, "Policy selected for enrollment: ", policyID)
+		if policyID != "" {
+			logInfo(streams, "Policy selected for enrollment: ", policyID)
+		}
+
 		cmdArgs, err := buildEnrollArgs(cfg, token, policyID)
 		if err != nil {
 			return err
@@ -349,11 +353,13 @@ func ensureServiceToken(streams *cli.IOStreams, cfg *setupConfig) error {
 	}
 
 	logInfo(streams, "Requesting service_token from Kibana.")
+
+	// Client is not passed in to this function because this function will use username/password and then
+	// all the following clients will use the created service token.
 	client, err := kibanaClient(cfg.Kibana, cfg.Kibana.Headers)
 	if err != nil {
 		return err
 	}
-
 	code, r, err := client.Connection.Request("POST", "/api/fleet/service-tokens", nil, nil, nil)
 	if err != nil {
 		return fmt.Errorf("request to get security token from Kibana failed: %w", err)