rtu_linux_services/upgrade.service/upgrade.go

package main

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"rtu_linux_services/servicelib"

	"hnyfkj.com.cn/rtu/linux/baseapp"
	"hnyfkj.com.cn/rtu/linux/netmgrd"
	"hnyfkj.com.cn/rtu/linux/utils/ftpclient"
)

type serviceUpgradeConfig struct {
	Name        string
	Enable      bool
	CheckVerURL string
	ClonePkgURL string
	InstallPath string
	ReplaceDirs []string
}

var upgradeServices = []serviceUpgradeConfig{
	{ // 1, 网络管理服务
		Name:        "networkd",
		Enable:      true,
		CheckVerURL: "",
		ClonePkgURL: "",
		InstallPath: "/opt/yfkj/networkd.service",
	},
	{ // 2, 时间同步服务
		Name:        "timesyncd",
		Enable:      true,
		CheckVerURL: "",
		ClonePkgURL: "",
		InstallPath: "/opt/yfkj/timesyncd.service",
	},
	{ // 3, 地理位置服务
		Name:        "gnss",
		Enable:      true,
		CheckVerURL: "",
		ClonePkgURL: "",
		InstallPath: "/opt/yfkj/gnss.service",
	},
	{ // 4, 相机采集服务
		Name:        "camera-capture",
		Enable:      true,
		CheckVerURL: "",
		ClonePkgURL: "",
		InstallPath: "/opt/yfkj/camera-capture.service",
		ReplaceDirs: []string{"lib"},
	},
	{ // 5, 应用安装服务
		Name:        "app-install",
		Enable:      true,
		CheckVerURL: "",
		ClonePkgURL: "",
		InstallPath: "/opt/yfkj/app-install.service",
	},
	{ // 6, 远程运维服务
		Name:        "sshd-mqtt-bridge",
		Enable:      true,
		CheckVerURL: "",
		ClonePkgURL: "",
		InstallPath: "/opt/yfkj/sshd-mqtt-bridge.service",
	},
	{ // 7, 管理后台服务
		Name:        "web-ui",
		Enable:      true,
		CheckVerURL: "",
		ClonePkgURL: "",
		InstallPath: "/opt/yfkj/web-ui.service",
	},
	{ // 8, 自动升级服务
		Name:        "upgrade",
		Enable:      true,
		CheckVerURL: "",
		ClonePkgURL: "",
		InstallPath: "/opt/yfkj/upgrade.service",
	},
}

func isNetworkOK() bool {
	_, _, tcpOK, httpOK, _ := netmgrd.CheckNetwork()
	return tcpOK || httpOK
}

func (p *program) upgradeLoop() {
	for !isNetworkOK() {
		select {
		case <-time.After(30 * time.Second):
		case <-p.ctx.Done():
			return
		}
	}

	p.checkUpgrade()

	ticker := time.NewTicker(24 * time.Hour)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if isNetworkOK() {
				p.checkUpgrade()
			}
		case <-p.ctx.Done():
			return
		}
	}
}

func (p *program) checkUpgrade() {
	baseapp.Logger.Infof("[%s] ### 开始自动升级检查 ###", p.name)
	defer baseapp.Logger.Infof("[%s] ### 自动升级检查结束 ###", p.name)

	for _, service := range upgradeServices {
		if !service.Enable || service.CheckVerURL == "" || service.ClonePkgURL == "" {
			continue
		}

		select {
		case <-p.ctx.Done():
			return
		default:
		}

		oldVersion, newVersion, upgraded, err := p.upgradeService(service)
		if err != nil {
			baseapp.Logger.Errorf("[%s] yfkj-%s.service 自动升级失败: %v!!", p.name, service.Name, err)
			continue
		}

		if upgraded {
			baseapp.Logger.Infof("[%s] yfkj-%s.service 自动升级成功: %s -> %s", p.name, service.Name, oldVersion, newVersion)
		}
	}
}

