Add check of disk space

Add request timeout and default client
Feature flag check

lib/client/api.go

@@ -691,16 +691,14 @@ func RetryWithRelogin(ctx context.Context, tc *TeleportClient, fn func() error,
 	// If needed, download the new version of {tsh, tctl} and re-exec. Make
 	// sure to exit this process with the same exit code as the child process.
 	//
-	// TODO(russjones): Consolidate this with updateAndRun.
-	fmt.Printf("--> Will check remote and if needed, will update binary and re-exec.\n")
-
 	toolsVersion, reexec, err := update.CheckRemote(ctx, tc.WebProxyAddr)
 	if err != nil {
 		return trace.Wrap(err)
 	}
 	if reexec {
 		// Download the version of client tools required by the cluster.
-		if err := update.Download(toolsVersion); err != nil {
+		err := update.Download(toolsVersion)
+		if err != nil {
 			return trace.Wrap(err)
 		}
 

tool/common/update/client.go

@@ -0,0 +1,61 @@
+/*
+ * Teleport
+ * Copyright (C) 2024  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package update
+
+import (
+	"crypto/tls"
+	"crypto/x509"
+	"net/http"
+	"net/url"
+	"time"
+
+	"golang.org/x/net/http/httpproxy"
+
+	apidefaults "github.com/gravitational/teleport/api/defaults"
+	tracehttp "github.com/gravitational/teleport/api/observability/tracing/http"
+	apiutils "github.com/gravitational/teleport/api/utils"
+)
+
+type downloadConfig struct {
+	// Insecure turns off TLS certificate verification when enabled.
+	Insecure bool
+	// Pool defines the set of root CAs to use when verifying server
+	// certificates.
+	Pool *x509.CertPool
+	// Timeout is a timeout for requests.
+	Timeout time.Duration
+}
+
+func newClient(cfg *downloadConfig) *http.Client {
+	rt := apiutils.NewHTTPRoundTripper(&http.Transport{
+		TLSClientConfig: &tls.Config{
+			InsecureSkipVerify: cfg.Insecure,
+			RootCAs:            cfg.Pool,
+		},
+		Proxy: func(req *http.Request) (*url.URL, error) {
+			return httpproxy.FromEnvironment().ProxyFunc()(req.URL)
+		},
+		IdleConnTimeout: apidefaults.DefaultIOTimeout,
+	}, nil)
+
+	return &http.Client{
+		Transport: tracehttp.NewTransport(rt),
+		Timeout:   cfg.Timeout,
+	}
+}

tool/common/update/feature_ent.go

@@ -0,0 +1,26 @@
+//go:build webassets_ent
+// +build webassets_ent
+
+/*
+ * Teleport
+ * Copyright (C) 2024  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package update
+
+func init() {
+	featureFlag |= FlagEnt
+}

tool/common/update/feature_fips.go

@@ -0,0 +1,26 @@
+//go:build fips
+// +build fips
+
+/*
+ * Teleport
+ * Copyright (C) 2024  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package update
+
+func init() {
+	featureFlag |= FlagFips
+}

tool/common/update/integration/main.go

@@ -1,3 +1,21 @@
+/*
+ * Teleport
+ * Copyright (C) 2024  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
 package main
 
 import (

tool/common/update/update.go

@@ -23,6 +23,8 @@ import (
 	"bytes"
 	"context"
 	"crypto/sha256"
+	"crypto/x509"
+	"encoding/hex"
 	"fmt"
 	"io"
 	"net/http"
@@ -46,11 +48,18 @@ import (
 
 const (
 	teleportToolsVersion = "TELEPORT_TOOLS_VERSION"
+	checksumHexLen       = 64
+	reservedFreeDisk     = 10 * 1024 * 1024
+
+	FlagEnt  = 1 << 0
+	FlagFips = 1 << 1
 )
 
 var (
-	pattern = regexp.MustCompile(`(?m)Teleport v(.*) git`)
-	baseUrl = "https://cdn.teleport.dev"
+	pattern       = regexp.MustCompile(`(?m)Teleport v(.*) git`)
+	baseUrl       = "https://cdn.teleport.dev"
+	defaultClient = newClient(&downloadConfig{})
+	featureFlag   int
 )
 
 // CheckLocal is run at client tool startup and will only perform local checks.
@@ -86,7 +95,16 @@ func CheckRemote(ctx context.Context, proxyAddr string) (string, bool, error) {
 		return "", false, nil
 	}
 
-	resp, err := webclient.Find(&webclient.Config{Context: ctx, ProxyAddr: proxyAddr})
+	certPool, err := x509.SystemCertPool()
+	if err != nil {
+		return "", false, trace.Wrap(err)
+	}
+	resp, err := webclient.Find(&webclient.Config{
+		Context:   ctx,
+		ProxyAddr: proxyAddr,
+		Pool:      certPool,
+		Timeout:   30 * time.Second,
+	})
 	if err != nil {
 		return "", false, trace.Wrap(err)
 	}
@@ -107,6 +125,7 @@ func CheckRemote(ctx context.Context, proxyAddr string) (string, bool, error) {
 	return "", false, nil
 }
 
+// Download downloads requested version package, unarchive and replace existing one.
 func Download(toolsVersion string) error {
 	// If the version of the running binary or the version downloaded to
 	// $TELEPORT_HOME/bin is the same as the requested version of client tools,
@@ -138,7 +157,6 @@ func Download(toolsVersion string) error {
 	return nil
 }
 
-// TODO(russjones): Add edition check here as well.
 func update(toolsVersion string) error {
 	dir, err := toolsDir()
 	if err != nil {
@@ -152,7 +170,7 @@ func update(toolsVersion string) error {
 	defer unlock()
 
 	// Get platform specific download URLs.
-	archiveURL, hashURL, err := urls(toolsVersion, "")
+	archiveURL, hashURL, err := urls(toolsVersion)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -180,7 +198,7 @@ func update(toolsVersion string) error {
 
 // urls returns the URL for the Teleport archive to download. The format is:
 // https://cdn.teleport.dev/teleport-{, ent-}v15.3.0-{linux, darwin, windows}-{amd64,arm64,arm,386}-{fips-}bin.tar.gz
-func urls(toolsVersion string, toolsEdition string) (string, string, error) {
+func urls(toolsVersion string) (string, string, error) {
 	var archive string
 
 	switch runtime.GOOS {
@@ -190,11 +208,11 @@ func urls(toolsVersion string, toolsEdition string) (string, string, error) {
 		archive = baseUrl + "/teleport-v" + toolsVersion + "-windows-amd64-bin.zip"
 	case "linux":
 		edition := ""
-		if toolsEdition == "ent" || toolsEdition == "fips" {
+		if featureFlag&FlagEnt != 0 || featureFlag&FlagFips != 0 {
 			edition = "ent-"
 		}
 		fips := ""
-		if toolsEdition == "fips" {
+		if featureFlag&FlagFips != 0 {
 			fips = "fips-"
 		}
 
@@ -217,7 +235,14 @@ func urls(toolsVersion string, toolsEdition string) (string, string, error) {
 }
 
 func downloadHash(url string) (string, error) {
-	resp, err := http.Get(url)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return "", trace.Wrap(err)
+	}
+	resp, err := defaultClient.Do(req)
 	if err != nil {
 		return "", trace.Wrap(err)
 	}
@@ -227,17 +252,20 @@ func downloadHash(url string) (string, error) {
 		return "", trace.BadParameter("request failed with: %v", resp.StatusCode)
 	}
 
-	body, err := io.ReadAll(resp.Body)
+	var buf bytes.Buffer
+	_, err = io.CopyN(&buf, resp.Body, checksumHexLen)
 	if err != nil {
 		return "", trace.Wrap(err)
 	}
-
-	// Hash is the first 64 bytes of the response.
-	return string(body)[0:64], nil
+	raw := buf.String()
+	if _, err = hex.DecodeString(raw); err != nil {
+		return "", trace.Wrap(err)
+	}
+	return raw, nil
 }
 
 func downloadArchive(url string, hash string) (string, error) {
-	resp, err := http.Get(url)
+	resp, err := defaultClient.Get(url)
 	if err != nil {
 		return "", trace.Wrap(err)
 	}
@@ -251,6 +279,12 @@ func downloadArchive(url string, hash string) (string, error) {
 		return "", trace.Wrap(err)
 	}
 
+	if resp.ContentLength != -1 {
+		if err := checkFreeSpace(dir, uint64(resp.ContentLength)); err != nil {
+			return "", trace.Wrap(err)
+		}
+	}
+
 	// Caller of this function will remove this file after the atomic swap has
 	// occurred.
 	f, err := os.CreateTemp(dir, "tmp-")
@@ -372,3 +406,16 @@ func toolName() (string, error) {
 
 	return filepath.Join(base, toolName), nil
 }
+
+// checkFreeSpace verifies that we have enough requested space at specific directory.
+func checkFreeSpace(path string, requested uint64) error {
+	free, err := freeDiskWithReserve(path)
+	if err != nil {
+		return trace.Errorf("failed to calculate free disk in %q: %v", path, err)
+	}
+	// Bail if there's not enough free disk space at the target.
+	if requested > free {
+		return trace.Errorf("%q needs %d additional bytes of disk space", path, requested-free)
+	}
+	return nil
+}

tool/common/update/update_darwin.go

@@ -57,7 +57,7 @@ func replace(path string, hash string) error {
 		// The first time a signed and notarized binary macOS application is run,
 		// execution is paused while it gets sent to Apple to verify. Once Apple
 		// approves the binary, the "com.apple.macl" extended attribute is added
-		// and the process is allow to execute. This process is not concurrent, any
+		// and the process is allowed to execute. This process is not concurrent, any
 		// other operations (like moving the application) on the application during
 		// this time will lead to the application being sent SIGKILL.
 		//

tool/common/update/update_linux.go

@@ -66,6 +66,11 @@ func replace(path string, _ string) error {
 			continue
 		}
 
+		// Verify that we have enough space for uncompressed file.
+		if err := checkFreeSpace(tempDir, uint64(header.Size)); err != nil {
+			return trace.Wrap(err)
+		}
+
 		dest := filepath.Join(dir, strings.TrimPrefix(header.Name, "teleport/"))
 		t, err := renameio.TempFile(tempDir, dest)
 		if err != nil {

tool/common/update/update_test.go

@@ -120,6 +120,13 @@ func TestUpdateInterruptSignal(t *testing.T) {
 		os.Environ(),
 		fmt.Sprintf("%s=%s", teleportToolsVersion, testVersions[1]),
 	)
+	err = cmd.Start()
+	require.NoError(t, err, "failed to start updater")
+
+	errChan := make(chan error)
+	go func() {
+		errChan <- cmd.Wait()
+	}()
 
 	// By setting the limit request next test http serving file going blocked until unlock is sent.
 	lock := make(chan struct{})
@@ -128,12 +135,9 @@ func TestUpdateInterruptSignal(t *testing.T) {
 		lock:  lock,
 	})
 
-	errChan := make(chan error)
-	go func() {
-		errChan <- cmd.Run()
-	}()
-
 	select {
+	case err := <-errChan:
+		require.NoError(t, err)
 	case <-time.After(5 * time.Second):
 		t.Errorf("failed to wait till the download is started")
 	case <-lock: