internal/agent/selfupdate/selfupdate.go

Ref: main v0.0.5   Size: 10.6 KiB   History

// Package selfupdate replaces the running agent binary with a
// server-instructed release and re-execs. Every step is level-triggered
// retry-safe: a failure leaves the current binary running and untouched, and
// the next snapshot carrying the offer tries again. The previous binary is
// kept beside the new one as "<exe>.prev" for manual recovery.
//
// The artifact a manifest names is either a bare binary or a release tarball
// carrying the agent as <bundle-dir>/eitri-agent. Which one it is comes from
// the bytes — gzip magic — never from the URL's suffix, which a mirror or a
// redirect is free to dress up however it likes.
package selfupdate

import (
	"archive/tar"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"
	"syscall"
	"time"
)

// agentMember is the file a release tarball carries the agent as. Bundles lay
// it out under a versioned directory, so the member is matched on basename.
const agentMember = "eitri-agent"

// Update names the target binary: its version (for logging), artifact URL,
// and expected sha256 (hex).
type Update struct {
	Version, URL, SHA256 string
}

// Applier performs the swap. Zero value works in production; tests inject the
// seams.
type Applier struct {
	// HTTP is the download client (nil ⇒ a 10-minute-timeout default,
	// mirroring the image fetch bound).
	HTTP *http.Client
	// Exec replaces the process image (nil ⇒ syscall.Exec). Tests capture it.
	Exec func(argv0 string, argv, env []string) error
	// ExePath resolves the running binary's path (nil ⇒ os.Executable).
	ExePath func() (string, error)
}

func (a *Applier) httpClient() *http.Client {
	if a.HTTP != nil {
		return a.HTTP
	}
	return &http.Client{Timeout: 10 * time.Minute}
}

func (a *Applier) execFn() func(argv0 string, argv, env []string) error {
	if a.Exec != nil {
		return a.Exec
	}
	return syscall.Exec
}

// Apply downloads, verifies, swaps, and re-execs. On any error before the
// final rename the running binary is untouched; after the rename the process
// re-execs (or returns the exec error — at that point <exe> is already the
// new binary and <exe>.prev the old one).
//
// ctx governs the download only (via http.NewRequestWithContext); if it is
// ever cancelled mid-download, the download aborts, Apply returns an error,
// the binary is untouched, and the next attempt starts over — a killed
// download is never resumed. Whether that happens on every reconnect or only
// on full process shutdown is the CALLER's choice of which ctx to pass;
// syncclient.Client.maybeUpgrade deliberately uses a ctx that outlives one
// QUIC session so a reconnect blip doesn't abandon a half-finished download.
func (a *Applier) Apply(ctx context.Context, u Update) error {
	exePath := os.Executable
	if a.ExePath != nil {
		exePath = a.ExePath
	}
	exe, err := exePath()
	if err != nil {
		return fmt.Errorf("resolve executable: %w", err)
	}
	dir := filepath.Dir(exe)

	// Sweep stale temps from a crashed prior attempt: the deferred os.Remove
	// below never runs if THAT process was killed mid-download, so litter can
	// accumulate across restarts. Best-effort; a removal failure here is not
	// fatal to this attempt.
	if stale, globErr := filepath.Glob(filepath.Join(dir, ".eitri-agent-upgrade-*")); globErr == nil {
		for _, f := range stale {
			_ = os.Remove(f)
		}
	}

	// Idempotency: if <exe> ALREADY has the target sha, the swap already
	// happened — most likely syscall.Exec failed on a previous Apply AFTER
	// the rename succeeded, so the running process is still the old version
	// and the server keeps re-offering the same upgrade. Skip straight to
	// exec: re-running download+copy+rename here would overwrite <exe>.prev
	// with the NEW binary (this exe IS the new binary already), destroying
	// the actual previous-version recovery copy for no reason.
	if u.SHA256 != "" {
		if sum, sumErr := sha256File(exe); sumErr == nil && sum == u.SHA256 {
			return a.execFn()(exe, os.Args, os.Environ())
		}
	}

	// The temp file lives NEXT TO the binary (same filesystem) so the final
	// os.Rename is atomic — the state dir may be a different mount.
	tmp, err := os.CreateTemp(dir, ".eitri-agent-upgrade-*")
	if err != nil {
		return fmt.Errorf("create temp beside binary: %w", err)
	}
	defer os.Remove(tmp.Name()) // no-op after the successful rename

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.URL, nil)
	if err != nil {
		tmp.Close()
		return err
	}
	resp, err := a.httpClient().Do(req)
	if err != nil {
		tmp.Close()
		return fmt.Errorf("download %s: %w", u.URL, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		tmp.Close()
		return fmt.Errorf("download %s: HTTP %d", u.URL, resp.StatusCode)
	}
	h := sha256.New()
	if _, err := io.Copy(io.MultiWriter(tmp, h), resp.Body); err != nil {
		tmp.Close()
		return fmt.Errorf("download %s: %w", u.URL, err)
	}
	// Durability: fsync the downloaded bytes before Close. Unlike imagecache
	// (content-addressed, self-healing on re-fetch), a torn agent binary
	// can't heal itself — these fsyncs (here, in copyFile, and on the parent
	// dir after the rename below) earn the "failure leaves a runnable
	// binary" invariant this package's doc comment claims.
	if err := tmp.Sync(); err != nil {
		tmp.Close()
		return fmt.Errorf("sync downloaded temp: %w", err)
	}
	if err := tmp.Close(); err != nil {
		return err
	}
	if got := hex.EncodeToString(h.Sum(nil)); got != u.SHA256 {
		return fmt.Errorf("sha256 mismatch for %s: got %s want %s", u.URL, got, u.SHA256)
	}

	// The manifest's sha covers the artifact as served, so it is checked above
	// over the whole download — a tarball is unpacked only once its bytes are
	// known to be the ones the server named.
	src := tmp.Name()
	gzipped, err := isGzip(src)
	if err != nil {
		return err
	}
	if gzipped {
		member, err := extractAgent(dir, src)
		if err != nil {
			return fmt.Errorf("extract %s from %s: %w", agentMember, u.URL, err)
		}
		defer os.Remove(member) // no-op after the successful rename
		// The idempotent-resume check above cannot fire for a tarball: the sha
		// it compares describes the archive, not the binary inside it. Spot the
		// already-swapped case here instead, so a retry after a failed exec
		// still leaves <exe>.prev pointing at the genuinely previous version.
		if same, cmpErr := sameFile(exe, member); cmpErr == nil && same {
			return a.execFn()(exe, os.Args, os.Environ())
		}
		src = member
	}

	if err := os.Chmod(src, 0o755); err != nil {
		return err
	}
	// Copy (not rename) the current binary to .prev: a crash between the two
	// steps must leave <exe> present and runnable.
	if err := copyFile(exe, exe+".prev"); err != nil {
		return fmt.Errorf("preserve .prev: %w", err)
	}
	if err := os.Rename(src, exe); err != nil {
		return fmt.Errorf("swap binary: %w", err)
	}
	// Durability: fsync the directory entry itself, not just the file data —
	// on many filesystems a rename's directory-entry update is not durable
	// until the containing directory is fsynced. Best-effort: the swap has
	// already happened either way, so a failure here doesn't unwind it.
	if df, derr := os.Open(dir); derr == nil {
		_ = df.Sync()
		_ = df.Close()
	}
	return a.execFn()(exe, os.Args, os.Environ())
}

