diff --git a/integration/autoupdate/tools/updater/main.go b/integration/autoupdate/tools/updater/main.go index 8a12dbbd1c9f7..1aa591609362c 100644 --- a/integration/autoupdate/tools/updater/main.go +++ b/integration/autoupdate/tools/updater/main.go @@ -65,7 +65,7 @@ func main() { } // Re-execute client tools with the correct version of client tools. - code, err := updater.Exec() + code, err := updater.Exec(os.Args[1:]) if err != nil { log.Fatalf("Failed to re-exec client tool: %v\n", err) } else { diff --git a/lib/autoupdate/tools/helper.go b/lib/autoupdate/tools/helper.go index 79069c3d3c90b..4c7a3c4dfe088 100644 --- a/lib/autoupdate/tools/helper.go +++ b/lib/autoupdate/tools/helper.go @@ -36,7 +36,7 @@ import ( // with the updated version. // If $TELEPORT_HOME/bin contains downloaded client tools, it always re-executes // using the version from the home directory. -func CheckAndUpdateLocal(ctx context.Context, currentVersion string) error { +func CheckAndUpdateLocal(ctx context.Context, currentVersion string, reExecArgs ...string) error { toolsDir, err := Dir() if err != nil { slog.WarnContext(ctx, "Client tools update is disabled", "error", err) @@ -51,7 +51,10 @@ func CheckAndUpdateLocal(ctx context.Context, currentVersion string) error { return trace.Wrap(err) } if reExec { - return trace.Wrap(updateAndReExec(ctx, updater, toolsVersion)) + if len(reExecArgs) == 0 { + reExecArgs = os.Args[1:] + } + return trace.Wrap(updateAndReExec(ctx, updater, toolsVersion, reExecArgs)) } return nil @@ -64,7 +67,7 @@ func CheckAndUpdateLocal(ctx context.Context, currentVersion string) error { // with the updated version. // If $TELEPORT_HOME/bin contains downloaded client tools, it always re-executes // using the version from the home directory. -func CheckAndUpdateRemote(ctx context.Context, currentVersion string, proxy string, insecure bool) error { +func CheckAndUpdateRemote(ctx context.Context, currentVersion string, proxy string, insecure bool, reExecArgs ...string) error { toolsDir, err := Dir() if err != nil { slog.WarnContext(ctx, "Client tools update is disabled", "error", err) @@ -81,13 +84,16 @@ func CheckAndUpdateRemote(ctx context.Context, currentVersion string, proxy stri return trace.Wrap(err) } if reExec { - return trace.Wrap(updateAndReExec(ctx, updater, toolsVersion)) + if len(reExecArgs) == 0 { + reExecArgs = os.Args[1:] + } + return trace.Wrap(updateAndReExec(ctx, updater, toolsVersion, reExecArgs)) } return nil } -func updateAndReExec(ctx context.Context, updater *Updater, toolsVersion string) error { +func updateAndReExec(ctx context.Context, updater *Updater, toolsVersion string, args []string) error { ctxUpdate, cancel := stacksignal.GetSignalHandler().NotifyContext(ctx) defer cancel() // Download the version of client tools required by the cluster. This @@ -99,7 +105,7 @@ func updateAndReExec(ctx context.Context, updater *Updater, toolsVersion string) } // Re-execute client tools with the correct version of client tools. - code, err := updater.Exec() + code, err := updater.Exec(args) if err != nil && !errors.Is(err, os.ErrNotExist) { slog.DebugContext(ctx, "Failed to re-exec client tool", "error", err) os.Exit(code) diff --git a/lib/autoupdate/tools/updater.go b/lib/autoupdate/tools/updater.go index 8bad07c395391..20860cf852b2d 100644 --- a/lib/autoupdate/tools/updater.go +++ b/lib/autoupdate/tools/updater.go @@ -61,7 +61,7 @@ const ( ) var ( - // // pattern is template for response on version command for client tools {tsh, tctl}. + // pattern is template for response on version command for client tools {tsh, tctl}. pattern = regexp.MustCompile(`(?m)Teleport v(.*) git`) ) @@ -327,7 +327,7 @@ func (u *Updater) update(ctx context.Context, pkg packageURL, pkgName string) er } // Exec re-executes tool command with same arguments and environ variables. -func (u *Updater) Exec() (int, error) { +func (u *Updater) Exec(args []string) (int, error) { path, err := toolName(u.toolsDir) if err != nil { return 0, trace.Wrap(err) @@ -336,7 +336,7 @@ func (u *Updater) Exec() (int, error) { env := append(os.Environ(), teleportToolsVersionEnv+"=off") if runtime.GOOS == constants.WindowsOS { - cmd := exec.Command(path, os.Args[1:]...) + cmd := exec.Command(path, args...) cmd.Env = env cmd.Stdin = os.Stdin cmd.Stdout = os.Stdout @@ -348,7 +348,7 @@ func (u *Updater) Exec() (int, error) { return cmd.ProcessState.ExitCode(), nil } - if err := syscall.Exec(path, append([]string{path}, os.Args[1:]...), env); err != nil { + if err := syscall.Exec(path, append([]string{path}, args...), env); err != nil { return 0, trace.Wrap(err) } @@ -389,11 +389,6 @@ func (u *Updater) downloadHash(ctx context.Context, url string) ([]byte, error) // downloadArchive downloads the archive package by `url` and writes content to the writer interface, // return calculated sha256 hash sum of the content. func (u *Updater) downloadArchive(ctx context.Context, url string, f io.Writer) ([]byte, error) { - // Display a progress bar before initiating the update request to inform the user that - // an update is in progress, allowing them the option to cancel before actual response - // which might be delayed with slow internet connection or complete isolation to CDN. - pw, finish := newProgressWriter(10) - defer finish() req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) if err != nil { return nil, trace.Wrap(err) @@ -416,6 +411,9 @@ func (u *Updater) downloadArchive(ctx context.Context, url string, f io.Writer) } } + pw, finish := newProgressWriter(10) + defer finish() + h := sha256.New() // It is a little inefficient to download the file to disk and then re-load // it into memory to unarchive later, but this is safer as it allows client diff --git a/tool/tsh/common/tsh.go b/tool/tsh/common/tsh.go index 8c51c89641eda..78837615433d4 100644 --- a/tool/tsh/common/tsh.go +++ b/tool/tsh/common/tsh.go @@ -708,7 +708,7 @@ func initLogger(cf *CLIConf) { // // DO NOT RUN TESTS that call Run() in parallel (unless you taken precautions). func Run(ctx context.Context, args []string, opts ...CliOption) error { - if err := tools.CheckAndUpdateLocal(ctx, teleport.Version); err != nil { + if err := tools.CheckAndUpdateLocal(ctx, teleport.Version, args...); err != nil { return trace.Wrap(err) } @@ -1483,7 +1483,7 @@ func Run(ctx context.Context, args []string, opts ...CliOption) error { case sessionsList.FullCommand(): err = onListSessions(&cf) case login.FullCommand(): - err = onLogin(&cf) + err = onLogin(&cf, args...) case logout.FullCommand(): err = onLogout(&cf) case show.FullCommand(): @@ -1834,7 +1834,7 @@ func serializeVersion(format string, proxyVersion string, proxyPublicAddress str } // onLogin logs in with remote proxy and gets signed certificates -func onLogin(cf *CLIConf) error { +func onLogin(cf *CLIConf, reExecArgs ...string) error { autoRequest := true // special case: --request-roles=no disables auto-request behavior. if cf.DesiredRoles == "no" { @@ -1875,7 +1875,7 @@ func onLogin(cf *CLIConf) error { // The user is not logged in and has typed in `tsh --proxy=... login`, if // the running binary needs to be updated, update and re-exec. if profile == nil { - if err := tools.CheckAndUpdateRemote(cf.Context, teleport.Version, tc.WebProxyAddr, tc.InsecureSkipVerify); err != nil { + if err := tools.CheckAndUpdateRemote(cf.Context, teleport.Version, tc.WebProxyAddr, tc.InsecureSkipVerify, reExecArgs...); err != nil { return trace.Wrap(err) } } @@ -1893,7 +1893,7 @@ func onLogin(cf *CLIConf) error { // The user has typed `tsh login`, if the running binary needs to // be updated, update and re-exec. - if err := tools.CheckAndUpdateRemote(cf.Context, teleport.Version, tc.WebProxyAddr, tc.InsecureSkipVerify); err != nil { + if err := tools.CheckAndUpdateRemote(cf.Context, teleport.Version, tc.WebProxyAddr, tc.InsecureSkipVerify, reExecArgs...); err != nil { return trace.Wrap(err) } @@ -1913,7 +1913,7 @@ func onLogin(cf *CLIConf) error { // The user has typed `tsh login`, if the running binary needs to // be updated, update and re-exec. - if err := tools.CheckAndUpdateRemote(cf.Context, teleport.Version, tc.WebProxyAddr, tc.InsecureSkipVerify); err != nil { + if err := tools.CheckAndUpdateRemote(cf.Context, teleport.Version, tc.WebProxyAddr, tc.InsecureSkipVerify, reExecArgs...); err != nil { return trace.Wrap(err) } @@ -1989,7 +1989,7 @@ func onLogin(cf *CLIConf) error { default: // The user is logged in and has typed in `tsh --proxy=... login`, if // the running binary needs to be updated, update and re-exec. - if err := tools.CheckAndUpdateRemote(cf.Context, teleport.Version, tc.WebProxyAddr, tc.InsecureSkipVerify); err != nil { + if err := tools.CheckAndUpdateRemote(cf.Context, teleport.Version, tc.WebProxyAddr, tc.InsecureSkipVerify, reExecArgs...); err != nil { return trace.Wrap(err) } }