first commit

agent/internal/upgrade/upgrade.go

@@ -0,0 +1,400 @@
+package upgrade
+
+import (
+	"bytes"
+	"compress/gzip"
+	"context"
+	"crypto/rand"
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"io"
+	"math/big"
+	"net/http"
+	"net/url"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"regexp"
+	"strings"
+	"time"
+
+	"incudal-agent/internal/config"
+	"incudal-agent/internal/panel"
+)
+
+const (
+	defaultServiceName      = "incudal-agent"
+	defaultMaxDownloadBytes = 64 * 1024 * 1024
+)
+
+var ErrUpgradeInProgress = errors.New("agent upgrade already in progress")
+var systemdServiceNamePattern = regexp.MustCompile(`^[A-Za-z0-9_.@-]+$`)
+
+type RestartFunc func(ctx context.Context, serviceName string) error
+
+type Runner struct {
+	BinaryPath       string
+	BackupPath       string
+	LockPath         string
+	ServiceName      string
+	AllowedBaseURL   string
+	HTTPClient       *http.Client
+	Restart          RestartFunc
+	MaxDownloadBytes int64
+}
+
+func DefaultRunner(cfg config.Config) *Runner {
+	binaryPath, err := os.Executable()
+	if err != nil || binaryPath == "" {
+		binaryPath = "/usr/local/bin/incudal-agent"
+	}
+
+	return &Runner{
+		BinaryPath:       binaryPath,
+		BackupPath:       binaryPath + ".bak",
+		LockPath:         defaultLockPath(),
+		ServiceName:      defaultServiceName,
+		AllowedBaseURL:   cfg.PanelURL,
+		HTTPClient:       &http.Client{Timeout: cfg.RequestTimeout},
+		Restart:          restartSystemdService,
+		MaxDownloadBytes: defaultMaxDownloadBytes,
+	}
+}
+
+func RandomJitter(max time.Duration) time.Duration {
+	if max <= 0 {
+		return 0
+	}
+	limit := big.NewInt(int64(max))
+	value, err := rand.Int(rand.Reader, limit)
+	if err != nil {
+		return 0
+	}
+	return time.Duration(value.Int64())
+}
+
+func (runner *Runner) Apply(ctx context.Context, instruction panel.UpgradeInstruction, currentVersion string) error {
+	if !instruction.Available {
+		return nil
+	}
+	if instruction.Version == "" {
+		return errors.New("upgrade version is required")
+	}
+	if instruction.Version == currentVersion {
+		return nil
+	}
+	if instruction.URL == "" {
+		return errors.New("upgrade URL is required")
+	}
+	if instruction.SHA256 == "" {
+		return errors.New("upgrade sha256 is required")
+	}
+	if err := runner.validateUpgradeURL(instruction.URL); err != nil {
+		return err
+	}
+
+	unlock, err := acquireLock(runner.lockPath())
+	if err != nil {
+		return err
+	}
+	defer unlock()
+
+	packageBytes, err := runner.download(ctx, instruction.URL)
+	if err != nil {
+		return err
+	}
+	if err := verifySHA256(packageBytes, instruction.SHA256); err != nil {
+		return err
+	}
+
+	binaryBytes := packageBytes
+	if instruction.Gzip {
+		binaryBytes, err = gunzip(packageBytes)
+		if err != nil {
+			return err
+		}
+	}
+
+	tempPath, err := runner.writeTempBinary(binaryBytes)
+	if err != nil {
+		return err
+	}
+	if err := runner.replaceBinary(tempPath); err != nil {
+		_ = os.Remove(tempPath)
+		return err
+	}
+
+	if err := runner.restart(ctx); err != nil {
+		return fmt.Errorf("restart agent after upgrade: %w", err)
+	}
+	return nil
+}
+
+func (runner *Runner) validateUpgradeURL(rawURL string) error {
+	upgradeURL, err := url.Parse(rawURL)
+	if err != nil || upgradeURL.Scheme == "" || upgradeURL.Host == "" {
+		return fmt.Errorf("upgrade URL is invalid: %s", rawURL)
+	}
+	if upgradeURL.Scheme != "http" && upgradeURL.Scheme != "https" {
+		return fmt.Errorf("upgrade URL scheme is not allowed: %s", upgradeURL.Scheme)
+	}
+
+	baseURL, err := url.Parse(strings.TrimRight(runner.AllowedBaseURL, "/"))
+	if err != nil || baseURL.Scheme == "" || baseURL.Host == "" {
+		return fmt.Errorf("panel URL is invalid: %s", runner.AllowedBaseURL)
+	}
+	if !strings.EqualFold(upgradeURL.Scheme, baseURL.Scheme) || !strings.EqualFold(upgradeURL.Host, baseURL.Host) {
+		return errors.New("upgrade URL is outside panel origin")
+	}
+	return nil
+}
+
+func (runner *Runner) download(ctx context.Context, rawURL string) ([]byte, error) {
+	client := runner.HTTPClient
+	if client == nil {
+		client = &http.Client{Timeout: 30 * time.Second}
+	}
+
+	request, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
+	if err != nil {
+		return nil, err
+	}
+	response, err := client.Do(request)
+	if err != nil {
+		return nil, err
+	}
+	defer response.Body.Close()
+
+	if response.StatusCode < 200 || response.StatusCode >= 300 {
+		body, _ := io.ReadAll(io.LimitReader(response.Body, 4096))
+		return nil, fmt.Errorf("download upgrade failed: status=%d body=%s", response.StatusCode, string(body))
+	}
+
+	limit := runner.MaxDownloadBytes
+	if limit <= 0 {
+		limit = defaultMaxDownloadBytes
+	}
+	body, err := io.ReadAll(io.LimitReader(response.Body, limit+1))
+	if err != nil {
+		return nil, err
+	}
+	if int64(len(body)) > limit {
+		return nil, fmt.Errorf("upgrade package exceeds %d bytes", limit)
+	}
+	return body, nil
+}
+
+func (runner *Runner) writeTempBinary(binaryBytes []byte) (string, error) {
+	binaryPath := runner.binaryPath()
+	tempFile, err := os.CreateTemp(filepath.Dir(binaryPath), ".incudal-agent-upgrade-*")
+	if err != nil {
+		return "", err
+	}
+	tempPath := tempFile.Name()
+	defer tempFile.Close()
+
+	if _, err := tempFile.Write(binaryBytes); err != nil {
+		_ = os.Remove(tempPath)
+		return "", err
+	}
+	if err := tempFile.Chmod(0755); err != nil {
+		_ = os.Remove(tempPath)
+		return "", err
+	}
+	return tempPath, nil
+}
+
+func (runner *Runner) replaceBinary(tempPath string) error {
+	binaryPath := runner.binaryPath()
+	backupPath := runner.backupPath()
+
+	if _, err := os.Stat(binaryPath); err == nil {
+		_ = os.Remove(backupPath)
+		if err := copyFile(binaryPath, backupPath); err != nil {
+			return fmt.Errorf("backup current agent: %w", err)
+		}
+	}
+
+	if err := os.Rename(tempPath, binaryPath); err != nil {
+		return fmt.Errorf("replace agent binary: %w", err)
+	}
+	return nil
+}
+
+func (runner *Runner) rollback() error {
+	backupPath := runner.backupPath()
+	if _, err := os.Stat(backupPath); err != nil {
+		return err
+	}
+	return os.Rename(backupPath, runner.binaryPath())
+}
+
+func (runner *Runner) restart(ctx context.Context) error {
+	if runner.Restart == nil {
+		return nil
+	}
+	serviceName := runner.ServiceName
+	if serviceName == "" {
+		serviceName = defaultServiceName
+	}
+	return runner.Restart(ctx, serviceName)
+}
+
+func (runner *Runner) binaryPath() string {
+	if runner.BinaryPath != "" {
+		return runner.BinaryPath
+	}
+	return "/usr/local/bin/incudal-agent"
+}
+
+func (runner *Runner) backupPath() string {
+	if runner.BackupPath != "" {
+		return runner.BackupPath
+	}
+	return runner.binaryPath() + ".bak"
+}
+
+func (runner *Runner) lockPath() string {
+	if runner.LockPath != "" {
+		return runner.LockPath
+	}
+	return defaultLockPath()
+}
+
+func defaultLockPath() string {
+	if info, err := os.Stat("/run"); err == nil && info.IsDir() {
+		return "/run/incudal-agent-upgrade.lock"
+	}
+	return filepath.Join(os.TempDir(), "incudal-agent-upgrade.lock")
+}
+
+func acquireLock(lockPath string) (func(), error) {
+	file, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0600)
+	if err != nil {
+		if errors.Is(err, os.ErrExist) {
+			return nil, ErrUpgradeInProgress
+		}
+		return nil, err
+	}
+	_, _ = fmt.Fprintf(file, "%d\n", os.Getpid())
+	_ = file.Close()
+
+	return func() {
+		_ = os.Remove(lockPath)
+	}, nil
+}
+
+func verifySHA256(payload []byte, expected string) error {
+	sum := sha256.Sum256(payload)
+	actual := hex.EncodeToString(sum[:])
+	if !strings.EqualFold(actual, expected) {
+		return fmt.Errorf("upgrade sha256 mismatch: expected=%s actual=%s", expected, actual)
+	}
+	return nil
+}
+
+func gunzip(payload []byte) ([]byte, error) {
+	reader, err := gzip.NewReader(bytes.NewReader(payload))
+	if err != nil {
+		return nil, err
+	}
+	defer reader.Close()
+	return io.ReadAll(reader)
+}
+
+func copyFile(source string, target string) error {
+	sourceFile, err := os.Open(source)
+	if err != nil {
+		return err
+	}
+	defer sourceFile.Close()
+
+	info, err := sourceFile.Stat()
+	if err != nil {
+		return err
+	}
+	mode := info.Mode().Perm()
+	if mode == 0 {
+		mode = 0755
+	}
+
+	targetFile, err := os.OpenFile(target, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, mode)
+	if err != nil {
+		return err
+	}
+	defer targetFile.Close()
+
+	if _, err := io.Copy(targetFile, sourceFile); err != nil {
+		return err
+	}
+	return targetFile.Chmod(mode)
+}
+
+func restartSystemdService(ctx context.Context, serviceName string) error {
+	if !systemdServiceNamePattern.MatchString(serviceName) {
+		return fmt.Errorf("invalid systemd service name: %s", serviceName)
+	}
+
+	if err := scheduleSystemdRestart(ctx, serviceName); err == nil {
+		return nil
+	}
+
+	systemctlPath, err := exec.LookPath("systemctl")
+	if err != nil {
+		return err
+	}
+
+	// 不等待 systemctl 完成。Agent 正在重启自身，等待子进程会在服务停止时
+	// 收到 SIGTERM，旧逻辑会误判失败并回滚已替换的新二进制。
+	command := exec.CommandContext(ctx, systemctlPath, "restart", serviceName)
+	return command.Start()
+}
+
+func scheduleSystemdRestart(ctx context.Context, serviceName string) error {
+	systemdRunPath, err := exec.LookPath("systemd-run")
+	if err != nil {
+		return err
+	}
+	systemctlPath, err := exec.LookPath("systemctl")
+	if err != nil {
+		return err
+	}
+
+	unitName := fmt.Sprintf("incudal-agent-restart-%d", os.Getpid())
+	args := []string{
+		"--unit", unitName,
+		"--description", "Restart Incudal Agent after self-upgrade",
+		"--on-active=2s",
+		"--collect",
+		systemctlPath, "restart", serviceName,
+	}
+
+	command := exec.CommandContext(ctx, systemdRunPath, args...)
+	output, err := command.CombinedOutput()
+	if err == nil {
+		return nil
+	}
+
+	// 老版本 systemd 可能不支持 --collect，降级重试一次。
+	if strings.Contains(string(output), "unrecognized option '--collect'") ||
+		strings.Contains(string(output), "Unknown option --collect") {
+		args = []string{
+			"--unit", unitName,
+			"--description", "Restart Incudal Agent after self-upgrade",
+			"--on-active=2s",
+			systemctlPath, "restart", serviceName,
+		}
+		command = exec.CommandContext(ctx, systemdRunPath, args...)
+		output, err = command.CombinedOutput()
+	}
+	if err != nil {
+		trimmedOutput := strings.TrimSpace(string(output))
+		if trimmedOutput == "" {
+			return err
+		}
+		return fmt.Errorf("%w: %s", err, trimmedOutput)
+	}
+	return nil
+}