// isGzip reports whether the file at path opens with the gzip magic bytes. A
// file too short to carry them is not an archive, which is answer enough.
func isGzip(path string) (bool, error) {
	f, err := os.Open(path)
	if err != nil {
		return false, err
	}
	defer f.Close()
	var magic [2]byte
	if _, err := io.ReadFull(f, magic[:]); err != nil {
		if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
			return false, nil
		}
		return false, err
	}
	return magic[0] == 0x1f && magic[1] == 0x8b, nil
}

// extractAgent writes the archive's agent member to a fresh temp file in dir
// and returns its path. The temp shares the sweep prefix, so an abandoned one
// is litter a later attempt collects rather than litter that stays.
//
// Member names steer nothing: the destination is a path this function chooses,
// and an entry claiming an absolute or dot-dot name is skipped rather than
// mapped onto a basename it does not own. Only regular files qualify, so a
// symlink named eitri-agent is not a way to make the swap read some other file.
func extractAgent(dir, archive string) (string, error) {
	f, err := os.Open(archive)
	if err != nil {
		return "", err
	}
	defer f.Close()
	gz, err := gzip.NewReader(f)
	if err != nil {
		return "", err
	}
	defer gz.Close()
	tr := tar.NewReader(gz)
	for {
		hdr, err := tr.Next()
		if errors.Is(err, io.EOF) {
			return "", fmt.Errorf("no %s member in archive", agentMember)
		}
		if err != nil {
			return "", err
		}
		if hdr.Typeflag != tar.TypeReg || !isAgentMember(hdr.Name) {
			continue
		}
		out, err := os.CreateTemp(dir, ".eitri-agent-upgrade-*")
		if err != nil {
			return "", err
		}
		if _, err := io.Copy(out, tr); err != nil {
			out.Close()
			os.Remove(out.Name())
			return "", err
		}
		// Durability: same power-loss hole as the downloaded temp above.
		if err := out.Sync(); err != nil {
			out.Close()
			os.Remove(out.Name())
			return "", err
		}
		if err := out.Close(); err != nil {
			os.Remove(out.Name())
			return "", err
		}
		return out.Name(), nil
	}
}

// isAgentMember reports whether a tar entry name is the bundle's agent binary.
func isAgentMember(name string) bool {
	if strings.HasPrefix(name, "/") {
		return false
	}
	for _, element := range strings.Split(name, "/") {
		if element == ".." {
			return false
		}
	}
	return path.Base(name) == agentMember
}

// sameFile reports whether two paths hold identical bytes.
func sameFile(a, b string) (bool, error) {
	sumA, err := sha256File(a)
	if err != nil {
		return false, err
	}
	sumB, err := sha256File(b)
	if err != nil {
		return false, err
	}
	return sumA == sumB, nil
}

// sha256File hashes the file at path, hex-encoded.
func sha256File(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}

// copyFile copies src to dst (0755), truncating dst.
func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o755)
	if err != nil {
		return err
	}
	if _, err := io.Copy(out, in); err != nil {
		out.Close()
		return err
	}
	// Durability: see the fsync note on the downloaded temp above — the same
	// power-loss hole applies to the .prev recovery copy.
	if err := out.Sync(); err != nil {
		out.Close()
		return err
	}
	return out.Close()
}