Reef/selfupdate/selfupdate.go

// Package selfupdate lets a program replace itself with the newest release published on a
// Gitea server: it checks the latest release, downloads the archive for this platform,
// verifies it against the release's SHA256SUMS, checks that the new binary runs, swaps it
// in and restarts into it.
package selfupdate

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"
)

const maxArchive = 300 << 20

// Event is one step of an update, for showing progress: the stage ("checking", "downloading",
// "verifying", "installing", "restarting", "done" or "failed"), how far along it is (0-100, or -1 to
// leave the bar where it was) and, when there is something worth showing in a log, a line of text.
type Event struct {
	Stage string
	Pct   int
	Line  string
}

// Config says what to update and where from.
type Config struct {
	Server  string // e.g. https://git.veillax.com
	Repo    string // owner/name, e.g. Veillax/blacktip
	App     string // binary and archive prefix, e.g. blacktip
	Current string // the running version, e.g. v0.1.2; anything that isn't vX.Y.Z is never updated
	// Interval between checks (default 6 hours).
	Interval time.Duration
	Logf     func(format string, args ...any)
	// Progress, if set, is told about every step of an update (it is called from the updating goroutine).
	Progress func(Event)

	// The fields below exist for tests and unusual setups.
	HTTP    *http.Client
	GOOS    string
	GOARCH  string
	ExePath string // binary to replace (default: the running one)
	// Verify runs on the downloaded binary before it replaces the old one. The default
	// runs it with -version and checks that it reports the new tag.
	Verify func(path, tag string) error
	// Restart is called after the swap (default: re-exec the new binary in place, keeping the
	// process ID and environment).
	Restart func(exe string) error
}

func (c *Config) defaults() {
	if c.Interval <= 0 {
		c.Interval = 6 * time.Hour
	}
	if c.Logf == nil {
		c.Logf = func(string, ...any) {}
	}
	if c.HTTP == nil {
		c.HTTP = &http.Client{Timeout: 5 * time.Minute}
	}
	if c.GOOS == "" {
		c.GOOS = runtime.GOOS
	}
	if c.GOARCH == "" {
		c.GOARCH = runtime.GOARCH
	}
	if c.Verify == nil {
		c.Verify = verifyRuns(c.App)
	}
	if c.Restart == nil {
		c.Restart = Restart
	}
	c.Server = strings.TrimRight(c.Server, "/")
}

// Run checks shortly after start and then every Interval until ctx ends. When an update
// is installed, Restart takes over and Run normally does not return.
func Run(ctx context.Context, cfg Config) {
	cfg.defaults()
	if _, ok := parseVersion(cfg.Current); !ok {
		cfg.Logf("auto-update: this is a development build (%q), so it will not update itself", cfg.Current)
		return
	}
	if cfg.GOOS == "windows" {
		cfg.Logf("auto-update: not supported on Windows, because a running program can't be replaced; update it by hand")
		return
	}
	wait := time.Duration(30+rand.Intn(240)) * time.Second // spread a fleet's first checks out
	for {
		select {
		case <-ctx.Done():
			return
		case <-time.After(wait):
		}
		if _, err := CheckOnce(ctx, cfg); err != nil {
			cfg.Logf("auto-update: %v", err)
		}
		wait = cfg.Interval + time.Duration(rand.Int63n(int64(cfg.Interval/10)+1))
	}
}

func (c Config) say(stage string, pct int, line string) {
	if c.Progress != nil {
		c.Progress(Event{Stage: stage, Pct: pct, Line: line})
	}
}

// CheckOnce looks for a newer release and installs it. It returns the new tag when it
// installed one (and Restart returned, which only happens in tests or on failure).
func CheckOnce(ctx context.Context, cfg Config) (string, error) {
	cfg.defaults()
	tag, err := checkOnce(ctx, cfg)
	if err != nil {
		cfg.say("failed", -1, "Failed: "+err.Error())
	}
	return tag, err
}

