3x-ui-agent/internal/selfupdate/selfupdate_test.go

package selfupdate

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestParseSHA256Sidecar(t *testing.T) {
	sum := sha256.Sum256([]byte("payload"))
	hexSum := hex.EncodeToString(sum[:])
	cases := []struct {
		in   string
		want string
		ok   bool
	}{
		{hexSum + "  x-ui-linux-amd64\n", hexSum, true},
		{hexSum + "\n", hexSum, true},
		{strings.ToUpper(hexSum) + " *x-ui-linux-amd64", hexSum, true},
		{"", "", false},
		{"not-a-digest  file", "", false},
		{hexSum[:40] + "  file", "", false},
	}
	for _, tc := range cases {
		got, ok := ParseSHA256Sidecar([]byte(tc.in))
		if got != tc.want || ok != tc.ok {
			t.Fatalf("ParseSHA256Sidecar(%q) = %q, %v; want %q, %v", tc.in, got, ok, tc.want, tc.ok)
		}
	}
}

func TestAssetName(t *testing.T) {
	if name, err := AssetName("amd64"); err != nil || name != "x-ui-linux-amd64" {
		t.Fatalf("amd64: %q, %v", name, err)
	}
	if name, err := AssetName("arm64"); err != nil || name != "x-ui-linux-arm64" {
		t.Fatalf("arm64: %q, %v", name, err)
	}
	if _, err := AssetName("riscv64"); err == nil {
		t.Fatal("unknown arch must error")
	}
}

func TestSwapExecutableKeepsBackup(t *testing.T) {
	dir := t.TempDir()
	exe := filepath.Join(dir, "xui-agent")
	if err := os.WriteFile(exe, []byte("old-binary"), 0o755); err != nil {
		t.Fatal(err)
	}
	next := filepath.Join(dir, ".x-ui-update-1")
	if err := os.WriteFile(next, []byte("new-binary"), 0o644); err != nil {
		t.Fatal(err)
	}
	bak, err := SwapExecutable(exe, next)
	if err != nil {
		t.Fatalf("SwapExecutable: %v", err)
	}
	if bak != exe+".bak" {
		t.Fatalf("backup path = %q", bak)
	}
	got, _ := os.ReadFile(exe)
	if string(got) != "new-binary" {
		t.Fatalf("exe content = %q, want new-binary", got)
	}
	old, _ := os.ReadFile(bak)
	if string(old) != "old-binary" {
		t.Fatalf("bak content = %q, want old-binary", old)
	}
	info, _ := os.Stat(exe)
	if info.Mode().Perm()&0o100 == 0 {
		t.Fatal("new binary is not executable")
	}
	// Rollback restores the old binary.
	if err := RollbackExecutable(exe); err != nil {
		t.Fatalf("RollbackExecutable: %v", err)
	}
	got, _ = os.ReadFile(exe)
	if string(got) != "old-binary" {
		t.Fatalf("after rollback exe = %q", got)
	}
}

func TestDownloadVerified(t *testing.T) {
	payload := []byte("fake-release-binary")
	sum := sha256.Sum256(payload)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case strings.HasSuffix(r.URL.Path, ".sha256"):
			_, _ = w.Write([]byte(hex.EncodeToString(sum[:]) + "  x-ui-linux-amd64\n"))
		default:
			_, _ = w.Write(payload)
		}
	}))
	defer srv.Close()

	dest := filepath.Join(t.TempDir(), "dl")
	if err := DownloadVerified(context.Background(), srv.Client(), srv.URL+"/x-ui-linux-amd64", srv.URL+"/x-ui-linux-amd64.sha256", dest); err != nil {
		t.Fatalf("DownloadVerified: %v", err)
	}
	got, _ := os.ReadFile(dest)
	if string(got) != string(payload) {
		t.Fatalf("downloaded content = %q", got)
	}
}

func TestDownloadVerifiedRejectsBadChecksum(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, ".sha256") {
			_, _ = w.Write([]byte(strings.Repeat("0", 64) + "  x-ui-linux-amd64\n"))
			return
		}
		_, _ = w.Write([]byte("tampered"))
	}))
	defer srv.Close()

	dest := filepath.Join(t.TempDir(), "dl")
	err := DownloadVerified(context.Background(), srv.Client(), srv.URL+"/a", srv.URL+"/a.sha256", dest)
	if err == nil || !strings.Contains(err.Error(), "sha256 mismatch") {
		t.Fatalf("want sha256 mismatch error, got %v", err)
	}
	if _, statErr := os.Stat(dest); !os.IsNotExist(statErr) {
		t.Fatal("partial download must be removed on mismatch")
	}
}