func (p *program) upgradeService(service serviceUpgradeConfig) (string, string, bool, error) {
	currentVersion, err := servicelib.ReadVersionFile(service.InstallPath)
	if err != nil {
		return "", "", false, fmt.Errorf("读取当前版本信息失败: %w", err)
	}

	info, err := getUpgradeInfo(p.ctx, service.CheckVerURL, 30*time.Second)
	if err != nil {
		return "", "", false, fmt.Errorf("获取远程升级信息失败: %w", err)
	}

	var currentVersionCode, newVersionCode baseapp.VersionCode

	if err := currentVersionCode.FromString(currentVersion); err != nil {
		return "", "", false, fmt.Errorf("解析当前版本信息失败: %w", err)
	}

	if err := newVersionCode.FromString(info.Version); err != nil {
		return "", "", false, fmt.Errorf("解析远程版本信息失败: %w", err)
	}

	if newVersionCode <= currentVersionCode { // 只允许升级，不允许降级
		return currentVersion, info.Version, false, nil
	}

	currentDir, err := getCurrentVersionDir(service.InstallPath, service.Name)
	if err != nil {
		return "", "", false, fmt.Errorf("获取当前版本目录失败: %w", err)
	}

	targetName := nextVersionDir(currentDir)
	targetDir := filepath.Join(service.InstallPath, targetName)
	pkg := filepath.Join(os.TempDir(), service.Name+".tar.gz")
	defer os.Remove(pkg)

	if err := downloadUpgradePackage(p.ctx, service.ClonePkgURL, pkg, 10*time.Minute); err != nil {
		return currentVersion, "", false, fmt.Errorf("下载服务升级包失败: %w", err)
	}

	if err := servicelib.CheckFileSHA256(pkg, info.SHA256); err != nil {
		return currentVersion, "", false, fmt.Errorf("校验服务升级包失败: %w", err)
	}

	if err := installPackage(pkg, targetDir, service); err != nil {
		return currentVersion, "", false, fmt.Errorf("安装服务升级包失败: %w", err)
	}

	if err := switchServiceVersion(service.Name, service.InstallPath, targetName); err != nil {
		return currentVersion, "", false, fmt.Errorf("切换服务版本失败: %w", err)
	}

	return currentVersion, info.Version, true, nil
}

type upgradeInfo struct {
	Version string `json:"version"`
	SHA256  string `json:"sha256"`
}

// 获取远程升级信息，返回版本号和升级包SHA256校检值
func getUpgradeInfo(ctx context.Context, rawURL string, timeout time.Duration) (upgradeInfo, error) {
	u, err := servicelib.ParseUpgradeURL(rawURL)
	if err != nil {
		return upgradeInfo{}, err
	}

	switch u.Scheme {
	case "ftp":
		return getFtpUpgradeInfo(ctx, u, timeout)
	default:
		return upgradeInfo{}, fmt.Errorf("暂不支持的升级协议: %s", u.Scheme)
	}
}

func getFtpUpgradeInfo(ctx context.Context, u *servicelib.UpgradeURL, timeout time.Duration) (upgradeInfo, error) {
	localFile, err := ftpclient.DownloadFileFromFtp(ctx, u.Host, u.User, u.Pass, u.Path, timeout)
	if err != nil {
		return upgradeInfo{}, err
	}

	defer os.Remove(localFile)

	data, err := os.ReadFile(localFile)
	if err != nil {
		return upgradeInfo{}, err
	}

	var info upgradeInfo
	if err := json.Unmarshal(data, &info); err != nil {
		return upgradeInfo{}, err
	}

	info.Version = strings.TrimSpace(info.Version)
	info.SHA256 = strings.TrimSpace(info.SHA256)

	if info.Version == "" {
		return upgradeInfo{}, fmt.Errorf("版本号为空")
	}

	if info.SHA256 == "" {
		return upgradeInfo{}, fmt.Errorf("校验值为空")
	}

	if len(info.SHA256) != 64 {
		return upgradeInfo{}, fmt.Errorf("校验值长度错误")
	}

	if _, err := hex.DecodeString(info.SHA256); err != nil {
		return upgradeInfo{}, fmt.Errorf("校验值格式错误")
	}

	return info, nil
}

// 下载升级包.tar.gz文件，下载完成后存入 localFile
func downloadUpgradePackage(ctx context.Context, rawURL string, localFile string, timeout time.Duration) error {
	u, err := servicelib.ParseUpgradeURL(rawURL)
	if err != nil {
		return err
	}

	switch u.Scheme {
	case "ftp":
		return downloadFtpUpgradePackage(ctx, u, localFile, timeout)
	default:
		return fmt.Errorf("暂不支持的升级协议: %s", u.Scheme)
	}
}

func downloadFtpUpgradePackage(ctx context.Context, u *servicelib.UpgradeURL, localFile string, timeout time.Duration) error {
	file, err := ftpclient.DownloadFileFromFtp(ctx, u.Host, u.User, u.Pass, u.Path, timeout)
	if err != nil {
		return err
	}

	defer os.Remove(file)

	return os.Rename(file, localFile)
}

// 获取当前软链接指向的版本目录，返回: a、b、c 之一
func getCurrentVersionDir(installPath string, serviceName string) (string, error) {
	currentLink := filepath.Join(installPath, serviceName)

	target, err := os.Readlink(currentLink)
	if err != nil {
		return "", err
	}

	if !filepath.IsAbs(target) {
		target = filepath.Join(installPath, target)
	}

	target = filepath.Clean(target)

	name := filepath.Base(target)
	switch name {
	case "a", "b", "c":
		return target, nil
	default:
		return "", fmt.Errorf("当前软链接指向无效的版本目录: %s", target)
	}
}

