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

package agent

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/mhsanaei/3x-ui/v3/internal/config"
	"github.com/mhsanaei/3x-ui/v3/internal/logger"
	"github.com/mhsanaei/3x-ui/v3/internal/selfupdate"
)

// One-click agent update (item 20). The master records an intent on the
// agent row; the intent rides the existing config poll (no new channel).
// The agent downloads the release binary for its architecture from the
// master's release base, verifies it against the published .sha256
// sidecar, swaps it into place keeping the old binary as .bak, and
// re-execs itself. Any failure before the swap leaves the old binary
// running untouched; the outcome is reported on heartbeats, and the new
// A-series version appearing at the master is the success signal.

// updateIntent is the master's update request as delivered on the config
// response.
type updateIntent struct {
	RequestedAt int64  `json:"requested_at"`
	Target      string `json:"target"`   // A-form display target, e.g. "A1.0.1"
	Tag         string `json:"tag"`      // release tag, e.g. "V1.0.1"
	BaseURL     string `json:"base_url"` // release download base (trailing slash)
}

// agentUpdateStatus mirrors the master's heartbeat updateStatus object.
type agentUpdateStatus struct {
	RequestedAt int64  `json:"requestedAt"`
	State       string `json:"state"` // downloading | swapped | failed | current
	Error       string `json:"error,omitempty"`
}

// updateRetryBackoff suppresses immediate re-downloads after a failure for
// the same intent (the master clears failed intents on report, so this
// only guards a master that keeps re-offering the same requested_at).
const updateRetryBackoff = 10 * time.Minute

// execSelf replaces the current process image with the binary at path.
// A package variable so tests can stub the restart.
var execSelf = func(path string) error {
	return syscall.Exec(path, append([]string{path}, os.Args[1:]...), os.Environ())
}

func (a *Agent) setUpdateReport(st agentUpdateStatus) {
	a.updateMu.Lock()
	cp := st
	a.updateReport = &cp
	a.updateMu.Unlock()
}

func (a *Agent) currentUpdateReport() *agentUpdateStatus {
	a.updateMu.Lock()
	defer a.updateMu.Unlock()
	if a.updateReport == nil {
		return nil
	}
	cp := *a.updateReport
	return &cp
}

// maybeHandleUpdate is called on every config poll with the intent the
// master attached (nil when none). Terminal reports are dropped once the
// master withdraws the intent, so heartbeats stop repeating them.
func (a *Agent) maybeHandleUpdate(intent *updateIntent) {
	if intent == nil {
		a.updateMu.Lock()
		if a.updateReport != nil &&
			(a.updateReport.State == "failed" || a.updateReport.State == "current") {
			a.updateReport = nil
		}
		a.updateMu.Unlock()
		return
	}
	a.updateMu.Lock()
	if a.updateRunning || intent.RequestedAt == a.updateHandled ||
		(intent.RequestedAt == a.updateFailedReq && time.Since(a.updateFailedAt) < updateRetryBackoff) {
		a.updateMu.Unlock()
		return
	}
	a.updateRunning = true
	a.updateMu.Unlock()
	go a.runUpdate(*intent)
}

// runUpdate executes one intent end to end. It runs synchronously in its
// own goroutine (spawned by maybeHandleUpdate) — tests call it directly.
func (a *Agent) runUpdate(intent updateIntent) {
	defer func() {
		a.updateMu.Lock()
		a.updateRunning = false
		a.updateMu.Unlock()
	}()
	a.setUpdateReport(agentUpdateStatus{RequestedAt: intent.RequestedAt, State: "downloading"})

	fail := func(err error) {
		a.updateMu.Lock()
		a.updateHandled = intent.RequestedAt
		a.updateFailedReq = intent.RequestedAt
		a.updateFailedAt = time.Now()
		a.updateMu.Unlock()
		a.setUpdateReport(agentUpdateStatus{RequestedAt: intent.RequestedAt, State: "failed", Error: err.Error()})
		logger.Warningf("agent update to %s failed: %v (old version keeps running)", intent.Target, err)
	}

	// Already at (or beyond) the target — e.g. the intent is re-delivered
	// after the swap-restart before the master closes it out. Report
	// "current" so the master clears the intent without a download.
	if cmp, ok := config.CompareDisplayVersion(config.GetAgentDisplayVersion(), intent.Target); ok && cmp >= 0 {
		a.updateMu.Lock()
		a.updateHandled = intent.RequestedAt
		a.updateMu.Unlock()
		a.setUpdateReport(agentUpdateStatus{RequestedAt: intent.RequestedAt, State: "current"})
		return
	}

	exePath, err := a.selfExePath()
	if err != nil {
		fail(fmt.Errorf("locate own binary: %w", err))
		return
	}
	asset, err := selfupdate.AssetName(runtime.GOARCH)
	if err != nil {
		fail(err)
		return
	}
	base := intent.BaseURL
	if base == "" {
		base = selfupdate.BaseURL()
	}
	base = strings.TrimSuffix(base, "/") + "/" + intent.Tag + "/"

	tmp := filepath.Join(filepath.Dir(exePath), fmt.Sprintf(".x-ui-update-%d", intent.RequestedAt))
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()
	var transport http.RoundTripper
	if a.httpClient != nil {
		transport = a.httpClient.Transport
	}
	client := &http.Client{Transport: transport, Timeout: 10 * time.Minute}
	if err := selfupdate.DownloadVerified(ctx, client, base+asset, base+asset+".sha256", tmp); err != nil {
		_ = os.Remove(tmp)
		fail(err)
		return
	}
	if _, err := selfupdate.SwapExecutable(exePath, tmp); err != nil {
		_ = os.Remove(tmp)
		fail(err)
		return
	}

	a.updateMu.Lock()
	a.updateHandled = intent.RequestedAt
	a.updateMu.Unlock()
	a.setUpdateReport(agentUpdateStatus{RequestedAt: intent.RequestedAt, State: "swapped"})
	logger.Infof("agent update to %s downloaded, verified and swapped in; restarting", intent.Target)
	// Best effort: flush the "swapped" report before the image is replaced
	// (success is ultimately proven by the new version's telemetry).
	_ = a.sendHeartbeat()
	a.stopXray()
	if err := execSelf(exePath); err != nil {
		// The exec failed before replacing the image: roll the old binary
		// back into place and keep running it.
		if rbErr := selfupdate.RollbackExecutable(exePath); rbErr != nil {
			logger.Warningf("agent update rollback failed: %v", rbErr)
		}
		fail(fmt.Errorf("restart into new binary: %w", err))
	}
}

// selfExePath resolves the agent's own binary path (os.Executable unless a
// test overrode it).
func (a *Agent) selfExePath() (string, error) {
	if a.exePathOverride != "" {
		return a.exePathOverride, nil
	}
	return selfupdate.ExecutablePath()
}