3x-ui-agent/internal/agent/update_test.go

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" }