u-desktop/updater/main.go

package main

import (
	"crypto/sha256"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

func main() {
	packageURL := flag.String("package-url", "", "下载地址")
	sha256Hash := flag.String("sha256", "", "SHA256 校验值")
	target := flag.String("target", "", "目标 exe 路径")
	ver := flag.String("version", "", "目标版本号")
	force := flag.Bool("force", false, "强制升级")
	backupDir := flag.String("backup-dir", "", "备份目录")
	flag.Parse()

	logPath := filepath.Join(filepath.Dir(*target), "updater.log")
	logFile, err := os.Create(logPath)
	if err != nil {
		fmt.Println("无法创建日志文件:", err)
		os.Exit(1)
	}
	defer logFile.Close()
	log.SetOutput(logFile)
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)

	log.Printf("updater 启动: version=%s target=%s force=%v", *ver, *target, *force)

	if *packageURL == "" || *target == "" {
		log.Fatal("缺少必要参数: --package-url, --target")
	}
	if *sha256Hash == "" {
		log.Fatal("缺少 --sha256 参数")
	}

	// 等待主程序退出
	time.Sleep(2 * time.Second)

	// 下载
	tmpFile, err := downloadWithRetry(*packageURL, 3)
	if err != nil {
		log.Fatal("下载失败: ", err)
	}
	defer os.Remove(tmpFile)
	log.Println("下载完成:", tmpFile)

	// 校验
	if err := verifySHA256(tmpFile, *sha256Hash); err != nil {
		log.Fatal("校验失败: ", err)
	}
	log.Println("SHA256 校验通过")

	// 备份
	if *backupDir == "" {
		*backupDir = filepath.Join(filepath.Dir(*target), "backup")
	}
	os.MkdirAll(*backupDir, 0755)
	backupPath := filepath.Join(*backupDir, fmt.Sprintf("old-%s.exe", time.Now().Format("20060102-150405")))
	if err := copyFile(*target, backupPath); err != nil {
		log.Println("备份失败(继续):", err)
	} else {
		log.Println("已备份:", backupPath)
	}

	// 替换: zip 包解压，exe 直接覆盖
	if strings.HasSuffix(strings.ToLower(*packageURL), ".zip") {
		if err := extractAndReplace(tmpFile, *target); err != nil {
			rollback(backupPath, *target)
			log.Fatal("解压替换失败: ", err)
		}
	} else {
		if err := copyFile(tmpFile, *target); err != nil {
			rollback(backupPath, *target)
			log.Fatal("替换失败: ", err)
		}
	}
	log.Println("替换完成")

	// 启动新版本
	cmd := exec.Command(*target)
	if err := cmd.Start(); err != nil {
		rollback(backupPath, *target)
		log.Fatal("启动新版本失败: ", err)
	}

	// 监听新进程 5 秒
	time.Sleep(5 * time.Second)
	if cmd.Process == nil {
		rollback(backupPath, *target)
		log.Fatal("新进程异常")
	}

	// 检查新进程是否仍在运行
	// NOTE: Windows 下无法直接用 cmd.Process.Signals，用简单的进程存活性判断
	if err := cmd.Process.Release(); err != nil {
		log.Println("新进程已退出，回滚")
		rollback(backupPath, *target)
		os.Exit(1)
	}

	log.Println("升级成功, updater 退出")
}

func downloadWithRetry(url string, maxRetries int) (string, error) {
	var lastErr error
	for i := 0; i < maxRetries; i++ {
		if i > 0 {
			log.Printf("重试下载 (%d/%d)...", i+1, maxRetries)
			time.Sleep(5 * time.Second)
		}
		tmpFile, err := os.CreateTemp("", "updater-*.download")
		if err != nil {
			return "", fmt.Errorf("创建临时文件失败: %w", err)
		}
		tmpPath := tmpFile.Name()

		resp, err := http.Get(url)
		if err != nil {
			tmpFile.Close()
			os.Remove(tmpPath)
			lastErr = err
			continue
		}

		if resp.StatusCode != 200 {
			tmpFile.Close()
			os.Remove(tmpPath)
			resp.Body.Close()
			lastErr = fmt.Errorf("HTTP %d", resp.StatusCode)
			continue
		}

		_, err = io.Copy(tmpFile, resp.Body)
		resp.Body.Close()
		tmpFile.Close()
		if err != nil {
			os.Remove(tmpPath)
			lastErr = err
			continue
		}
		return tmpPath, nil
	}
	return "", fmt.Errorf("下载失败(重试 %d 次): %w", maxRetries, lastErr)
}

func verifySHA256(path, expected string) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()

	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return err
	}

	actual := fmt.Sprintf("%x", h.Sum(nil))
	if actual != strings.ToLower(expected) {
		return fmt.Errorf("SHA256 不匹配: 期望 %s, 实际 %s", expected, actual)
	}
	return nil
}

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()

	if _, err := io.Copy(out, in); err != nil {
		return err
	}
	// 复制可执行权限
	if fi, err := os.Stat(src); err == nil {
		out.Chmod(fi.Mode())
	}
	return nil
}

func extractAndReplace(zipPath, target string) error {
	// TODO: 如需 zip 解压支持，引入 archive/zip
	// 初期可直接分发 exe，此路径不执行
	return fmt.Errorf("zip 解压暂未实现，请使用 exe 分发")
}

func rollback(backupPath, target string) {
	if backupPath == "" {
		return
	}
	if _, err := os.Stat(backupPath); err != nil {
		return
	}
	log.Println("回滚: 从备份恢复", backupPath)
	if err := copyFile(backupPath, target); err != nil {
		log.Println("回滚失败:", err)
		return
	}
	cmd := exec.Command(target)
	cmd.Start()
	log.Println("已回滚并启动旧版本")
}