spb/internetpressure

1package update23import (4	"context"5	"crypto/sha256"6	"encoding/hex"7	"errors"8	"io"9	"log/slog"10	"os"11	"path/filepath"12	"testing"1314	"internetpressure.io/probe-agent/internal/protocol"15)1617type fakeFetcher struct {18	latest    protocol.AgentLatest19	asset     []byte20	downloads int21}2223func (f *fakeFetcher) AgentLatest(context.Context) (*protocol.AgentLatest, error) {24	return &f.latest, nil25}26func (f *fakeFetcher) Download(_ context.Context, _ string, w io.Writer, _ int64) (int64, error) {27	f.downloads++28	n, err := w.Write(f.asset)29	return int64(n), err30}3132func TestUpdateReplacesBinary(t *testing.T) {33	dir := t.TempDir()34	exe := filepath.Join(dir, "bin", "ip-probe")35	os.MkdirAll(filepath.Dir(exe), 0o755)36	os.WriteFile(exe, []byte("old"), 0o755)3738	newBin := []byte("#!/bin/sh\necho new\n")39	sum := sha256.Sum256(newBin)40	f := &fakeFetcher{asset: newBin, latest: protocol.AgentLatest{Version: "0.2.0",41		Assets: map[string]protocol.Asset{Platform(): {URL: "https://x/ip-probe", SHA256: hex.EncodeToString(sum[:])}}}}42	u := &Updater{Fetcher: f, Version: "0.1.0", DataDir: dir, ExePath: exe, Log: slog.New(slog.DiscardHandler)}4344	updated, err := u.Check(context.Background())45	if err != nil || !updated {46		t.Fatalf("updated=%v err=%v", updated, err)47	}48	got, _ := os.ReadFile(exe)49	if string(got) != string(newBin) {50		t.Fatalf("binary not replaced: %q", got)51	}52	if fi, _ := os.Stat(exe); fi.Mode().Perm() != 0o755 {53		t.Fatalf("mode %v", fi.Mode())54	}55	if _, err := os.Stat(filepath.Join(dir, tmpFile)); !errors.Is(err, os.ErrNotExist) {56		t.Fatal("tmp file left behind")57	}58	// Same version offered again (e.g. the new binary crashes and the old one is restored) → not re-attempted.59	os.WriteFile(exe, []byte("old"), 0o755)60	updated, err = u.Check(context.Background())61	if err != nil || updated || f.downloads != 1 {62		t.Fatalf("loop guard failed: updated=%v err=%v downloads=%d", updated, err, f.downloads)63	}64}6566func TestUpdateRejectsBadHashAndSameVersion(t *testing.T) {67	dir := t.TempDir()68	exe := filepath.Join(dir, "ip-probe")69	os.WriteFile(exe, []byte("old"), 0o755)70	f := &fakeFetcher{asset: []byte("evil"), latest: protocol.AgentLatest{Version: "0.3.0",71		Assets: map[string]protocol.Asset{Platform(): {URL: "u", SHA256: "00"}}}}72	u := &Updater{Fetcher: f, Version: "0.1.0", DataDir: dir, ExePath: exe, Log: slog.New(slog.DiscardHandler)}73	updated, err := u.Check(context.Background())74	if updated || err == nil {75		t.Fatalf("bad hash accepted: updated=%v err=%v", updated, err)76	}77	if got, _ := os.ReadFile(exe); string(got) != "old" {78		t.Fatal("binary modified despite bad hash")79	}80	// Same version → nothing to do, no download.81	f.latest.Version = "0.1.0"82	f.downloads = 083	if updated, err := u.Check(context.Background()); updated || err != nil || f.downloads != 0 {84		t.Fatalf("same version: updated=%v err=%v downloads=%d", updated, err, f.downloads)85	}86	// No asset for this platform → skip quietly.87	f.latest = protocol.AgentLatest{Version: "0.4.0", Assets: map[string]protocol.Asset{"plan9-mips": {URL: "u", SHA256: "00"}}}88	if updated, err := u.Check(context.Background()); updated || err != nil {89		t.Fatalf("missing platform: updated=%v err=%v", updated, err)90	}91}92