From b8dc66888d047355edd1938c71a2e8a0d5219616 Mon Sep 17 00:00:00 2001
From: Vadym Popov <vadym.popov@goteleport.com>
Date: Thu, 16 Jan 2025 22:01:36 -0800
Subject: [PATCH] Restore priority of env check over remote check

In case of double re-execution case we should stop second one to prevent re-execution loop
---
 .../autoupdate/tools/updater_tsh_test.go      |  6 ++--
 lib/autoupdate/tools/updater.go               | 32 +++++++++++++++++--
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/integration/autoupdate/tools/updater_tsh_test.go b/integration/autoupdate/tools/updater_tsh_test.go
index 661b6bec18df7..a318db603f241 100644
--- a/integration/autoupdate/tools/updater_tsh_test.go
+++ b/integration/autoupdate/tools/updater_tsh_test.go
@@ -96,7 +96,7 @@ func TestAliasLoginWithUpdater(t *testing.T) {
 	require.NoError(t, err)
 	configPath := filepath.Join(toolsDir, client.TSHConfigPath)
 	require.NoError(t, os.MkdirAll(filepath.Dir(configPath), 0700))
-	executable := filepath.Join(toolsDir, "tsh")
+	executable := filepath.Join(os.TempDir(), "tsh")
 	out, err := yaml.Marshal(client.TSHConfig{
 		Aliases: map[string]string{
 			"loginalice": fmt.Sprintf(
@@ -109,7 +109,7 @@ func TestAliasLoginWithUpdater(t *testing.T) {
 	require.NoError(t, os.WriteFile(configPath, out, 0600))
 
 	// Fetch compiled test binary and install to tools dir [v1.2.3].
-	err = tools.NewUpdater(toolsDir, testVersions[0], tools.WithBaseURL(baseURL)).Update(ctx, testVersions[0])
+	err = tools.NewUpdater(os.TempDir(), testVersions[0], tools.WithBaseURL(baseURL)).Update(ctx, testVersions[0])
 	require.NoError(t, err)
 
 	// Execute alias command which must be transformed to the login command.
@@ -122,7 +122,7 @@ func TestAliasLoginWithUpdater(t *testing.T) {
 	require.NoError(t, cmd.Run())
 
 	// Verify tctl status after login.
-	cmd = exec.CommandContext(ctx, filepath.Join(toolsDir, "tctl"), "status", "--insecure")
+	cmd = exec.CommandContext(ctx, filepath.Join(os.TempDir(), "tctl"), "status", "--insecure")
 	cmd.Env = os.Environ()
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
diff --git a/lib/autoupdate/tools/updater.go b/lib/autoupdate/tools/updater.go
index 9ca924b15d536..4fc70ff064049 100644
--- a/lib/autoupdate/tools/updater.go
+++ b/lib/autoupdate/tools/updater.go
@@ -155,11 +155,31 @@ func (u *Updater) CheckLocal() (version string, reExec bool, err error) {
 	return toolsVersion, true, nil
 }
 
-// CheckRemote checks against the Proxy Service to determine if client tools need updating by requesting
+// CheckRemote first checks the version set by the environment variable. If not set or disabled,
+// it checks against the Proxy Service to determine if client tools need updating by requesting
 // the `webapi/find` handler, which stores information about the required client tools version to
 // operate with this cluster. It returns the semantic version that needs updating and whether
 // re-execution is necessary, by re-execution flag we understand that update and re-execute is required.
 func (u *Updater) CheckRemote(ctx context.Context, proxyAddr string, insecure bool) (version string, reExec bool, err error) {
+	// Check if the user has requested a specific version of client tools.
+	requestedVersion := os.Getenv(teleportToolsVersionEnv)
+	switch requestedVersion {
+	// The user has turned off any form of automatic updates.
+	case "off":
+		return "", false, nil
+	// Requested version already the same as client version.
+	case u.localVersion:
+		return u.localVersion, false, nil
+	// No requested version, we continue.
+	case "":
+	// Requested version that is not the local one.
+	default:
+		if _, err := semver.NewVersion(requestedVersion); err != nil {
+			return "", false, trace.Wrap(err, "checking that request version is semantic")
+		}
+		return requestedVersion, true, nil
+	}
+
 	certPool, err := x509.SystemCertPool()
 	if err != nil {
 		return "", false, trace.Wrap(err)
@@ -317,7 +337,15 @@ func (u *Updater) Exec(args []string) (int, error) {
 	if err := os.Unsetenv(teleportToolsVersionEnv); err != nil {
 		return 0, trace.Wrap(err)
 	}
-	env := append(os.Environ(), teleportToolsVersionEnv+"=off")
+
+	env := os.Environ()
+	executablePath, err := os.Executable()
+	if err != nil {
+		return 0, trace.Wrap(err)
+	}
+	if path == executablePath {
+		env = append(env, teleportToolsVersionEnv+"=off")
+	}
 
 	if runtime.GOOS == constants.WindowsOS {
 		cmd := exec.Command(path, args...)