func checkOnce(ctx context.Context, cfg Config) (string, error) {
	cur, ok := parseVersion(cfg.Current)
	if !ok {
		return "", fmt.Errorf("%q is not a release version", cfg.Current)
	}
	cfg.say("checking", 2, "Checking "+cfg.Server+" for the newest release of "+cfg.App)
	rel, err := latestRelease(ctx, cfg)
	if err != nil {
		return "", err
	}
	if rel.Prerelease || rel.Draft {
		cfg.say("done", 100, "Already on the newest release ("+cfg.Current+")")
		return "", nil
	}
	latest, ok := parseVersion(rel.Tag)
	if !ok || !latest.newerThan(cur) {
		cfg.say("done", 100, "Already on the newest release ("+cfg.Current+")")
		return "", nil // up to date, or something we don't understand: never downgrade
	}

	if cfg.GOOS == "windows" {
		return "", errors.New("auto-update is not supported on Windows")
	}
	archive := fmt.Sprintf("%s_%s_%s_%s.tar.gz", cfg.App, rel.Tag, cfg.GOOS, cfg.GOARCH)
	var archiveURL, sumsURL string
	for _, a := range rel.Assets {
		switch a.Name {
		case archive:
			archiveURL = a.URL
		case "SHA256SUMS":
			sumsURL = a.URL
		}
	}
	if archiveURL == "" || sumsURL == "" {
		return "", fmt.Errorf("release %s has no %s with SHA256SUMS yet (it may still be building); will retry", rel.Tag, archive)
	}
	cfg.Logf("auto-update: %s is available (running %s); downloading %s", rel.Tag, cfg.Current, archive)
	cfg.say("checking", 5, rel.Tag+" is available (running "+cfg.Current+")")

	sums, err := get(ctx, cfg.HTTP, sumsURL, 1<<20, nil)
	if err != nil {
		return "", err
	}
	want := ""
	for _, line := range strings.Split(string(sums), "\n") {
		f := strings.Fields(line)
		if len(f) == 2 && strings.TrimPrefix(f[1], "*") == archive {
			want = strings.ToLower(f[0])
		}
	}
	if want == "" {
		return "", fmt.Errorf("SHA256SUMS does not list %s", archive)
	}
	cfg.say("downloading", 6, "Downloading "+archive)
	lastQuarter := 0
	data, err := get(ctx, cfg.HTTP, archiveURL, maxArchive, func(done, total int64) {
		if total <= 0 {
			cfg.say("downloading", -1, "")
			return
		}
		pct := 6 + int(80*done/total)
		line := ""
		if q := int(100 * done / total / 25); q > lastQuarter {
			lastQuarter = q
			line = fmt.Sprintf("Downloaded %.1f of %.1f MB (%d%%)", float64(done)/1e6, float64(total)/1e6, 100*done/total)
		}
		cfg.say("downloading", pct, line)
	})
	if err != nil {
		return "", err
	}
	cfg.say("verifying", 88, "Verifying the checksum")
	got := sha256.Sum256(data)
	if hex.EncodeToString(got[:]) != want {
		return "", fmt.Errorf("checksum mismatch for %s: not installing", archive)
	}
	cfg.say("verifying", 90, "The checksum matches")
	cfg.say("installing", 92, "Unpacking the new binary")
	bin, err := extract(data, cfg.App)
	if err != nil {
		return "", err
	}

	exe := cfg.ExePath
	if exe == "" {
		if exe, err = os.Executable(); err != nil {
			return "", err
		}
		if exe, err = filepath.EvalSymlinks(exe); err != nil {
			return "", err
		}
	}
	tmp := filepath.Join(filepath.Dir(exe), "."+filepath.Base(exe)+".new")
	if err := os.WriteFile(tmp, bin, 0o755); err != nil {
		return "", fmt.Errorf("can't write next to %s (is its folder writable by this user?): %w", exe, err)
	}
	cfg.say("installing", 94, "Checking that the new binary starts")
	if err := cfg.Verify(tmp, rel.Tag); err != nil {
		os.Remove(tmp)
		return "", fmt.Errorf("the new binary failed its check, keeping the old one: %w", err)
	}
	cfg.say("installing", 97, "Replacing "+exe)
	if err := os.Rename(tmp, exe); err != nil { // atomic; the running process keeps its old file
		os.Remove(tmp)
		return "", fmt.Errorf("replacing %s: %w", exe, err)
	}
	cfg.Logf("auto-update: installed %s, restarting", rel.Tag)
	cfg.say("restarting", 100, "Installed "+rel.Tag+"; restarting into it")
	if err := cfg.Restart(exe); err != nil {
		return rel.Tag, fmt.Errorf("installed %s but could not restart into it (it will start with the next restart): %w", rel.Tag, err)
	}
	return rel.Tag, nil
}

