Skip to content

Commit

Permalink
Fix auto-update re-exec arguments modified by aliases
Browse files Browse the repository at this point in the history
  • Loading branch information
vapopov committed Dec 13, 2024
1 parent 75c2cef commit ba62d79
Show file tree
Hide file tree
Showing 4 changed files with 27 additions and 23 deletions.
2 changes: 1 addition & 1 deletion integration/autoupdate/tools/updater/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
18 changes: 12 additions & 6 deletions lib/autoupdate/tools/helper.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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)
Expand Down
16 changes: 7 additions & 9 deletions lib/autoupdate/tools/updater.go
Original file line number Diff line number Diff line change
Expand Up @@ -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`)
)

Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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)
}

Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down
14 changes: 7 additions & 7 deletions tool/tsh/common/tsh.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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" {
Expand Down Expand Up @@ -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)
}
}
Expand All @@ -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)
}

Expand All @@ -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)
}

Expand Down Expand Up @@ -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)
}
}
Expand Down

0 comments on commit ba62d79

Please sign in to comment.