package agent import ( "crypto/sha256" "encoding/hex" "net/http" "net/http/httptest" "os" "path/filepath" "runtime" "strings" "sync/atomic" "testing" "github.com/mhsanaei/3x-ui/v3/internal/selfupdate" ) // Item 20: the agent consumes the master's update intent, downloads the // release asset for its architecture, verifies the sha256 sidecar, swaps // the binary keeping a .bak, and re-execs. Failures must leave the old // binary running and produce a "failed" report. func fakeReleaseServer(t *testing.T, payload []byte, goodSum bool, hits *atomic.Int32) *httptest.Server { t.Helper() asset, err := selfupdate.AssetName(runtime.GOARCH) if err != nil { t.Fatal(err) } sum := sha256.Sum256(payload) if !goodSum { sum = sha256.Sum256([]byte("something else")) } return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if hits != nil { hits.Add(1) } switch { case strings.HasSuffix(r.URL.Path, asset+".sha256"): _, _ = w.Write([]byte(hex.EncodeToString(sum[:]) + " " + asset + "\n")) case strings.HasSuffix(r.URL.Path, asset): _, _ = w.Write(payload) default: http.NotFound(w, r) } })) } func newUpdateTestAgent(t *testing.T, masterURL string) (*Agent, string) { t.Helper() dir := t.TempDir() exe := filepath.Join(dir, "xui-agent") if err := os.WriteFile(exe, []byte("old-binary"), 0o755); err != nil { t.Fatal(err) } a := &Agent{ config: &Config{MasterURL: masterURL, Token: "tok", NodeName: "n"}, httpClient: http.DefaultClient, exePathOverride: exe, } return a, exe } func stubExec(t *testing.T, err error) *atomic.Int32 { t.Helper() calls := &atomic.Int32{} orig := execSelf execSelf = func(string) error { calls.Add(1) return err } t.Cleanup(func() { execSelf = orig }) return calls } func TestAgentUpdateSwapsAndRestarts(t *testing.T) { srv := fakeReleaseServer(t, []byte("new-binary"), true, nil) defer srv.Close() a, exe := newUpdateTestAgent(t, srv.URL) execCalls := stubExec(t, nil) a.runUpdate(updateIntent{RequestedAt: 42, Target: "A2.0.1", Tag: "V2.0.1", BaseURL: srv.URL + "/"}) got, _ := os.ReadFile(exe) if string(got) != "new-binary" { t.Fatalf("exe after update = %q, want new-binary", got) } bak, _ := os.ReadFile(exe + ".bak") if string(bak) != "old-binary" { t.Fatalf(".bak = %q, want old-binary", bak) } if execCalls.Load() != 1 { t.Fatalf("execSelf calls = %d, want 1", execCalls.Load()) } if rep := a.currentUpdateReport(); rep == nil || rep.State != "swapped" { t.Fatalf("report = %+v, want swapped", rep) } } func TestAgentUpdateBadChecksumKeepsOldBinary(t *testing.T) { srv := fakeReleaseServer(t, []byte("new-binary"), false, nil) defer srv.Close() a, exe := newUpdateTestAgent(t, srv.URL) execCalls := stubExec(t, nil) a.runUpdate(updateIntent{RequestedAt: 43, Target: "A2.0.1", Tag: "V2.0.1", BaseURL: srv.URL + "/"}) got, _ := os.ReadFile(exe) if string(got) != "old-binary" { t.Fatalf("exe after failed update = %q, want old-binary", got) } if _, err := os.Stat(exe + ".bak"); !os.IsNotExist(err) { t.Fatal("no .bak may be created when verification fails") } if execCalls.Load() != 0 { t.Fatal("execSelf must not run when verification fails") } rep := a.currentUpdateReport() if rep == nil || rep.State != "failed" || !strings.Contains(rep.Error, "sha256 mismatch") { t.Fatalf("report = %+v, want failed with sha256 mismatch", rep) } } func TestAgentUpdateAlreadyCurrentSkipsDownload(t *testing.T) { var hits atomic.Int32 srv := fakeReleaseServer(t, []byte("new-binary"), true, &hits) defer srv.Close() a, exe := newUpdateTestAgent(t, srv.URL) stubExec(t, nil) // The test binary's display version defaults to A2.0.0. a.runUpdate(updateIntent{RequestedAt: 44, Target: "A2.0.0", Tag: "V2.0.0", BaseURL: srv.URL + "/"}) if hits.Load() != 0 { t.Fatalf("release server hit %d times for an already-current agent", hits.Load()) } got, _ := os.ReadFile(exe) if string(got) != "old-binary" { t.Fatalf("exe = %q, want untouched old-binary", got) } if rep := a.currentUpdateReport(); rep == nil || rep.State != "current" { t.Fatalf("report = %+v, want current", rep) } } func TestAgentUpdateExecFailureRollsBack(t *testing.T) { srv := fakeReleaseServer(t, []byte("new-binary"), true, nil) defer srv.Close() a, exe := newUpdateTestAgent(t, srv.URL) stubExec(t, errExecBoom) a.runUpdate(updateIntent{RequestedAt: 45, Target: "A2.0.1", Tag: "V2.0.1", BaseURL: srv.URL + "/"}) got, _ := os.ReadFile(exe) if string(got) != "old-binary" { t.Fatalf("exe after exec failure = %q, want rolled-back old-binary", got) } if rep := a.currentUpdateReport(); rep == nil || rep.State != "failed" { t.Fatalf("report = %+v, want failed", rep) } } func TestAgentUpdateSkippedOnUnsupportedOS(t *testing.T) { orig := selfUpdateSupported selfUpdateSupported = false t.Cleanup(func() { selfUpdateSupported = orig }) var hits atomic.Int32 srv := fakeReleaseServer(t, []byte("new-binary"), true, &hits) defer srv.Close() a, exe := newUpdateTestAgent(t, srv.URL) intent := &updateIntent{RequestedAt: 46, Target: "A2.0.1", Tag: "V2.0.1", BaseURL: srv.URL + "/"} a.maybeHandleUpdate(intent) // A repeated delivery of the same intent is also a no-op. a.maybeHandleUpdate(intent) if hits.Load() != 0 { t.Fatalf("release server hit %d times on an unsupported OS, want 0", hits.Load()) } got, _ := os.ReadFile(exe) if string(got) != "old-binary" { t.Fatalf("exe = %q, want untouched old-binary", got) } a.updateMu.Lock() handled, running := a.updateHandled, a.updateRunning a.updateMu.Unlock() if handled != 46 || running { t.Fatalf("handled=%d running=%v, want handled=46 running=false", handled, running) } if rep := a.currentUpdateReport(); rep != nil { t.Fatalf("report = %+v, want none for a skipped intent", rep) } } var errExecBoom = &execBoomError{} type execBoomError struct{} func (*execBoomError) Error() string { return "exec boom" }