type release struct {
	Tag        string `json:"tag_name"`
	Prerelease bool   `json:"prerelease"`
	Draft      bool   `json:"draft"`
	Assets     []struct {
		Name string `json:"name"`
		URL  string `json:"browser_download_url"`
	} `json:"assets"`
}

func latestRelease(ctx context.Context, cfg Config) (release, error) {
	var rel release
	body, err := get(ctx, cfg.HTTP, cfg.Server+"/api/v1/repos/"+cfg.Repo+"/releases/latest", 1<<20, nil)
	if err != nil {
		return rel, fmt.Errorf("checking for a new release: %w", err)
	}
	if err := json.Unmarshal(body, &rel); err != nil {
		return rel, fmt.Errorf("reading the latest release: %w", err)
	}
	return rel, nil
}

// Check reports the newest full release and whether it is newer than cfg.Current, without
// installing anything. A development build (a Current that isn't vX.Y.Z) is never "newer".
func Check(ctx context.Context, cfg Config) (latest string, newer bool, err error) {
	cfg.defaults()
	rel, err := latestRelease(ctx, cfg)
	if err != nil {
		return "", false, err
	}
	if rel.Prerelease || rel.Draft {
		return "", false, nil
	}
	l, ok := parseVersion(rel.Tag)
	if !ok {
		return rel.Tag, false, nil
	}
	cur, ok := parseVersion(cfg.Current)
	return rel.Tag, ok && l.newerThan(cur), nil
}

func get(ctx context.Context, c *http.Client, url string, max int64, onProgress func(done, total int64)) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "sharks-selfupdate")
	resp, err := c.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("%s: %s", url, resp.Status)
	}
	var b []byte
	buf := make([]byte, 64<<10)
	for {
		n, err := resp.Body.Read(buf)
		b = append(b, buf[:n]...)
		if int64(len(b)) > max {
			return nil, fmt.Errorf("%s is larger than %d bytes", url, max)
		}
		if onProgress != nil && n > 0 {
			onProgress(int64(len(b)), resp.ContentLength)
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
	}
	return b, nil
}

// extract finds the program inside a release archive (<pkg>/<app>).
func extract(data []byte, app string) ([]byte, error) {
	zr, err := gzip.NewReader(bytes.NewReader(data))
	if err != nil {
		return nil, err
	}
	tr := tar.NewReader(zr)
	for {
		h, err := tr.Next()
		if err == io.EOF {
			return nil, fmt.Errorf("no %s inside the archive", app)
		}
		if err != nil {
			return nil, err
		}
		if h.Typeflag == tar.TypeReg && path.Base(h.Name) == app && strings.Count(h.Name, "/") == 1 {
			return io.ReadAll(io.LimitReader(tr, maxArchive))
		}
	}
}

func verifyRuns(app string) func(path, tag string) error {
	return func(p, tag string) error {
		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
		defer cancel()
		out, err := exec.CommandContext(ctx, p, "-version").Output()
		if err != nil {
			return err
		}
		if want := app + " " + tag; strings.TrimSpace(string(out)) != want {
			return fmt.Errorf("it reports %q, expected %q", strings.TrimSpace(string(out)), want)
		}
		return nil
	}
}

type version [3]int

func parseVersion(s string) (version, bool) {
	var v version
	if !strings.HasPrefix(s, "v") || strings.ContainsAny(s, "-+") {
		return v, false // pre-releases and local builds are never part of auto-update
	}
	parts := strings.Split(s[1:], ".")
	if len(parts) != 3 {
		return v, false
	}
	for i, p := range parts {
		n, err := strconv.Atoi(p)
		if err != nil || n < 0 {
			return v, false
		}
		v[i] = n
	}
	return v, true
}

func (a version) newerThan(b version) bool {
	for i := range a {
		if a[i] != b[i] {
			return a[i] > b[i]
		}
	}
	return false
}