// a 为出厂版本，固定不变；升级版本在 b、c 之间轮转
func nextVersionDir(currentDir string) string {
	switch filepath.Base(currentDir) {
	case "a":
		return "b"
	case "b":
		return "c"
	case "c":
		return "b"
	default:
		return "b"
	}
}

// 升级包默认包含与服务同名的可执行程序，之后再解压
func installPackage(pkg string, targetDir string, service serviceUpgradeConfig) error {
	tempDir, err := os.MkdirTemp(service.InstallPath, ".upgrade-")
	if err != nil {
		return fmt.Errorf("创建临时目录未成功: %w", err)
	}
	defer os.RemoveAll(tempDir)

	cmd := exec.Command("tar", "-xzf", pkg, "-C", tempDir)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("解压%s升级包未成功: %w: %s", service.Name, err, strings.TrimSpace(string(output)))
	}

	program := filepath.Join(tempDir, service.Name)
	if _, err := os.Stat(program); err != nil {
		return fmt.Errorf("检查%s服务程序失败: %w", service.Name, err)
	}

	if err := os.MkdirAll(targetDir, 0755); err != nil {
		return fmt.Errorf("创建%s版本目录失败: %w", service.Name, err)
	}

	if err := os.Rename(program, filepath.Join(targetDir, service.Name)); err != nil {
		return fmt.Errorf("安装%s服务程序失败: %w", service.Name, err)
	}

	if err := replacePackageDirs(tempDir, service.InstallPath, service.ReplaceDirs); err != nil {
		return fmt.Errorf("安装%s服务目录失败: %w", service.Name, err)
	}

	return nil
}

// 替换升级包中的附加目录，可用于升级服务的依赖库等
func replacePackageDirs(tempDir string, targetBaseDir string, dirs []string) error {
	for _, name := range dirs {
		srcDir := filepath.Join(tempDir, name)
		if _, err := os.Stat(srcDir); err != nil {
			if os.IsNotExist(err) { //// 如果升级包中没有该目录
				continue
			}
			return fmt.Errorf("检查%s目录失败: %w", name, err)
		}

		targetDir := filepath.Join(targetBaseDir, name)
		newDir := filepath.Join(targetBaseDir, "."+name+"-new")
		oldDir := filepath.Join(targetBaseDir, "."+name+"-old")

		if err := os.RemoveAll(newDir); err != nil {
			return fmt.Errorf("清理%s目录失败: %w", name, err)
		}

		if err := os.RemoveAll(oldDir); err != nil {
			return fmt.Errorf("清理%s目录失败: %w", name, err)
		}

		if err := os.Rename(srcDir, newDir); err != nil {
			return fmt.Errorf("移动%s目录失败: %w", name, err)
		}

		if err := os.Rename(targetDir, oldDir); err != nil && !os.IsNotExist(err) {
			return fmt.Errorf("备份%s目录失败: %w", name, err)
		}

		if err := os.Rename(newDir, targetDir); err != nil {
			if _, statErr := os.Stat(oldDir); statErr == nil {
				_ = os.Rename(oldDir, targetDir) ///// 还原旧目录
			}
			return fmt.Errorf("替换%s目录失败: %w", name, err)
		}

		if err := os.RemoveAll(oldDir); err != nil {
			return fmt.Errorf("删除%s目录失败: %w", name, err)
		}
	}

	return nil
}

// 切换服务版本，新版本启动失败时自动回滚并重启服务
func switchServiceVersion(serviceName string, installPath string, targetName string) error {
	service := "yfkj-" + serviceName + ".service"
	currentLink := filepath.Join(installPath, serviceName)
	tempLink := filepath.Join(installPath, serviceName+".tmp")

	currentTarget, err := os.Readlink(currentLink)
	if err != nil {
		return err
	}

	if err := exec.Command("systemctl", "stop", service).Run(); err != nil {
		return err
	}

	_ = os.Remove(tempLink)

	if err := os.Symlink(targetName, tempLink); err != nil {
		return err
	}

	if err := os.Rename(tempLink, currentLink); err != nil {
		_ = os.Remove(tempLink)
		return err
	}

	if err := exec.Command("systemctl", "start", service).Run(); err != nil {
		startErr := err

		if err := os.Symlink(currentTarget, tempLink); err == nil {
			if err := os.Rename(tempLink, currentLink); err == nil {
				_ = exec.Command("systemctl", "start", service).Run()
			}
		}

		return startErr
	}

	return nil
}