Sentinel-EdgeAgent/updater/main.go

package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"sentinel/pkg/config"
	"sentinel/pkg/device"
	"sentinel/pkg/log"
	"sentinel/pkg/net"
	"sentinel/pkg/platform"
	"time"
)

func main() {
	log.Println("[daemon] 设备ID:", device.GetDeviceID())

	log.Init(config.LOG_FILE_DIC)

	starter := newMainProgramStarter()
	// 启动主程序
	initMainProcess(starter)
	ticker := time.NewTicker(config.DAEMON_UPDATE_CHECK_INTERVAL_SECONDS * time.Second)
	defer ticker.Stop()

	go func() {
		for range ticker.C {
			//log.Println("[daemon] 定时检查主程序运行状态")
			handleProcessCheck(starter)

			cfg := config.LoadConfig()
			if cfg.NeedUpdate {
				log.Println("[daemon] 配置文件标记需要更新，触发更新")
				cfg.NeedUpdate = false
				if err := config.WriteLocalConfig(cfg); err != nil {
					log.Println("[daemon] 写回配置文件失败:", err)
				}
				handleUpdate(starter)
			}
		}
	}()
	select {}
}

func initMainProcess(starter platform.MainProgramStarter) {
	if err := startMainProgram(starter); err != nil {
		log.Println("[daemon] 启动失败，回退到更新:", err)
		handleUpdate(starter)
	}
}

var lastState string
var lastPID int

func handleProcessCheck(starter platform.MainProgramStarter) {
	cfg := config.LoadConfig()

	// 1. PID 是否存在
	running, err := starter.IsProcessRunning(cfg.PID)
	if cfg.PID <= 0 || err != nil || !running {
		if lastState != "NOT_RUNNING" {
			log.Println("[daemon] 主程序未运行，尝试启动")
			lastState = "NOT_RUNNING"
			lastPID = 0
		}
		restartMain(starter)
		return
	}

	// 2. 语义健康判断
	switch cfg.ControlState {
	case "DEGRADED":
		if lastState != "DEGRADED" {
			log.Println("[daemon] 主程序存在错误，等待主程序自行处理")
			lastState = "DEGRADED"
			lastPID = cfg.PID
		}
		return
	case "CONTROL_LOST":
		if lastState != "CONTROL_LOST" {
			log.Println("[daemon] 主程序心跳超时，已失控，强制重启")
			lastState = "CONTROL_LOST"
			lastPID = cfg.PID
		}
		restartMain(starter)
		return
	default:
		// 3. 健康
		if lastState != "HEALTHY" || lastPID != cfg.PID {
			log.Println("[daemon] 主程序健康运行，PID:", cfg.PID)
			lastState = "HEALTHY"
			lastPID = cfg.PID
		}
	}
}

func restartMain(starter platform.MainProgramStarter) {
	cfg := config.LoadConfig()
	if cfg.PID > 0 {
		_ = starter.KillProcess(cfg.PID)
		time.Sleep(300 * time.Millisecond)
	}
	_ = startMainProgram(starter)
}

func handleUpdate(starter platform.MainProgramStarter) {
	cfg := config.LoadConfig()
	info, err := api.CheckUpdate(device.GetDeviceID())
	if err != nil {
		log.Println("[daemon] 检查更新失败:", err)
		return
	}

	if info.Version <= cfg.VersionCode {
		log.Println("[daemon] 没有新版本，跳过")
		return
	}

	log.Println("[daemon] 发现新版本 %d\n", info.Version)

	// 停止主程序
	if cfg.PID > 0 {
		_ = starter.KillProcess(cfg.PID)
		time.Sleep(500 * time.Millisecond)
	}

	if err := downloadAndReplace(starter, info.DownloadURL); err != nil {
		log.Println("[daemon] 更新失败:", err)
		return
	}

	if err := startMainProgram(starter); err != nil {
		log.Println("[daemon] 更新后启动失败:", err)
		return
	}

	log.Println("[daemon] 更新完成")
}

func startMainProgram(starter platform.MainProgramStarter) error {
	selfPath, _ := os.Executable()
	selfDir := filepath.Dir(selfPath)
	targetExe := filepath.Join(selfDir, starter.GetMainName())

	if _, err := os.Stat(targetExe); os.IsNotExist(err) {
		log.Println("[daemon] 主程序文件不存在，触发更新")
		handleUpdate(starter)
		return nil // 先不启动，等更新完再启动
	}

	err := starter.Start(targetExe)
	return err
}

func downloadAndReplace(
	starter platform.MainProgramStarter,
	downloadURL string,
) error {
	// 获取主程序路径
	selfPath, err := os.Executable()
	if err != nil {
		return fmt.Errorf("获取自身路径失败: %w", err)
	}
	selfDir := filepath.Dir(selfPath)
	targetExe := filepath.Join(selfDir, starter.GetMainName())
	outFile, err := os.Create(targetExe)
	if err != nil {
		return fmt.Errorf("创建目标文件失败: %w", err)
	}
	defer outFile.Close()

	client := &http.Client{Timeout: 30 * time.Second}
	resp, err := client.Get(downloadURL)
	if err != nil {
		return fmt.Errorf("下载新程序失败: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("下载失败, HTTP状态码: %d", resp.StatusCode)
	}

	if _, err := io.Copy(outFile, resp.Body); err != nil {
		return fmt.Errorf("写入目标文件失败: %w", err)
	}
	return nil
}