agent/internal/upgrade/upgrade_test.go

@@ -0,0 +1,230 @@
+package upgrade
+
+import (
+	"bytes"
+	"compress/gzip"
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"incudal-agent/internal/panel"
+)
+
+func TestApplyUpgradeReplacesBinaryAndRestarts(t *testing.T) {
+	tempDir := t.TempDir()
+	binaryPath := filepath.Join(tempDir, "incudal-agent")
+	if err := os.WriteFile(binaryPath, []byte("old-binary"), 0755); err != nil {
+		t.Fatalf("write current binary: %v", err)
+	}
+
+	nextBinary := []byte("new-binary")
+	packageBytes := gzipBytes(t, nextBinary)
+	sha := sha256Hex(packageBytes)
+	server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
+		response.WriteHeader(http.StatusOK)
+		_, _ = response.Write(packageBytes)
+	}))
+	defer server.Close()
+
+	restarted := false
+	runner := Runner{
+		BinaryPath:       binaryPath,
+		BackupPath:       binaryPath + ".bak",
+		LockPath:         filepath.Join(tempDir, "upgrade.lock"),
+		ServiceName:      "incudal-agent",
+		AllowedBaseURL:   server.URL,
+		HTTPClient:       server.Client(),
+		MaxDownloadBytes: 1024 * 1024,
+		Restart: func(_ context.Context, serviceName string) error {
+			if serviceName != "incudal-agent" {
+				t.Fatalf("unexpected service name: %s", serviceName)
+			}
+			restarted = true
+			return nil
+		},
+	}
+
+	err := runner.Apply(context.Background(), panel.UpgradeInstruction{
+		Available: true,
+		Version:   "v2",
+		URL:       server.URL + "/incudal-agent-linux-amd64.gz",
+		SHA256:    sha,
+		Gzip:      true,
+	}, "v1")
+	if err != nil {
+		t.Fatalf("apply upgrade: %v", err)
+	}
+	if !restarted {
+		t.Fatalf("restart was not called")
+	}
+
+	actual, err := os.ReadFile(binaryPath)
+	if err != nil {
+		t.Fatalf("read replaced binary: %v", err)
+	}
+	if string(actual) != string(nextBinary) {
+		t.Fatalf("binary mismatch: %q", string(actual))
+	}
+
+	backup, err := os.ReadFile(binaryPath + ".bak")
+	if err != nil {
+		t.Fatalf("read backup binary: %v", err)
+	}
+	if string(backup) != "old-binary" {
+		t.Fatalf("backup mismatch: %q", string(backup))
+	}
+}
+
+func TestApplyUpgradeRejectsBadSHA(t *testing.T) {
+	tempDir := t.TempDir()
+	binaryPath := filepath.Join(tempDir, "incudal-agent")
+	if err := os.WriteFile(binaryPath, []byte("old-binary"), 0755); err != nil {
+		t.Fatalf("write current binary: %v", err)
+	}
+
+	server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
+		response.WriteHeader(http.StatusOK)
+		_, _ = response.Write([]byte("payload"))
+	}))
+	defer server.Close()
+
+	restarted := false
+	runner := Runner{
+		BinaryPath:       binaryPath,
+		BackupPath:       binaryPath + ".bak",
+		LockPath:         filepath.Join(tempDir, "upgrade.lock"),
+		AllowedBaseURL:   server.URL,
+		HTTPClient:       server.Client(),
+		MaxDownloadBytes: 1024 * 1024,
+		Restart: func(context.Context, string) error {
+			restarted = true
+			return nil
+		},
+	}
+
+	err := runner.Apply(context.Background(), panel.UpgradeInstruction{
+		Available: true,
+		Version:   "v2",
+		URL:       server.URL + "/incudal-agent-linux-amd64.gz",
+		SHA256:    "0000000000000000000000000000000000000000000000000000000000000000",
+		Gzip:      true,
+	}, "v1")
+	if err == nil {
+		t.Fatalf("expected sha mismatch")
+	}
+	if restarted {
+		t.Fatalf("restart should not be called")
+	}
+
+	current, err := os.ReadFile(binaryPath)
+	if err != nil {
+		t.Fatalf("read current binary: %v", err)
+	}
+	if string(current) != "old-binary" {
+		t.Fatalf("current binary should stay unchanged: %q", string(current))
+	}
+}
+
+func TestApplyUpgradeDoesNotRollbackWhenSelfRestartIsInterrupted(t *testing.T) {
+	tempDir := t.TempDir()
+	binaryPath := filepath.Join(tempDir, "incudal-agent")
+	if err := os.WriteFile(binaryPath, []byte("old-binary"), 0755); err != nil {
+		t.Fatalf("write current binary: %v", err)
+	}
+
+	nextBinary := []byte("new-binary")
+	packageBytes := gzipBytes(t, nextBinary)
+	sha := sha256Hex(packageBytes)
+	server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
+		response.WriteHeader(http.StatusOK)
+		_, _ = response.Write(packageBytes)
+	}))
+	defer server.Close()
+
+	runner := Runner{
+		BinaryPath:       binaryPath,
+		BackupPath:       binaryPath + ".bak",
+		LockPath:         filepath.Join(tempDir, "upgrade.lock"),
+		ServiceName:      "incudal-agent",
+		AllowedBaseURL:   server.URL,
+		HTTPClient:       server.Client(),
+		MaxDownloadBytes: 1024 * 1024,
+		Restart: func(context.Context, string) error {
+			return errors.New("signal: terminated")
+		},
+	}
+
+	err := runner.Apply(context.Background(), panel.UpgradeInstruction{
+		Available: true,
+		Version:   "v2",
+		URL:       server.URL + "/incudal-agent-linux-amd64.gz",
+		SHA256:    sha,
+		Gzip:      true,
+	}, "v1")
+	if err == nil {
+		t.Fatalf("expected restart error")
+	}
+
+	actual, err := os.ReadFile(binaryPath)
+	if err != nil {
+		t.Fatalf("read replaced binary: %v", err)
+	}
+	if string(actual) != string(nextBinary) {
+		t.Fatalf("binary should stay replaced after restart interruption: %q", string(actual))
+	}
+}
+
+func TestApplyUpgradeRejectsDifferentOrigin(t *testing.T) {
+	tempDir := t.TempDir()
+	binaryPath := filepath.Join(tempDir, "incudal-agent")
+	if err := os.WriteFile(binaryPath, []byte("old-binary"), 0755); err != nil {
+		t.Fatalf("write current binary: %v", err)
+	}
+
+	runner := Runner{
+		BinaryPath:     binaryPath,
+		BackupPath:     binaryPath + ".bak",
+		LockPath:       filepath.Join(tempDir, "upgrade.lock"),
+		AllowedBaseURL: "https://panel.example",
+		Restart: func(context.Context, string) error {
+			t.Fatalf("restart should not be called")
+			return nil
+		},
+	}
+
+	err := runner.Apply(context.Background(), panel.UpgradeInstruction{
+		Available: true,
+		Version:   "v2",
+		URL:       "https://evil.example/incudal-agent-linux-amd64.gz",
+		SHA256:    "0000000000000000000000000000000000000000000000000000000000000000",
+		Gzip:      true,
+	}, "v1")
+	if err == nil {
+		t.Fatalf("expected origin validation error")
+	}
+}
+
+func gzipBytes(t *testing.T, payload []byte) []byte {
+	t.Helper()
+
+	var buffer bytes.Buffer
+	writer := gzip.NewWriter(&buffer)
+	if _, err := writer.Write(payload); err != nil {
+		t.Fatalf("gzip write: %v", err)
+	}
+	if err := writer.Close(); err != nil {
+		t.Fatalf("gzip close: %v", err)
+	}
+	return buffer.Bytes()
+}
+
+func sha256Hex(payload []byte) string {
+	sum := sha256.Sum256(payload)
+	return hex.EncodeToString(sum[:])
+}