fix(update): compare versions as semver and replace the binary safely

cmd/version.go

@@ -30,20 +30,24 @@ var updateCmd = &cobra.Command{
 			return errors.New("the update command must be run as root, please run 'sudo fly update'")
 		}
 
-		latestVersion, hasUpdate, err := utils.CheckForUpdates()
+		update, err := utils.CheckForUpdates(cmd.Context())
 		if err != nil {
 			return fmt.Errorf("checking for updates: %w", err)
 		}
 
-		if !hasUpdate {
+		latest := update.Release.TagName
+		switch {
+		case !update.Comparable:
+			fmt.Printf("This is not a release build (version %s). Latest release: %s\n", version.Version, latest)
+		case !update.Available:
 			fmt.Println("You are already running the latest version.")
 			return nil
+		default:
+			fmt.Printf("New version available: %s\n", latest)
 		}
 
-		fmt.Printf("New version available: %s\n", latestVersion)
-
 		if !yesFlag {
-			fmt.Print("Do you want to update? (y/n): ")
+			fmt.Printf("Do you want to install %s? (y/n): ", latest)
 			var response string
 			// An empty or unreadable answer cancels the update.
 			_, _ = fmt.Scanln(&response)
@@ -54,11 +58,11 @@ var updateCmd = &cobra.Command{
 		}
 
 		fmt.Println("Updating...")
-		if err := utils.SelfUpdate(); err != nil {
+		if err := utils.SelfUpdate(cmd.Context(), update.Release); err != nil {
 			return fmt.Errorf("updating: %w", err)
 		}
 
-		fmt.Println("Update successful. Please restart fly cli.")
+		fmt.Printf("Updated to %s.\n", latest)
 		return nil
 	},
 }

go.mod

@@ -7,6 +7,7 @@ toolchain go1.27.1
 require (
 	github.com/fatih/color v1.19.0
 	github.com/spf13/cobra v1.10.2
+	golang.org/x/mod v0.41.0
 	gopkg.in/yaml.v2 v2.4.0
 )
 

go.sum

@@ -14,6 +14,8 @@ github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An
 github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk=
 github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
+golang.org/x/mod v0.41.0 h1:qJmnOUb4YB+FsEuM3HcWucdZASCPGhsX6uljO6pog0c=
+golang.org/x/mod v0.41.0/go.mod h1:Ek9pY8RKWXwsWvd3rQiHYtMqkjSUV+s1Rj7j4H5Ur6o=
 golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo=
 golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=

internal/utils/version.go

@@ -1,19 +1,33 @@
 package utils
 
 import (
+	"archive/tar"
+	"compress/gzip"
+	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"os"
-	"os/exec"
+	"path"
 	"path/filepath"
+	"regexp"
 	"runtime"
+	"time"
 
 	"github.com/flywp/server-cli/internal/version"
+	"golang.org/x/mod/semver"
 )
 
-const GithubAPI = "https://api.github.com/repos/flywp/server-cli/releases/latest"
+// GithubAPI is the GitHub API URL of the latest release.
+var GithubAPI = "https://api.github.com/repos/flywp/server-cli/releases/latest"
+
+// httpClient limits each request, including the download of the binary.
+var httpClient = &http.Client{Timeout: 60 * time.Second}
+
+// maxBinarySize limits the size of the binary in a release archive.
+const maxBinarySize = 200 << 20
 
 type GithubRelease struct {
 	TagName string `json:"tag_name"`
@@ -23,136 +37,195 @@ type GithubRelease struct {
 	} `json:"assets"`
 }
 
-func CheckForUpdates() (string, bool, error) {
-	resp, err := http.Get(GithubAPI)
-	if err != nil {
-		return "", false, err
-	}
-	defer resp.Body.Close()
+// Update compares the latest release with the running version.
+type Update struct {
+	Release *GithubRelease
+	// Available is true when the release is newer than the running version.
+	Available bool
+	// Comparable is false when the running version is not built from a
+	// release tag, for example "dev".
+	Comparable bool
+}
 
-	body, err := io.ReadAll(resp.Body)
+// CheckForUpdates gets the latest release from GitHub and compares it with
+// the running version.
+func CheckForUpdates(ctx context.Context) (*Update, error) {
+	release, err := LatestRelease(ctx)
 	if err != nil {
-		return "", false, err
+		return nil, err
 	}
 
-	var release GithubRelease
-	if err := json.Unmarshal(body, &release); err != nil {
-		return "", false, err
+	available, comparable := isNewer(release.TagName, version.Version)
+	return &Update{Release: release, Available: available, Comparable: comparable}, nil
+}
+
+// describeSuffix matches what git describe adds after a tag: the number of
+// commits since the tag, the commit hash and "-dirty" for local changes.
+var describeSuffix = regexp.MustCompile(`(-\d+-g[0-9a-f]+)?(-dirty)?$`)
+
+// isNewer reports whether the release version latest is newer than current.
+// comparable is false when current is not built from a release tag.
+func isNewer(latest, current string) (newer, comparable bool) {
+	base := describeSuffix.ReplaceAllString(current, "")
+	if !semver.IsValid(base) {
+		return false, false
 	}
 
-	return release.TagName, release.TagName > version.Version, nil
+	return semver.Compare(latest, base) > 0, true
 }
 
-func SelfUpdate() error {
-	if os.Geteuid() != 0 {
-		return fmt.Errorf("the update command must be run as root")
+// LatestRelease returns the latest release from GitHub.
+func LatestRelease(ctx context.Context) (*GithubRelease, error) {
+	resp, err := get(ctx, GithubAPI)
+	if err != nil {
+		return nil, err
 	}
+	defer func() { _ = resp.Body.Close() }()
 
-	release, err := getLatestRelease()
-	if err != nil {
-		return fmt.Errorf("failed to get latest release: %w", err)
+	var release GithubRelease
+	if err := json.NewDecoder(resp.Body).Decode(&release); err != nil {
+		return nil, fmt.Errorf("reading release information: %w", err)
 	}
 
-	assetURL := getAssetURL(release)
-	if assetURL == "" {
-		return fmt.Errorf("no suitable binary found for this system (OS: %s, ARCH: %s)", runtime.GOOS, runtime.GOARCH)
+	if !semver.IsValid(release.TagName) {
+		return nil, fmt.Errorf("latest release has an invalid version %q", release.TagName)
 	}
 
-	// Create a temporary directory
-	tmpDir, err := os.MkdirTemp("", "fly-cli-update")
+	return &release, nil
+}
+
+// get sends a GET request and returns the response if its status is 200 OK.
+func get(ctx context.Context, url string) (*http.Response, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
 	if err != nil {
-		return fmt.Errorf("failed to create temp directory: %w", err)
+		return nil, err
 	}
-	defer os.RemoveAll(tmpDir)
+	req.Header.Set("Accept", "application/vnd.github+json")
+	req.Header.Set("User-Agent", "fly-cli/"+version.Version)
 
-	// Download the archive
-	resp, err := http.Get(assetURL)
+	resp, err := httpClient.Do(req)
 	if err != nil {
-		return fmt.Errorf("failed to download update: %w", err)
+		return nil, err
 	}
-	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
-		return fmt.Errorf("failed to download update: HTTP %d", resp.StatusCode)
-	}
+		_ = resp.Body.Close()
 
-	// Create the archive file
-	archivePath := filepath.Join(tmpDir, "update.tar.gz")
-	out, err := os.Create(archivePath)
-	if err != nil {
-		return fmt.Errorf("failed to create archive file: %w", err)
+		limited := resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusTooManyRequests
+		if limited && resp.Header.Get("X-RateLimit-Remaining") == "0" {
+			return nil, errors.New("the GitHub API rate limit is exceeded, try again later")
+		}
+		return nil, fmt.Errorf("unexpected response from %s: %s", url, resp.Status)
 	}
 
-	// Write the body to file
-	_, err = io.Copy(out, resp.Body)
-	out.Close()
-	if err != nil {
-		return fmt.Errorf("failed to write archive file: %w", err)
-	}
+	return resp, nil
+}
 
-	// Extract the archive
-	binaryName := fmt.Sprintf("fly-%s-%s", runtime.GOOS, runtime.GOARCH)
-	cmd := exec.Command("tar", "-xzf", archivePath, "-C", tmpDir)
-	if err := cmd.Run(); err != nil {
-		return fmt.Errorf("failed to extract archive: %w", err)
+// SelfUpdate replaces the running binary with the binary from release.
+func SelfUpdate(ctx context.Context, release *GithubRelease) error {
+	assetURL := assetURL(release, runtime.GOOS, runtime.GOARCH)
+	if assetURL == "" {
+		return fmt.Errorf("no suitable binary found for this system (OS: %s, ARCH: %s)", runtime.GOOS, runtime.GOARCH)
 	}
 
-	// Get the current executable path
 	exe, err := os.Executable()
 	if err != nil {
-		return fmt.Errorf("failed to get current executable path: %w", err)
+		return fmt.Errorf("finding the current executable: %w", err)
 	}
 	exe, err = filepath.EvalSymlinks(exe)
 	if err != nil {
-		return fmt.Errorf("failed to resolve symlinks: %w", err)
+		return fmt.Errorf("resolving symlinks: %w", err)
 	}
 
-	// Make the new binary executable
-	extractedBinary := filepath.Join(tmpDir, binaryName)
-	if err := os.Chmod(extractedBinary, 0755); err != nil {
-		return fmt.Errorf("failed to make binary executable: %w", err)
+	resp, err := get(ctx, assetURL)
+	if err != nil {
+		return fmt.Errorf("downloading update: %w", err)
 	}
+	defer func() { _ = resp.Body.Close() }()
 
-	// Rename the temporary file to the executable name
-	if err := os.Rename(extractedBinary, exe); err != nil {
-		return fmt.Errorf("failed to replace old binary: %w", err)
-	}
+	return replaceBinary(exe, resp.Body, binaryName(runtime.GOOS, runtime.GOARCH))
+}
 
-	return nil
+// binaryName is the name of the binary in a release archive. Releases must
+// keep this name: installed versions of fly look for it.
+func binaryName(goos, goarch string) string {
+	return fmt.Sprintf("fly-%s-%s", goos, goarch)
 }
 
-func getLatestRelease() (*GithubRelease, error) {
-	resp, err := http.Get(GithubAPI)
-	if err != nil {
-		return nil, err
+// assetURL returns the download URL of the release archive for goos and
+// goarch, or "" if the release has none.
+func assetURL(release *GithubRelease, goos, goarch string) string {
+	if goos != "linux" {
+		return ""
 	}
-	defer resp.Body.Close()
 
-	body, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return nil, err
+	expectedName := binaryName(goos, goarch) + ".tar.gz"
+	for _, asset := range release.Assets {
+		if asset.Name == expectedName {
+			return asset.BrowserDownloadURL
+		}
 	}
 
-	var release GithubRelease
-	if err := json.Unmarshal(body, &release); err != nil {
-		return nil, err
+	return ""
+}
+
+// replaceBinary extracts the file name from the tar.gz archive and puts it in
+// place of exe. The new binary is written to a temporary file in the same
+// directory and then renamed, so exe is never incomplete and the rename
+// does not cross filesystems.
+func replaceBinary(exe string, archive io.Reader, name string) error {
+	gz, err := gzip.NewReader(archive)
+	if err != nil {
+		return fmt.Errorf("reading archive: %w", err)
 	}
+	defer func() { _ = gz.Close() }()
 
-	return &release, nil
-}
+	tr := tar.NewReader(gz)
+	for {
+		hdr, err := tr.Next()
+		if errors.Is(err, io.EOF) {
+			return fmt.Errorf("archive does not contain %s", name)
+		}
+		if err != nil {
+			return fmt.Errorf("reading archive: %w", err)
+		}
 
-func getAssetURL(release *GithubRelease) string {
-	arch := runtime.GOARCH
-	if runtime.GOOS != "linux" {
-		return ""
+		if hdr.Typeflag != tar.TypeReg || path.Clean(hdr.Name) != name {
+			continue
+		}
+		if hdr.Size > maxBinarySize {
+			return fmt.Errorf("%s in archive is too large (%d bytes)", name, hdr.Size)
+		}
+
+		return writeBinary(exe, tr)
 	}
+}
 
-	expectedName := fmt.Sprintf("fly-linux-%s.tar.gz", arch)
-	for _, asset := range release.Assets {
-		if asset.Name == expectedName {
-			return asset.BrowserDownloadURL
+// writeBinary writes r to a temporary file next to exe and renames it to exe.
+func writeBinary(exe string, r io.Reader) (err error) {
+	tmp, err := os.CreateTemp(filepath.Dir(exe), ".fly-update-*")
+	if err != nil {
+		return fmt.Errorf("creating temporary file: %w", err)
+	}
+	defer func() {
+		if err != nil {
+			_ = os.Remove(tmp.Name())
 		}
+	}()
+
+	if _, err = io.Copy(tmp, r); err != nil {
+		_ = tmp.Close()
+		return fmt.Errorf("writing new binary: %w", err)
+	}
+	if err = tmp.Close(); err != nil {
+		return fmt.Errorf("writing new binary: %w", err)
+	}
+	if err = os.Chmod(tmp.Name(), 0o755); err != nil {
+		return fmt.Errorf("making binary executable: %w", err)
+	}
+	if err = os.Rename(tmp.Name(), exe); err != nil {
+		return fmt.Errorf("replacing binary: %w", err)
 	}
 
-	return ""
+	return nil
 }