Export selfupdate.Restart

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

selfupdate/restart_unix.go

@@ -0,0 +1,13 @@
+//go:build !windows
+
+package selfupdate
+
+import (
+	"os"
+	"syscall"
+)
+
+// Restart replaces this process with the new binary, with the same arguments and environment.
+func Restart(exe string) error {
+	return syscall.Exec(exe, os.Args, os.Environ())
+}

selfupdate/restart_windows.go

@@ -0,0 +1,7 @@
+package selfupdate
+
+import "errors"
+
+func Restart(string) error {
+	return errors.New("restarting in place is not supported on Windows")
+}

selfupdate/selfupdate.go

@@ -0,0 +1,299 @@
+// 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
+
+// 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)
+
+	// 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))
+	}
+}
+
+// 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()
+	cur, ok := parseVersion(cfg.Current)
+	if !ok {
+		return "", fmt.Errorf("%q is not a release version", cfg.Current)
+	}
+	var rel 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"`
+	}
+	body, err := get(ctx, cfg.HTTP, cfg.Server+"/api/v1/repos/"+cfg.Repo+"/releases/latest", 1<<20)
+	if err != nil {
+		return "", fmt.Errorf("checking for a new release: %w", err)
+	}
+	if err := json.Unmarshal(body, &rel); err != nil {
+		return "", fmt.Errorf("reading the latest release: %w", err)
+	}
+	if rel.Prerelease || rel.Draft {
+		return "", nil
+	}
+	latest, ok := parseVersion(rel.Tag)
+	if !ok || !latest.newerThan(cur) {
+		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)
+
+	sums, err := get(ctx, cfg.HTTP, sumsURL, 1<<20)
+	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)
+	}
+	data, err := get(ctx, cfg.HTTP, archiveURL, maxArchive)
+	if err != nil {
+		return "", err
+	}
+	got := sha256.Sum256(data)
+	if hex.EncodeToString(got[:]) != want {
+		return "", fmt.Errorf("checksum mismatch for %s: not installing", archive)
+	}
+	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)
+	}
+	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)
+	}
+	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)
+	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
+}
+
+func get(ctx context.Context, c *http.Client, url string, max 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)
+	}
+	b, err := io.ReadAll(io.LimitReader(resp.Body, max+1))
+	if err != nil {
+		return nil, err
+	}
+	if int64(len(b)) > max {
+		return nil, fmt.Errorf("%s is larger than %d bytes", url, max)
+	}
+	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
+}

selfupdate/selfupdate_test.go

