spb/internetpressure

1// Package update implements the optional self-update: GET /agent/latest, download the asset for GOOS-GOARCH,2// verify its sha256, atomically replace the running binary and ask the caller to exit 0 (the supervisor —3// launchd / systemd — restarts the new binary).4package update56import (7	"context"8	"crypto/sha256"9	"encoding/hex"10	"errors"11	"fmt"12	"io"13	"log/slog"14	"os"15	"path/filepath"16	"runtime"17	"strings"18	"time"1920	"internetpressure.io/probe-agent/internal/client"21	"internetpressure.io/probe-agent/internal/protocol"22)2324// Defaults.25const (26	CheckEvery      = 6 * time.Hour27	FirstCheckAfter = 2 * time.Minute28	DownloadTimeout = 5 * time.Minute29	MaxAssetBytes   = 256 << 2030	lastAttemptFile = "update.last"31	tmpFile         = "update.tmp"32)3334// Fetcher is the subset of client.Client used (interface for tests).35type Fetcher interface {36	AgentLatest(ctx context.Context) (*protocol.AgentLatest, error)37	Download(ctx context.Context, url string, w io.Writer, limit int64) (int64, error)38}3940// Updater checks for and applies releases.41type Updater struct {42	Fetcher Fetcher43	Version string44	DataDir string45	// ExePath is the binary to replace; "" → os.Executable() (resolved through symlinks).46	ExePath string47	Log     *slog.Logger48}4950// Platform key used in the assets map.51func Platform() string { return runtime.GOOS + "-" + runtime.GOARCH }5253// Check fetches /agent/latest and applies an update when one is available. Returns true when the binary was54// replaced (the caller must exit 0).55func (u *Updater) Check(ctx context.Context) (bool, error) {56	log := u.Log57	if log == nil {58		log = slog.Default()59	}60	latest, err := u.Fetcher.AgentLatest(ctx)61	if err != nil {62		return false, err63	}64	if latest.Version == "" || latest.Version == u.Version {65		return false, nil66	}67	asset, ok := latest.Assets[Platform()]68	if !ok || asset.URL == "" || asset.SHA256 == "" {69		log.Info("update available but no asset for this platform", "version", latest.Version, "platform", Platform())70		return false, nil71	}72	// Downgrade/loop guard: never attempt the same version twice (covers a bad release that keeps crashing73	// and being "updated" to, or a server offering an older version).74	if last := u.lastAttempt(); last == latest.Version {75		log.Debug("skipping update already attempted", "version", latest.Version)76		return false, nil77	}78	if err := u.recordAttempt(latest.Version); err != nil {79		return false, err80	}81	log.Info("downloading update", "from", u.Version, "to", latest.Version, "url", asset.URL)8283	exe := u.ExePath84	if exe == "" {85		p, err := os.Executable()86		if err != nil {87			return false, err88		}89		if r, err := filepath.EvalSymlinks(p); err == nil {90			p = r91		}92		exe = p93	}9495	tmp := filepath.Join(u.DataDir, tmpFile)96	if err := u.download(ctx, asset, tmp); err != nil {97		os.Remove(tmp)98		return false, err99	}100	if err := os.Chmod(tmp, 0o755); err != nil {101		os.Remove(tmp)102		return false, err103	}104	// Atomic replace: rename within the same filesystem. If data_dir is on another filesystem, stage a copy105	// next to the binary first.106	if err := os.Rename(tmp, exe); err != nil {107		staged := exe + ".new"108		if cerr := copyFile(tmp, staged, 0o755); cerr != nil {109			os.Remove(tmp)110			return false, fmt.Errorf("stage update: %w", cerr)111		}112		os.Remove(tmp)113		if err := os.Rename(staged, exe); err != nil {114			os.Remove(staged)115			return false, fmt.Errorf("replace binary: %w", err)116		}117	}118	log.Info("update installed, exiting for restart", "version", latest.Version, "binary", exe)119	return true, nil120}121122func (u *Updater) download(ctx context.Context, asset protocol.Asset, dst string) error {123	ctx, cancel := context.WithTimeout(ctx, DownloadTimeout)124	defer cancel()125	f, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o700)126	if err != nil {127		return err128	}129	h := sha256.New()130	n, err := u.Fetcher.Download(ctx, asset.URL, io.MultiWriter(f, h), MaxAssetBytes)131	if cerr := f.Close(); err == nil {132		err = cerr133	}134	if err != nil {135		return fmt.Errorf("download: %w", err)136	}137	if n == 0 {138		return errors.New("download: empty asset")139	}140	got := hex.EncodeToString(h.Sum(nil))141	if !strings.EqualFold(got, strings.TrimSpace(asset.SHA256)) {142		return fmt.Errorf("download: sha256 mismatch (got %s, want %s)", got, asset.SHA256)143	}144	return nil145}146147func (u *Updater) lastAttempt() string {148	b, err := os.ReadFile(filepath.Join(u.DataDir, lastAttemptFile))149	if err != nil {150		return ""151	}152	return strings.TrimSpace(string(b))153}154155func (u *Updater) recordAttempt(v string) error {156	return os.WriteFile(filepath.Join(u.DataDir, lastAttemptFile), []byte(v+"\n"), 0o600)157}158159func copyFile(src, dst string, mode os.FileMode) error {160	in, err := os.Open(src)161	if err != nil {162		return err163	}164	defer in.Close()165	out, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, mode)166	if err != nil {167		return err168	}169	if _, err := io.Copy(out, in); err != nil {170		out.Close()171		return err172	}173	return out.Close()174}175176// Run checks after FirstCheckAfter and then every CheckEvery; onUpdated is called once the binary was replaced.177func (u *Updater) Run(ctx context.Context, onUpdated func()) {178	log := u.Log179	if log == nil {180		log = slog.Default()181	}182	wait := FirstCheckAfter183	for {184		select {185		case <-ctx.Done():186			return187		case <-time.After(wait):188		}189		wait = CheckEvery190		updated, err := u.Check(ctx)191		if err != nil {192			var he *client.HTTPError193			if errors.As(err, &he) && he.Status == 404 {194				log.Debug("no release feed", "err", err)195			} else {196				log.Warn("update check failed", "err", err)197			}198			continue199		}200		if updated && onUpdated != nil {201			onUpdated()202			return203		}204	}205}206