@@ -0,0 +1,137 @@
+package selfupdate
+
+import (
+	"archive/tar"
+	"bytes"
+	"compress/gzip"
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+)
+
+func makeArchive(t *testing.T, pkg, app, content string) []byte {
+	t.Helper()
+	var buf bytes.Buffer
+	gz := gzip.NewWriter(&buf)
+	tw := tar.NewWriter(gz)
+	for name, body := range map[string]string{pkg + "/README.md": "hi", pkg + "/" + app: content} {
+		if err := tw.WriteHeader(&tar.Header{Name: name, Mode: 0o755, Size: int64(len(body)), Typeflag: tar.TypeReg}); err != nil {
+			t.Fatal(err)
+		}
+		tw.Write([]byte(body))
+	}
+	tw.Close()
+	gz.Close()
+	return buf.Bytes()
+}
+
+// fakeGitea serves one release with the given tag.
+func fakeGitea(t *testing.T, tag string, prerelease bool, corrupt bool) *httptest.Server {
+	archive := fmt.Sprintf("demo_%s_linux_amd64.tar.gz", tag)
+	data := makeArchive(t, strings.TrimSuffix(archive, ".tar.gz"), "demo", "NEW-BINARY "+tag)
+	sum := sha256.Sum256(data)
+	sums := hex.EncodeToString(sum[:]) + "  " + archive + "\n"
+	if corrupt {
+		sums = strings.Repeat("0", 64) + "  " + archive + "\n"
+	}
+	var srv *httptest.Server
+	mux := http.NewServeMux()
+	mux.HandleFunc("/api/v1/repos/Veillax/demo/releases/latest", func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintf(w, `{"tag_name":%q,"prerelease":%v,"assets":[{"name":%q,"browser_download_url":"%s/dl/%s"},{"name":"SHA256SUMS","browser_download_url":"%s/dl/SHA256SUMS"}]}`,
+			tag, prerelease, archive, srv.URL, archive, srv.URL)
+	})
+	mux.HandleFunc("/dl/SHA256SUMS", func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, sums) })
+	mux.HandleFunc("/dl/"+archive, func(w http.ResponseWriter, r *http.Request) { w.Write(data) })
+	srv = httptest.NewServer(mux)
+	t.Cleanup(srv.Close)
+	return srv
+}
+
+func setup(t *testing.T, srv *httptest.Server, current string) (Config, string, *string) {
+	exe := filepath.Join(t.TempDir(), "demo")
+	os.WriteFile(exe, []byte("OLD-BINARY"), 0o755)
+	restarted := new(string)
+	return Config{Server: srv.URL, Repo: "Veillax/demo", App: "demo", Current: current, GOOS: "linux", GOARCH: "amd64", ExePath: exe,
+		Verify:  func(p, tag string) error { return nil },
+		Restart: func(e string) error { *restarted = e; return nil }}, exe, restarted
+}
+
+func TestInstallsANewerRelease(t *testing.T) {
+	cfg, exe, restarted := setup(t, fakeGitea(t, "v0.2.0", false, false), "v0.1.5")
+	tag, err := CheckOnce(context.Background(), cfg)
+	if err != nil || tag != "v0.2.0" {
+		t.Fatalf("tag=%q err=%v", tag, err)
+	}
+	if b, _ := os.ReadFile(exe); string(b) != "NEW-BINARY v0.2.0" {
+		t.Fatalf("binary is %q", b)
+	}
+	if *restarted != exe {
+		t.Fatalf("restart called with %q", *restarted)
+	}
+	if _, err := os.Stat(filepath.Join(filepath.Dir(exe), ".demo.new")); err == nil {
+		t.Fatal("temporary file left behind")
+	}
+}
+
+func TestLeavesThingsAloneWhenItShould(t *testing.T) {
+	for _, c := range []struct {
+		name, current, tag string
+		pre                bool
+	}{
+		{"same version", "v0.2.0", "v0.2.0", false},
+		{"never downgrades", "v0.3.0", "v0.2.0", false},
+		{"pre-release is skipped", "v0.1.0", "v0.2.0", true},
+		{"dev build", "dev", "v0.2.0", false},
+		{"local build", "v0.2.0-local", "v0.9.0", false},
+	} {
+		cfg, exe, restarted := setup(t, fakeGitea(t, c.tag, c.pre, false), c.current)
+		tag, _ := CheckOnce(context.Background(), cfg)
+		if b, _ := os.ReadFile(exe); tag != "" || string(b) != "OLD-BINARY" || *restarted != "" {
+			t.Errorf("%s: changed things (tag=%q binary=%q)", c.name, tag, b)
+		}
+	}
+}
+
+func TestRefusesABadChecksum(t *testing.T) {
+	cfg, exe, restarted := setup(t, fakeGitea(t, "v0.2.0", false, true), "v0.1.0")
+	if _, err := CheckOnce(context.Background(), cfg); err == nil || !strings.Contains(err.Error(), "checksum") {
+		t.Fatalf("err=%v", err)
+	}
+	if b, _ := os.ReadFile(exe); string(b) != "OLD-BINARY" || *restarted != "" {
+		t.Fatal("installed a binary that failed its checksum")
+	}
+}
+
+func TestKeepsTheOldBinaryWhenTheNewOneFailsItsCheck(t *testing.T) {
+	cfg, exe, restarted := setup(t, fakeGitea(t, "v0.2.0", false, false), "v0.1.0")
+	cfg.Verify = func(p, tag string) error { return fmt.Errorf("segfault") }
+	if _, err := CheckOnce(context.Background(), cfg); err == nil {
+		t.Fatal("expected an error")
+	}
+	if b, _ := os.ReadFile(exe); string(b) != "OLD-BINARY" || *restarted != "" {
+		t.Fatal("replaced the binary despite a failed check")
+	}
+	if _, err := os.Stat(filepath.Join(filepath.Dir(exe), ".demo.new")); err == nil {
+		t.Fatal("temporary file left behind")
+	}
+}
+
+func TestVersionOrdering(t *testing.T) {
+	for _, c := range []struct {
+		a, b string
+		want bool
+	}{{"v0.1.10", "v0.1.9", true}, {"v1.0.0", "v0.99.99", true}, {"v0.1.2", "v0.1.2", false}, {"v0.1.1", "v0.2.0", false}} {
+		a, _ := parseVersion(c.a)
+		b, _ := parseVersion(c.b)
+		if a.newerThan(b) != c.want {
+			t.Errorf("%s newer than %s: want %v", c.a, c.b, c.want)
+		}
+	}
+}